mirrors/vllm
0
0 Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-10-28 09:45:26 +08:00
Code Activity

Merge main; take main's disaggregated and fault-tolerance steps

Browse Source 
#63609 disabled the steps that install kv_connectors.txt at runtime, so
the dependency declarations this PR added are dropped. The PR now only
removes the two run_all_patterns entries.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: khluu <khluu000@gmail.com>
This commit is contained in:
khluuandClaude Opus 5.5 committed 2026-10-26 21:49:08 +00:00
2 parents 0437520df0 41f5146b5b
commit 31f3a6df1f
849 files changed
+71642 -13703
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.agents/skills/model-microbenchmark/SKILL.md
+3 -2
View File
@@ -7,8 +7,9 @@ description: Time a vLLM model's forward pass on exact batch shapes (decode, pre
Builds the batch without the scheduler, runs it through MRV2's real
`execute_model`, and times the forward as a one-shot CUDA graph (prefill
included). Sampling, the speculator, and CPU step overhead are excluded, so
expect results somewhat below a serving step.
included). The configured drafter is included by default; use
`--exclude-drafter` for target-only timing. All draft tokens are accepted.
Target sampling, rejection sampling, and CPU step overhead are excluded.
## Usage
.buildkite/hardware_tests/ascend_npu.yaml
+101 -100
View File
@@ -1,102 +1,103 @@
group: Hardware
depends_on: ~
steps:
- label: "Ascend NPU Test"
key: ascend-npu-test
soft_fail: true
no_plugin: true
device: ascend_npu
timeout_in_minutes: 30
source_file_dependencies:
- vllm/
- .buildkite/scripts/hardware_ci/run-npu-test.sh
- "!vllm/platforms/cuda.py"
- "!vllm/platforms/xpu.py"
- "!vllm/platforms/tpu.py"
- "!vllm/platforms/cpu.py"
- "!vllm/platforms/zen_cpu.py"
- "!vllm/cute_utils/"
- "!vllm/tilelang_utils/"
- "!vllm/device_allocator/"
- "!vllm/vllm_flash_attn/"
- "!vllm/kernels/"
- "!vllm/model_executor/kernels/"
- "!vllm/engine/"
- "!vllm/entrypoints/"
- "!vllm/ray/"
- "!vllm/v1/attention/backends/flash_attn.py"
- "!vllm/v1/attention/backends/flash_attn_diffkv.py"
- "!vllm/v1/attention/backends/flashinfer.py"
- "!vllm/v1/attention/backends/triton_attn.py"
- "!vllm/v1/attention/backends/triton_attn_diffkv.py"
- "!vllm/v1/attention/backends/rocm_attn.py"
- "!vllm/v1/attention/backends/rocm_aiter_fa.py"
- "!vllm/v1/attention/backends/rocm_aiter_unified_attn.py"
- "!vllm/v1/attention/backends/flex_attention.py"
- "!vllm/v1/attention/backends/hpc_attn.py"
- "!vllm/v1/attention/backends/mamba1_attn.py"
- "!vllm/v1/attention/backends/mamba2_attn.py"
- "!vllm/v1/attention/backends/mamba_attn.py"
- "!vllm/v1/attention/backends/short_conv_attn.py"
- "!vllm/v1/attention/backends/turboquant_attn.py"
- "!vllm/v1/attention/backends/cpu_attn.py"
- "!vllm/v1/attention/backends/fa_utils.py"
- "!vllm/distributed/device_communicators/pynccl.py"
- "!vllm/distributed/device_communicators/pynccl_wrapper.py"
- "!vllm/distributed/device_communicators/pynccl_allocator.py"
- "!vllm/distributed/device_communicators/xpu_communicator.py"
- "!vllm/distributed/device_communicators/cpu_communicator.py"
- "!vllm/distributed/device_communicators/cuda_wrapper.py"
- "!vllm/distributed/device_communicators/flashinfer_all_reduce.py"
- "!vllm/distributed/device_communicators/aiter_custom_all_reduce.py"
- "!vllm/distributed/device_communicators/custom_all_reduce.py"
- "!vllm/distributed/device_communicators/quick_all_reduce.py"
- "!vllm/distributed/device_communicators/all2all.py"
- "!vllm/distributed/device_communicators/all_reduce_utils.py"
- "!vllm/distributed/device_communicators/shm_object_storage.py"
- "!vllm/distributed/device_communicators/symm_mem.py"
- "!vllm/distributed/device_communicators/mnnvl_compat.py"
- "!vllm/third_party/flashmla/"
- "!vllm/third_party/pynvml.py"
- "!vllm/_custom_ops.py"
- "!vllm/_xpu_ops.py"
- "!vllm/_aiter_ops"
- "!vllm/collect_env.py"
- "!vllm/connections.py"
- "!vllm/env_override.py"
- "!vllm/logits_process.py"
- "!vllm/logprobs.py"
- "!vllm/model_inspection.py"
- "!vllm/outputs.py"
- "!vllm/scalar_type.py"
- "!vllm/version.py"
- "!vllm/assets/"
- "!vllm/benchmarks/"
- "!vllm/inputs/"
- "!vllm/ir/"
- "!vllm/parser/"
- "!vllm/plugins/"
- "!vllm/reasoning/"
- "!vllm/renderers/"
- "!vllm/tokenizers/"
- "!vllm/tool_parsers/"
- "!vllm/tracing/"
- "!vllm/usage/"
- "!tests/"
- "!benchmarks/"
- "!docs/"
- "!examples/"
- "!vllm/model_executor/layers/quantization/auto_awq.py"
- "!vllm/model_executor/layers/quantization/auto_gptq.py"
- "!vllm/model_executor/layers/quantization/inc/"
- "!vllm/model_executor/layers/quantization/awq_triton.py"
- "!vllm/model_executor/layers/quantization/experts_int8.py"
- "!vllm/model_executor/layers/quantization/fbgemm_fp8.py"
- "!vllm/model_executor/layers/quantization/fp_quant.py"
- "!vllm/model_executor/layers/quantization/humming.py"
- "!vllm/model_executor/layers/quantization/moe_wna16.py"
- "!vllm/model_executor/layers/quantization/qutlass_utils.py"
- "!vllm/model_executor/layers/quantization/quark/"
- "!vllm/model_executor/layers/quantization/turboquant/"
commands:
- bash .buildkite/scripts/hardware_ci/run-npu-test.sh
steps: []
# Disabled until the Ascend packages installation is locked.
# - label: "Ascend NPU Test"
# key: ascend-npu-test
# soft_fail: true
# no_plugin: true
# device: ascend_npu
# timeout_in_minutes: 30
# source_file_dependencies:
# - vllm/
# - .buildkite/scripts/hardware_ci/run-npu-test.sh
# - "!vllm/platforms/cuda.py"
# - "!vllm/platforms/xpu.py"
# - "!vllm/platforms/tpu.py"
# - "!vllm/platforms/cpu.py"
# - "!vllm/platforms/zen_cpu.py"
# - "!vllm/cute_utils/"
# - "!vllm/tilelang_utils/"
# - "!vllm/device_allocator/"
# - "!vllm/vllm_flash_attn/"
# - "!vllm/kernels/"
# - "!vllm/model_executor/kernels/"
# - "!vllm/engine/"
# - "!vllm/entrypoints/"
# - "!vllm/ray/"
# - "!vllm/v1/attention/backends/flash_attn.py"
# - "!vllm/v1/attention/backends/flash_attn_diffkv.py"
# - "!vllm/v1/attention/backends/flashinfer.py"
# - "!vllm/v1/attention/backends/triton_attn.py"
# - "!vllm/v1/attention/backends/triton_attn_diffkv.py"
# - "!vllm/v1/attention/backends/rocm_attn.py"
# - "!vllm/v1/attention/backends/rocm_aiter_fa.py"
# - "!vllm/v1/attention/backends/rocm_aiter_unified_attn.py"
# - "!vllm/v1/attention/backends/flex_attention.py"
# - "!vllm/v1/attention/backends/hpc_attn.py"
# - "!vllm/v1/attention/backends/mamba1_attn.py"
# - "!vllm/v1/attention/backends/mamba2_attn.py"
# - "!vllm/v1/attention/backends/mamba_attn.py"
# - "!vllm/v1/attention/backends/short_conv_attn.py"
# - "!vllm/v1/attention/backends/turboquant_attn.py"
# - "!vllm/v1/attention/backends/cpu_attn.py"
# - "!vllm/v1/attention/backends/fa_utils.py"
# - "!vllm/distributed/device_communicators/pynccl.py"
# - "!vllm/distributed/device_communicators/pynccl_wrapper.py"
# - "!vllm/distributed/device_communicators/pynccl_allocator.py"
# - "!vllm/distributed/device_communicators/xpu_communicator.py"
# - "!vllm/distributed/device_communicators/cpu_communicator.py"
# - "!vllm/distributed/device_communicators/cuda_wrapper.py"
# - "!vllm/distributed/device_communicators/flashinfer_all_reduce.py"
# - "!vllm/distributed/device_communicators/aiter_custom_all_reduce.py"
# - "!vllm/distributed/device_communicators/custom_all_reduce.py"
# - "!vllm/distributed/device_communicators/quick_all_reduce.py"
# - "!vllm/distributed/device_communicators/all2all.py"
# - "!vllm/distributed/device_communicators/all_reduce_utils.py"
# - "!vllm/distributed/device_communicators/shm_object_storage.py"
# - "!vllm/distributed/device_communicators/symm_mem.py"
# - "!vllm/distributed/device_communicators/mnnvl_compat.py"
# - "!vllm/third_party/flashmla/"
# - "!vllm/third_party/pynvml.py"
# - "!vllm/_custom_ops.py"
# - "!vllm/_xpu_ops.py"
# - "!vllm/_aiter_ops"
# - "!vllm/collect_env.py"
# - "!vllm/connections.py"
# - "!vllm/env_override.py"
# - "!vllm/logits_process.py"
# - "!vllm/logprobs.py"
# - "!vllm/model_inspection.py"
# - "!vllm/outputs.py"
# - "!vllm/scalar_type.py"
# - "!vllm/version.py"
# - "!vllm/assets/"
# - "!vllm/benchmarks/"
# - "!vllm/inputs/"
# - "!vllm/ir/"
# - "!vllm/parser/"
# - "!vllm/plugins/"
# - "!vllm/reasoning/"
# - "!vllm/renderers/"
# - "!vllm/tokenizers/"
# - "!vllm/tool_parsers/"
# - "!vllm/tracing/"
# - "!vllm/usage/"
# - "!tests/"
# - "!benchmarks/"
# - "!docs/"
# - "!examples/"
# - "!vllm/model_executor/layers/quantization/auto_awq.py"
# - "!vllm/model_executor/layers/quantization/auto_gptq.py"
# - "!vllm/model_executor/layers/quantization/inc/"
# - "!vllm/model_executor/layers/quantization/awq_triton.py"
# - "!vllm/model_executor/layers/quantization/experts_int8.py"
# - "!vllm/model_executor/layers/quantization/fbgemm_fp8.py"
# - "!vllm/model_executor/layers/quantization/fp_quant.py"
# - "!vllm/model_executor/layers/quantization/humming.py"
# - "!vllm/model_executor/layers/quantization/moe_wna16.py"
# - "!vllm/model_executor/layers/quantization/qutlass_utils.py"
# - "!vllm/model_executor/layers/quantization/quark/"
# - "!vllm/model_executor/layers/quantization/turboquant/"
# commands:
# - bash .buildkite/scripts/hardware_ci/run-npu-test.sh
.buildkite/hardware_tests/cpu.yaml
+3
View File
@@ -26,6 +26,8 @@ steps:
- tests/kernels/test_cpu_fused_sampling.py
- tests/kernels/test_awq_int4_to_int8.py
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/quantization/test_cpu_fp8_quant.py
- vllm/model_executor/kernels/linear/scaled_mm/cpu.py
- tests/kernels/quantization/test_zen_da8w4.py
- vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
- vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py
@@ -53,6 +55,7 @@ steps:
tests/kernels/test_cpu_fused_sampling.py \
tests/kernels/test_awq_int4_to_int8.py \
tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \
tests/kernels/quantization/test_cpu_fp8_quant.py \
tests/kernels/quantization/test_zen_da8w4.py \
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \
tests/kernels/mamba/test_causal_conv1d.py \
.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+46 -44
View File
@@ -15,28 +15,29 @@ steps:
limit: 2
- exit_status: -10 # Agent was lost
limit: 2
- label: ":intel: (BMG) Example Test"
key: xpu-example-test
depends_on:
- image-build-xpu
timeout_in_minutes: 50
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
# Disabled until tblib transitive dependencies are locked.
# - label: ":intel: (BMG) Example Test"
# key: xpu-example-test
# depends_on:
# - image-build-xpu
# timeout_in_minutes: 50
# optional: true
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 2+
# mem: 24+
# no_plugin: true
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# source_file_dependencies:
# - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
# - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
- label: ":intel: (BMG) V1 test"
key: xpu-v1-test
depends_on:
@@ -83,28 +84,29 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh w8a8-fp8-linear'
- label: ":intel: (BMG) Server test"
key: xpu-server-test
depends_on:
- image-build-xpu
timeout_in_minutes: 45
optional: true
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'
# Disabled until the av installation is locked.
# - label: ":intel: (BMG) Server test"
# key: xpu-server-test
# depends_on:
# - image-build-xpu
# timeout_in_minutes: 45
# optional: true
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 16+
# no_plugin: true
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# source_file_dependencies:
# - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
# - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'
- label: ":intel: (BMG) Quantization test"
key: xpu-quantization-test
depends_on:
.buildkite/intel_jobs/basic_correctness_intel.yaml
+1 -1
View File
@@ -26,4 +26,4 @@ steps:
'cd tests &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
pytest -v -s basic_correctness/cpu_offload &&
pytest -v -s basic_correctness/memory'
pytest -v -s basic_correctness/memory --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem"'
.buildkite/intel_jobs/engine_intel.yaml
+30
View File
@@ -65,6 +65,36 @@ steps:
pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py --ignore v1/engine/test_core_engine_actor_manager.py &&
pytest -v -s test_config.py'
- label: ":intel: (BMG) E2E Core"
key: e2e-core-1-gpu
timeout_in_minutes: 40
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/v1/core/
- vllm/v1/worker/
- vllm/v1/attention/backends/gdn_attn.py
- vllm/v1/attention/backends/mamba_attn.py
- vllm/model_executor/layers/mamba/
- vllm/model_executor/models/qwen3_next
- vllm/config/cache.py
- vllm/platforms/xpu.py
- tests/v1/e2e/general/test_mamba_prefix_cache.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py'
- label: ":intel: (BMG) V1 e2e (2 GPUs)"
key: v1-e2e-2-gpus
timeout_in_minutes: 30
.buildkite/intel_jobs/misc_intel.yaml
+66 -64
View File
@@ -30,71 +30,73 @@ steps:
pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py &&
bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh'
- label: ":intel: (BMG) V1 Core + KV + Metrics"
key: v1-core-kv-metrics-intel
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/v1/core
- tests/v1/executor
- tests/v1/kv_offload
- tests/v1/worker
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install -r requirements/kv_connectors.txt &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
cd tests &&
pytest -v -s v1/executor &&
VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit
--ignore=v1/ec_connector/unit/test_ec_cpu_connector.py'
# Disabled until the KV connectors installation is locked.
# - label: ":intel: (BMG) V1 Core + KV + Metrics"
# key: v1-core-kv-metrics-intel
# timeout_in_minutes: 45
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 16+
# no_plugin: true
# working_dir: "."
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# VLLM_TEST_DEVICE: "xpu"
# source_file_dependencies:
# - vllm/
# - tests/v1/core
# - tests/v1/executor
# - tests/v1/kv_offload
# - tests/v1/worker
# - tests/v1/kv_connector/unit
# - tests/v1/ec_connector/unit
# - tests/v1/metrics
# - tests/entrypoints/openai/correctness/test_lmeval.py
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'pip install -r requirements/kv_connectors.txt &&
# export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
# cd tests &&
# pytest -v -s v1/executor &&
# VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit
# --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py'
- label: ":intel: (BMG) EC CPU Connector (XPU)"
key: ec-cpu-connector-xpu
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/v1/core
- tests/v1/executor
- tests/v1/kv_offload
- tests/v1/worker
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install -r requirements/kv_connectors.txt &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
cd tests &&
pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py'
# Disabled until the KV connectors installation is locked.
# - label: ":intel: (BMG) EC CPU Connector (XPU)"
# key: ec-cpu-connector-xpu
# timeout_in_minutes: 45
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 24+
# no_plugin: true
# working_dir: "."
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# VLLM_TEST_DEVICE: "xpu"
# source_file_dependencies:
# - vllm/
# - tests/v1/core
# - tests/v1/executor
# - tests/v1/kv_offload
# - tests/v1/worker
# - tests/v1/kv_connector/unit
# - tests/v1/ec_connector/unit
# - tests/v1/metrics
# - tests/entrypoints/openai/correctness/test_lmeval.py
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'pip install -r requirements/kv_connectors.txt &&
# export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
# cd tests &&
# pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py'
- label: ":intel: (BMG) V1 Sample + Logits"
key: v1-sample-logits-intel
.buildkite/intel_jobs/model_executor_intel.yaml
+57 -55
View File
@@ -1,59 +1,61 @@
group: Model Executor Intel
depends_on:
- image-build-xpu
steps:
- label: ":intel: (BMG) Model Executor"
key: model-executor-intel
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'apt-get update && apt-get install -y curl libsodium23 &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
export PYTHONFAULTHANDLER=1 &&
cd tests &&
pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --ignore="model_executor/model_loader/test_weight_cache.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"'
steps: []
# Disabled until the apt packages installation is locked.
# - label: ":intel: (BMG) Model Executor"
# key: model-executor-intel
# timeout_in_minutes: 45
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 16+
# no_plugin: true
# working_dir: "."
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# VLLM_TEST_DEVICE: "xpu"
# source_file_dependencies:
# - vllm/engine/arg_utils.py
# - vllm/config/model.py
# - vllm/model_executor
# - tests/model_executor
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'apt-get update && apt-get install -y curl libsodium23 &&
# export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
# export PYTHONFAULTHANDLER=1 &&
# cd tests &&
# pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --ignore="model_executor/model_loader/test_weight_cache.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"'
- label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)"
key: tensorizer-model-loader-tests-xpu
timeout_in_minutes: 45
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'apt-get update && apt-get install -y curl libsodium23 &&
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
export PYTHONFAULTHANDLER=1 &&
cd tests &&
pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"'
# Disabled until the apt packages installation is locked.
# - label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)"
# key: tensorizer-model-loader-tests-xpu
# timeout_in_minutes: 45
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 24+
# no_plugin: true
# working_dir: "."
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# VLLM_TEST_DEVICE: "xpu"
# source_file_dependencies:
# - vllm/engine/arg_utils.py
# - vllm/config/model.py
# - vllm/model_executor
# - tests/model_executor
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'apt-get update && apt-get install -y curl libsodium23 &&
# export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
# export PYTHONFAULTHANDLER=1 &&
# cd tests &&
# pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"'
.buildkite/intel_jobs/models_language_intel.yaml
+58
View File
@@ -0,0 +1,58 @@
group: Models - Language Intel
depends_on:
- image-build-xpu
steps:
- label: ":intel: (BMG) Language Models (Hybrid)"
key: language-models-hybrid
timeout_in_minutes: 40
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/v1/core/
- vllm/v1/worker/
- vllm/v1/attention/backends/mamba1_attn.py
- vllm/v1/attention/backends/mamba_attn.py
- vllm/model_executor/layers/mamba/
- vllm/model_executor/models/jamba.py
- vllm/model_executor/models/mamba.py
- vllm/config/cache.py
- vllm/platforms/xpu.py
- tests/models/language/generation/test_hybrid.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s models/language/generation/test_hybrid.py -k apc'
- label: ":intel: (BMG) Language Models (Pooling)"
key: language-models-pooling
timeout_in_minutes: 30
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
optional: true
source_file_dependencies:
- vllm/
- tests/models/language/pooling
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s models/language/pooling/test_auto_prefix_cache_support.py'
.buildkite/intel_jobs/test-intel.yaml
+59 -57
View File
@@ -76,63 +76,65 @@ steps:
python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 &&
python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096
'
- label: ":intel: (BMG) V1 test"
key: xpu-v1-test
depends_on:
- image-build-xpu
timeout_in_minutes: 70
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests && pip install dpctl &&
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py --ignore=v1/engine/test_core_engine_actor_manager.py &&
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/e2e/general/test_correctness_sliding_window.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]'
parallelism: 2
- label: ":intel: (BMG) Qwen3-30B + V1 KV Connector"
key: xpu-qwen3-30b-v1-kv-connector
depends_on:
- image-build-xpu
timeout_in_minutes: 90
device: intel_gpu
agent_tags:
label: production
gpu: 2+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pip install dpctl &&
pytest -v -s v1/kv_connector/unit/test_multi_connector.py &&
pytest -v -s v1/kv_connector/unit/test_example_connector.py &&
cd .. &&
VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192'
# Disabled until the dpctl installation is locked.
# - label: ":intel: (BMG) V1 test"
# key: xpu-v1-test
# depends_on:
# - image-build-xpu
# timeout_in_minutes: 70
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 1+
# mem: 16+
# no_plugin: true
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# source_file_dependencies:
# - vllm/
# - .buildkite/intel_jobs/test-intel.yaml
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'cd tests && pip install dpctl &&
# pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
# pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py --ignore=v1/engine/test_core_engine_actor_manager.py &&
# pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
# pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
# pytest -v -s v1/structured_output &&
# pytest -v -s v1/test_serial_utils.py &&
# pytest -v -s v1/e2e/general/test_correctness_sliding_window.py &&
# pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head &&
# pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]'
# parallelism: 2
# Disabled until the dpctl installation is locked.
# - label: ":intel: (BMG) Qwen3-30B + V1 KV Connector"
# key: xpu-qwen3-30b-v1-kv-connector
# depends_on:
# - image-build-xpu
# timeout_in_minutes: 90
# device: intel_gpu
# agent_tags:
# label: production
# gpu: 2+
# mem: 24+
# no_plugin: true
# env:
# REGISTRY: "public.ecr.aws/q9t5s3a7"
# REPO: "vllm-ci-test-repo"
# source_file_dependencies:
# - vllm/
# - .buildkite/intel_jobs/test-intel.yaml
# commands:
# - >-
# bash .buildkite/scripts/hardware_ci/run-intel-test.sh
# 'cd tests &&
# pip install dpctl &&
# pytest -v -s v1/kv_connector/unit/test_multi_connector.py &&
# pytest -v -s v1/kv_connector/unit/test_example_connector.py &&
# cd .. &&
# VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192'
- label: ":intel: (BMG) Server test"
key: xpu-server-test
depends_on:
.buildkite/release-pipeline.yaml
+87 -45
View File
@@ -300,58 +300,80 @@ steps:
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images'
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-x86-ubuntu2404
# ---- AWS EFA variants: docker/Dockerfile.efa on the images built above ----
- label: "Build release image - x86_64 - CUDA 13.0 - EFA"
depends_on:
- build-release-image-x86
id: build-release-image-x86-efa
agents:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359 \
--target vllm-openai \
$(bash .buildkite/scripts/docker-build-metadata-args.sh efa) \
--build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" release-images'
-f docker/Dockerfile.efa docker/
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" release-images'
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-arm64-ubuntu2404
- label: "Build release image - aarch64 - CUDA 13.0 - EFA"
depends_on:
- build-release-image-arm64
id: build-release-image-arm64-efa
agents:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-b8b5f17a7d9ccfc25bbc5cf17b3fcea12964a042 \
--target vllm-openai \
$(bash .buildkite/scripts/docker-build-metadata-args.sh efa) \
--build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" release-images'
-f docker/Dockerfile.efa docker/
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" release-images'
- label: "Build release image - x86_64 - Rubin CUDA 13.4 - EFA"
depends_on:
- build-release-image-x86-rubin-cuda-13-4
id: build-release-image-x86-rubin-cuda-13-4-efa
if: build.env("NIGHTLY") == "1"
agents:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu134-efa) \
--build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134 \
--progress plain \
-f docker/Dockerfile.efa docker/
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" release-images'
- label: "Build release image - aarch64 - Rubin CUDA 13.4 - EFA"
depends_on:
- build-release-image-arm64-rubin-cuda-13-4
id: build-release-image-arm64-rubin-cuda-13-4-efa
if: build.env("NIGHTLY") == "1"
agents:
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu134-efa) \
--build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134 \
--progress plain \
-f docker/Dockerfile.efa docker/
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" release-images'
- label: ":docker: Build release image - x86_64 - XPU"
depends_on: ~
id: build-xpu-release-image
@@ -446,19 +468,35 @@ steps:
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" release-manifests'
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
# ---- AWS EFA variant manifests ----
- label: "Create multi-arch manifest - CUDA 13.0 - EFA"
depends_on:
- build-release-image-x86-ubuntu2404
- build-release-image-arm64-ubuntu2404
id: create-multi-arch-manifest-ubuntu2404
- build-release-image-x86-efa
- build-release-image-arm64-efa
id: create-multi-arch-manifest-efa
agents:
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-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" release-manifests'
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-efa --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 - EFA" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa" release-manifests'
- label: "Create multi-arch manifest - Rubin CUDA 13.4 - EFA"
depends_on:
- build-release-image-x86-rubin-cuda-13-4-efa
- build-release-image-arm64-rubin-cuda-13-4-efa
id: create-multi-arch-manifest-rubin-cuda-13-4-efa
if: build.env("NIGHTLY") == "1"
agents:
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-cu134-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu134-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu134-efa --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134-efa"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4 - EFA" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134-efa" release-manifests'
- label: "Create manifest - XPU"
depends_on:
- build-xpu-release-image
@@ -814,7 +852,6 @@ steps:
key: block-publish-release-images
depends_on:
- create-multi-arch-manifest
- create-multi-arch-manifest-ubuntu2404
- create-manifest-xpu
- build-rocm100-release-image
- build-rocm72-release-image
@@ -827,6 +864,10 @@ steps:
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
# Wait for the EFA images, but don't let an EFA failure block publishing
# the others; the cuda-13-0-efa publish job fails on a partial build.
- step: create-multi-arch-manifest-efa
allow_failure: true
- label: "Publish {{matrix}} release images to DockerHub"
depends_on:
@@ -838,7 +879,8 @@ steps:
- "bash .buildkite/scripts/publish-release-images.sh {{matrix}}"
matrix:
- "cuda-13-0"
- "cuda-13-0-ubuntu-24-04"
# AWS EFA variant; skipped if no EFA image was built, fails on a partial build.
- "cuda-13-0-efa"
- "rocm"
- "rocm72"
- "xpu"
.buildkite/scripts/ci-bake-rocm.sh
+67 -1
View File
@@ -864,6 +864,50 @@ should_export_rocm_smoke() {
|| "${TARGET}" == "smoke-test-rocm-ci" ]]
}
should_export_kernel_symbol_map() {
[[ "${VLLM_KERNEL_SYMBOL_MAP:-0}" == "1" ]] || return 1
case "${TARGET}" in
test-rocm-ci|test-rocm-ci-with-wheel|test-rocm-ci-with-artifacts|\
export-wheel-rocm|csrc-rocm-ci|smoke-test-rocm-ci|kernel-symbol-map-rocm)
return 0 ;;
*) return 1 ;;
esac
}
export_kernel_symbol_map() {
should_export_kernel_symbol_map || return 0
echo "--- :world_map: Exporting ROCm kernel symbol map"
rm -rf ./kernel-symbol-map-rocm || return 1
docker buildx bake \
"${BAKE_ALLOW_ARGS[@]}" "${BAKE_FILES[@]}" \
--progress "${BUILDKIT_PROGRESS:-plain}" \
kernel-symbol-map-rocm || return 1
upload_kernel_symbol_map_if_present
}
export_kernel_symbol_map_from_cache() {
should_export_kernel_symbol_map || return 0
# A separate shell keeps setup failures fatal inside the optional export.
FORCE_BUILD=1 VLLM_KERNEL_SYMBOL_MAP=1 UPLOAD_ROCM_WHEEL_ARTIFACTS=0 \
CI_HCL_SOURCE="${CI_HCL_PATH}" VLLM_BAKE_FILE="${VLLM_BAKE_FILE}" \
BUILDER_NAME="${BUILDER_NAME}" \
bash "${BASH_SOURCE[0]}" kernel-symbol-map-rocm \
|| echo "ROCm kernel symbol map export from cache failed; continuing without it" >&2
}
upload_kernel_symbol_map_if_present() {
should_export_kernel_symbol_map || return 0
local map_dir="./kernel-symbol-map-rocm"
local map_name="kernel_symbol_map.rocm.json.gz"
if [[ ! -s "${map_dir}/${map_name}" ]]; then
echo "ROCm kernel symbol map export is missing: ${map_dir}/${map_name}" >&2
return 1
fi
if [[ "${BUILDKITE:-false}" == "true" ]]; then
(cd "${map_dir}" && buildkite-agent artifact upload "${map_name}")
fi
}
verify_rocm_smoke_export() {
local marker="./build/rocm-smoke-export/vllm-smoke-ok"
local expected_smoke_id="${BUILDKITE_BUILD_ID:-local}"
@@ -1453,7 +1497,8 @@ maybe_skip_existing_image() {
return 0
fi
if ! is_ci_base_target \
&& { should_upload_wheel_artifacts || should_export_rocm_smoke; }; then
&& { should_upload_wheel_artifacts || should_export_rocm_smoke \
|| [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; }; then
echo "Local-output targets always run for the current build"
return 0
fi
@@ -1501,11 +1546,13 @@ maybe_skip_existing_image() {
echo "Commit image already exists: ${IMAGE_TAG}"
echo "Skipping build"
export_kernel_symbol_map_from_cache
exit 0
fi
echo "Image already exists: ${IMAGE_TAG}"
echo "Skipping build"
export_kernel_symbol_map_from_cache
exit 0
}
@@ -1806,6 +1853,7 @@ EOF
uses_rocm_csrc_cache() {
case "${TARGET}" in
csrc-rocm-ci \
| kernel-symbol-map-rocm \
| test-rocm-ci \
| test-rocm-ci-with-wheel \
| test-rocm-ci-with-artifacts \
@@ -2252,6 +2300,15 @@ EOF
cat <<EOF
}
target "kernel-symbol-map-rocm" {
cache-from = concat(
get_cache_from_rocm_csrc(),
EOF
write_hcl_string_list " " "${csrc_content_cache_from[@]}"
cat <<EOF
)
}
target "test-rocm-ci" {
cache-from = concat(
get_cache_from_rocm(),
@@ -2724,12 +2781,21 @@ main() {
# from an earlier build or retry.
rm -rf ./build/rocm-smoke-export
fi
if [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; then
rm -rf ./kernel-symbol-map-rocm
fi
seed_dependency_caches_if_needed
run_bake
verify_rocm_smoke_export
promote_stable_ci_base_tag
publish_ci_base_handoff_ref
upload_wheel_artifacts_if_present
if [[ "${TARGET}" == "kernel-symbol-map-rocm" ]]; then
upload_kernel_symbol_map_if_present
else
export_kernel_symbol_map \
|| echo "ROCm kernel symbol map export or upload failed; continuing without it" >&2
fi
}
if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
.buildkite/scripts/publish-release-images.sh
+65 -42
View File
@@ -10,9 +10,9 @@ set -euo pipefail
TARGET="${1:-all}"
case "${TARGET}" in
cuda-13-0 | cuda-13-0-ubuntu-24-04 | rocm | rocm72 | xpu | cpu | all) ;;
cuda-13-0 | cuda-13-0-efa | rocm | rocm72 | xpu | cpu | all) ;;
*)
echo "Usage: $0 {cuda-13-0|cuda-13-0-ubuntu-24-04|rocm|rocm72|xpu|cpu|all}"
echo "Usage: $0 {cuda-13-0|cuda-13-0-efa|rocm|rocm72|xpu|cpu|all}"
exit 2
;;
esac
@@ -21,6 +21,63 @@ target_enabled() {
[ "${TARGET}" = "all" ] || [ "${TARGET}" = "$1" ]
}
# Promote one image family from ECR staging to DockerHub: per-arch tags plus the
# latest and versioned multi-arch manifests. $1 is the tag suffix shared by both
# registries: "" for the default CUDA family, else e.g. "-efa".
publish_image_family() {
local suffix="$1" arch
local staging="public.ecr.aws/q9t5s3a7/vllm-release-repo"
local latest="vllm/vllm-openai:latest${suffix}"
local versioned="vllm/vllm-openai:v${RELEASE_VERSION}${suffix}"
for arch in x86_64 aarch64; do
docker pull "${staging}:${COMMIT}-${arch}${suffix}"
done
for arch in x86_64 aarch64; do
docker tag "${staging}:${COMMIT}-${arch}${suffix}" "vllm/vllm-openai:latest-${arch}${suffix}"
docker tag "${staging}:${COMMIT}-${arch}${suffix}" "vllm/vllm-openai:v${RELEASE_VERSION}-${arch}${suffix}"
docker push "vllm/vllm-openai:latest-${arch}${suffix}"
docker push "vllm/vllm-openai:v${RELEASE_VERSION}-${arch}${suffix}"
done
docker manifest rm "${latest}" || true
docker manifest rm "${versioned}" || true
docker manifest create "${latest}" \
"vllm/vllm-openai:latest-x86_64${suffix}" \
"vllm/vllm-openai:latest-aarch64${suffix}"
docker manifest create "${versioned}" \
"vllm/vllm-openai:v${RELEASE_VERSION}-x86_64${suffix}" \
"vllm/vllm-openai:v${RELEASE_VERSION}-aarch64${suffix}"
docker manifest push "${latest}"
docker manifest push "${versioned}"
}
# Publish an opt-in family all-or-nothing, like the CPU images below: skip when
# neither arch was built, fail on a partial build so it is not silenced.
publish_optional_family() {
local suffix="$1" arch missing=() present=()
local staging="public.ecr.aws/q9t5s3a7/vllm-release-repo"
for arch in x86_64 aarch64; do
if docker manifest inspect "${staging}:${COMMIT}-${arch}${suffix}" >/dev/null 2>&1; then
present+=("${arch}")
else
missing+=("${arch}")
fi
done
if [ "${#present[@]}" -eq 0 ]; then
echo "WARNING: no ${suffix} images found in ECR for ${COMMIT}, skipping the ${suffix} publish"
elif [ "${#missing[@]}" -eq 0 ]; then
publish_image_family "${suffix}"
else
echo "ERROR: partial ${suffix} build (present: ${present[*]}; missing: ${missing[*]})." >&2
echo " Refusing to publish to avoid per-arch and multi-arch tags disagreeing." >&2
exit 1
fi
}
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
@@ -41,25 +98,7 @@ aws ecr-public get-login-password --region us-east-1 | \
# ---- CUDA (default: 13.0) ----
if target_enabled cuda-13-0; then
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64"
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64"
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" vllm/vllm-openai:latest-x86_64
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64"
docker push vllm/vllm-openai:latest-x86_64
docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64"
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" vllm/vllm-openai:latest-aarch64
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64"
docker push vllm/vllm-openai:latest-aarch64
docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64"
docker manifest rm vllm/vllm-openai:latest || true
docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}" || true
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64"
docker manifest push vllm/vllm-openai:latest
docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}"
publish_image_family ""
ZSTD_DIGEST=$(docker buildx imagetools inspect \
"public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" \
@@ -69,28 +108,12 @@ if target_enabled cuda-13-0; then
"vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-zstd"
fi
# ---- Ubuntu 24.04 (CUDA 13.0) ----
# ---- AWS EFA variant (docker/Dockerfile.efa) ----
# No cu134 entry: Rubin CUDA 13.4 is nightly-only and is not promoted to
# DockerHub for the base variant either.
if target_enabled cuda-13-0-ubuntu-24-04; then
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404"
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404"
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404" vllm/vllm-openai:latest-x86_64-ubuntu2404
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404"
docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404"
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404" vllm/vllm-openai:latest-aarch64-ubuntu2404
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404"
docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404"
docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404" || true
docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404"
docker manifest push vllm/vllm-openai:latest-ubuntu2404
docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404"
if target_enabled cuda-13-0-efa; then
publish_optional_family "-efa"
fi
# ---- ROCm ----
.buildkite/scripts/rocm/aiter-nightly-overlay.sh
+6 -2
View File
@@ -18,7 +18,7 @@
#
# Exit status:
# 10 no wheel fits this ci_base
# 11 the wheel does not install
# 11 the wheel does not install or its modules do not load
# Anything else is unexpected and passed through as is.
# Buildkite retries this step once on exit status 1 or 2, which it treats as
# transient Docker/registry failures (amd_infra_retry in the template).
@@ -78,6 +78,10 @@ ARG CI_BASE_IMAGE
FROM ${CI_BASE_IMAGE}
ARG AITER_WHEEL_URL
RUN python3 -m pip install --no-cache-dir "$AITER_WHEEL_URL"
# A wheel built against a newer glibc or libstdc++ than ci_base installs fine
# but fails every test that loads it; fail here instead.
RUN aiter="$(python3 -c 'import importlib.util as u; print(u.find_spec("aiter").origin)')" \
&& ! find "${aiter%/*}" -name '*.so' -exec ldd {} + 2>&1 | grep "version .* not found"
EOF
echo "--- :docker: Installing ${wheel_url##*/}"
@@ -90,7 +94,7 @@ if ! docker buildx build \
-t "${BUILD_TAG}" \
"${scratch}"; then
fail 11 "the wheel does not install" \
"Installing ${wheel_url##*/} over ${ci_base} failed; see this step's log."
"Installing or loading ${wheel_url##*/} over ${ci_base} failed; see this step's log."
fi
docker push "${BUILD_TAG}"
.buildkite/scripts/run-rust-frontend-cargo-ci.sh
+1 -1
View File
@@ -22,7 +22,7 @@ export RUSTUP_HOME="${RUSTUP_HOME:-$HOME/.rustup}"
export PATH="$CARGO_HOME/bin:$PATH"
CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}"
UV_VERSION="${UV_VERSION:-0.11.28}"
UV_VERSION="${UV_VERSION:-0.13.4}"
UV_SCRIPT_PYTHON_VERSION="${UV_SCRIPT_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
.buildkite/scripts/upload-release-wheels-pypi.sh
+1 -1
View File
@@ -41,7 +41,7 @@ set -x # avoid printing secrets above
# install uv if not already available
if ! command -v uv &> /dev/null; then
curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
curl -LsSf https://astral.sh/uv/0.13.4/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
fi
.buildkite/test-amd.yaml
+986 -963
View File
File diff suppressed because it is too large. Load diff
.buildkite/test_areas/benchmarks.yaml
+1 -1
View File
@@ -34,7 +34,7 @@ steps:
- benchmarks/attention_benchmarks/
- vllm/v1/attention/
commands:
- python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k"
- python3 benchmarks/attention_benchmarks/benchmark.py --backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k"
mirror:
amd:
label: ":amd: (MI355) Attention Benchmark Smoke"
.buildkite/test_areas/disaggregated.yaml
+491 -483
View File
@@ -1,209 +1,215 @@
group: Disaggregated
depends_on:
- image-build
steps:
- label: ":nvidia: (L4) [Nightly] Distributed NixlConnector PD accuracy"
key: distributed-nixlconnector-pd-accuracy-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 55
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- GPU_MEMORY_UTILIZATION=0.6 SLEEP_WAKE=1 bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI355) Distributed NixlConnector PD accuracy"
dind: false
device: mi355_4
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
steps: []
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] Distributed NixlConnector PD accuracy"
# key: distributed-nixlconnector-pd-accuracy-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 55
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# - GPU_MEMORY_UTILIZATION=0.6 SLEEP_WAKE=1 bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI355) Distributed NixlConnector PD accuracy"
# dind: false
# device: mi355_4
# timeout_in_minutes: 60
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: ":nvidia: (L4) [Nightly] Distributed FlashInfer NixlConnector PD accuracy"
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 55
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
dind: false
device: mi355_4
num_devices: 4
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- tests/v1/kv_connector/rocm_pd_accuracy_utils.py
- tests/utils.py
- vllm/v1/attention/backends/rocm_aiter_unified_attn.py
- vllm/platforms/rocm.py
- requirements/common.txt
- requirements/kv_connectors_rocm.txt
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# Preserve all six model/TP pairs with native AITER attention.
- pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] Distributed FlashInfer NixlConnector PD accuracy"
# key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 55
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
# dind: false
# device: mi355_4
# num_devices: 4
# timeout_in_minutes: 60
# working_dir: "/vllm-workspace/tests"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
# - tests/utils.py
# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
# - vllm/platforms/rocm.py
# - requirements/common.txt
# - requirements/kv_connectors_rocm.txt
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# # Preserve all six model/TP pairs with native AITER attention.
# - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
- label: ":nvidia: (L4) [Nightly] Push NixlConnector PP prefill PD accuracy"
key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- tests/v1/kv_connector/nixl_push_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy"
dind: false
device: mi300_4
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- tests/v1/kv_connector/nixl_push_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] Push NixlConnector PP prefill PD accuracy"
# key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 30
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - tests/v1/kv_connector/nixl_push_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy"
# dind: false
# device: mi300_4
# timeout_in_minutes: 55
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - tests/v1/kv_connector/nixl_push_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh
- label: ":nvidia: (L4) [Nightly] DP EP Distributed NixlConnector PD accuracy"
key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy"
dind: false
device: mi355_4
timeout_in_minutes: 40
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] DP EP Distributed NixlConnector PD accuracy"
# key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 30
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy"
# dind: false
# device: mi355_4
# timeout_in_minutes: 40
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: ":nvidia: (L4) [Nightly] CrossLayer KV layout Distributed NixlConnector PD accuracy"
key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 55
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy"
dind: false
device: mi300_4
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] CrossLayer KV layout Distributed NixlConnector PD accuracy"
# key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 55
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy"
# dind: false
# device: mi300_4
# timeout_in_minutes: 60
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: ":nvidia: (L4) [Nightly] Hybrid SSM NixlConnector PD accuracy"
key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy"
dind: false
device: mi300_4
timeout_in_minutes: 80
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] Hybrid SSM NixlConnector PD accuracy"
# key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 60
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy"
# dind: false
# device: mi300_4
# timeout_in_minutes: 80
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
# Per-PR core of the nightly-only 4-GPU sweeps above: hetero-TP (both
# directions), MoE DP-EP decode, push-mode PP prefill, sleep/wake, MLA,
@@ -211,296 +217,298 @@ steps:
# KV sharing on a 2-GPU pod. TP>1 configs share GPUs between prefill and
# decode (gpu-memory-utilization 0.35 per engine); TP1/TP1 architecture
# configs keep one dedicated GPU per side.
- label: ":nvidia: (L4) NixlConnector PD accuracy core"
key: nixlconnector-pd-accuracy-core-2-gpus
timeout_in_minutes: 90
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
# No AMD mirror: the MI300/MI355 mirrors above are optional: false and
# already cover these paths on PRs.
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- tests/v1/kv_connector/nixl_push_integration/
- requirements/kv_connectors.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_core_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) NixlConnector PD accuracy core"
# key: nixlconnector-pd-accuracy-core-2-gpus
# timeout_in_minutes: 90
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# # No AMD mirror: the MI300/MI355 mirrors above are optional: false and
# # already cover these paths on PRs.
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - tests/v1/kv_connector/nixl_push_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_core_test.sh
- label: ":nvidia: (L4) NixlConnector PD edge case"
key: nixlconnector-pd-edge-cases-2-gpus
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors.txt
env:
PREFILL_GPU_ID: "0"
DECODE_GPU_ID: "1"
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
mirror:
amd:
label: ":amd: (MI300) NixlConnector PD edge case"
dind: false
device: mi300_2
timeout_in_minutes: 65
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) NixlConnector PD edge case"
# key: nixlconnector-pd-edge-cases-2-gpus
# timeout_in_minutes: 40
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/core/sched/
# - tests/v1/kv_connector/nixl_integration/
# env:
# PREFILL_GPU_ID: "0"
# DECODE_GPU_ID: "1"
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
# mirror:
# amd:
# label: ":amd: (MI300) NixlConnector PD edge case"
# dind: false
# device: mi300_2
# timeout_in_minutes: 65
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/core/sched/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh
- label: ":nvidia: (L4) Hybrid SSM NixlConnector PD prefix cache"
key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
timeout_in_minutes: 25
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- vllm/v1/core/kv_cache_coordinator.py
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
mirror:
amd:
label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache"
dind: false
device: mi355_2
timeout_in_minutes: 25
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- vllm/v1/core/kv_cache_coordinator.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) Hybrid SSM NixlConnector PD prefix cache"
# key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
# timeout_in_minutes: 25
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/core/sched/
# - vllm/v1/core/kv_cache_coordinator.py
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache"
# dind: false
# device: mi355_2
# timeout_in_minutes: 25
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/core/sched/
# - vllm/v1/core/kv_cache_coordinator.py
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy"
key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
mirror:
amd:
label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD accuracy"
dind: false
device: mi355_2
timeout_in_minutes: 40
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors_rocm.txt
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy"
# key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
# timeout_in_minutes: 40
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD accuracy"
# dind: false
# device: mi355_2
# timeout_in_minutes: 40
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
# - tests/v1/kv_connector/nixl_integration/
# - requirements/kv_connectors_rocm.txt
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance"
key: nixlconnector-pd-spec-decode-acceptance-2-gpus
timeout_in_minutes: 45
device: a100
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
mirror:
amd:
label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance"
dind: false
device: mi355_2
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance"
# key: nixlconnector-pd-spec-decode-acceptance-2-gpus
# timeout_in_minutes: 45
# device: a100
# working_dir: "/vllm-workspace/tests"
# num_devices: 2
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/worker/kv_connector_model_runner_mixin.py
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance"
# dind: false
# device: mi355_2
# timeout_in_minutes: 45
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/v1/worker/kv_connector_model_runner_mixin.py
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD edge cases"
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
timeout_in_minutes: 25
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
mirror:
amd:
label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases"
dind: false
device: mi355_2
timeout_in_minutes: 25
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD edge cases"
# key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
# timeout_in_minutes: 25
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
# - tests/v1/kv_connector/nixl_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases"
# dind: false
# device: mi355_2
# timeout_in_minutes: 25
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
# P TP 4 - D DPEP 4 test case for DSv4-Flash
- label: ":nvidia: (H200) [Nightly] DSv4-Flash Disaggregated DP EP"
key: dsv4-flash-disaggregated
timeout_in_minutes: 60
device: h200
optional: true
working_dir: "/vllm-workspace/tests"
num_devices: 8
env:
ENABLE_HMA_FLAG: "1"
DP_EP: "1"
GPU_MEMORY_UTILIZATION: "0.85"
PREFILLER_TP_SIZE: "4"
DECODER_TP_SIZE: "4"
PREFILL_BLOCK_SIZE: "256"
DECODE_BLOCK_SIZE: "256"
MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
VLLM_ENGINE_READY_TIMEOUT_S: "1800"
VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
# The V4 indexer only supports block-outermost layouts; let it pick.
PREFILLER_KV_LAYOUT: ""
DECODER_KV_LAYOUT: ""
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
mirror:
amd:
label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP"
dind: false
device: mi355_8
timeout_in_minutes: 85
depends_on:
- image-build-amd
source_file_dependencies:
- requirements/kv_connectors_rocm.txt
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/models/deepseek_v4/
- vllm/v1/attention/
- tests/v1/kv_connector/nixl_integration/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# A mirror does not inherit the parent env, so it is repeated here in
# the parent's order. The layouts must stay empty: unset, the script
# defaults to HND, which the V4 indexer rejects.
- >-
VLLM_ROCM_USE_AITER=1
ENABLE_HMA_FLAG=1
DP_EP=1
GPU_MEMORY_UTILIZATION=0.85
PREFILLER_TP_SIZE=4
DECODER_TP_SIZE=4
PREFILL_BLOCK_SIZE=256
DECODE_BLOCK_SIZE=256
MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash
VLLM_ENGINE_READY_TIMEOUT_S=1800
VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8
PREFILLER_KV_LAYOUT=
DECODER_KV_LAYOUT=
bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (H200) [Nightly] DSv4-Flash Disaggregated DP EP"
# key: dsv4-flash-disaggregated
# timeout_in_minutes: 60
# device: h200
# optional: true
# working_dir: "/vllm-workspace/tests"
# num_devices: 8
# env:
# ENABLE_HMA_FLAG: "1"
# DP_EP: "1"
# GPU_MEMORY_UTILIZATION: "0.85"
# PREFILLER_TP_SIZE: "4"
# DECODER_TP_SIZE: "4"
# PREFILL_BLOCK_SIZE: "256"
# DECODE_BLOCK_SIZE: "256"
# MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
# VLLM_ENGINE_READY_TIMEOUT_S: "1800"
# VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
# # The V4 indexer only supports block-outermost layouts; let it pick.
# PREFILLER_KV_LAYOUT: ""
# DECODER_KV_LAYOUT: ""
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP"
# dind: false
# device: mi355_8
# timeout_in_minutes: 85
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - requirements/kv_connectors_rocm.txt
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/model_executor/layers/fused_moe/
# - vllm/model_executor/layers/quantization/
# - vllm/models/deepseek_v4/
# - vllm/v1/attention/
# - tests/v1/kv_connector/nixl_integration/
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# # A mirror does not inherit the parent env, so it is repeated here in
# # the parent's order. The layouts must stay empty: unset, the script
# # defaults to HND, which the V4 indexer rejects.
# - >-
# VLLM_ROCM_USE_AITER=1
# ENABLE_HMA_FLAG=1
# DP_EP=1
# GPU_MEMORY_UTILIZATION=0.85
# PREFILLER_TP_SIZE=4
# DECODER_TP_SIZE=4
# PREFILL_BLOCK_SIZE=256
# DECODE_BLOCK_SIZE=256
# MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash
# VLLM_ENGINE_READY_TIMEOUT_S=1800
# VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8
# PREFILLER_KV_LAYOUT=
# DECODER_KV_LAYOUT=
# bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
- label: ":nvidia: (H200) [Nightly] Kimi-Linear-48B-A3B Disaggregated DP EP"
key: kimi-linear-disaggregated
timeout_in_minutes: 60
device: h200
optional: true
working_dir: "/vllm-workspace/tests"
num_devices: 4
env:
ENABLE_HMA_FLAG: "1"
DP_EP: "1"
GPU_MEMORY_UTILIZATION: "0.9"
VLLM_SSM_CONV_STATE_LAYOUT: "DS"
PREFILLER_TP_SIZE: "2"
DECODER_TP_SIZE: "2"
MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct"
# Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima.
PREFILL_BLOCK_SIZE: "2048"
DECODE_BLOCK_SIZE: "2048"
VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768"
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
mirror:
amd:
label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
dind: false
device: mi355_4
num_devices: 4
timeout_in_minutes: 105
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/platforms/rocm.py
- vllm/models/kimi_k3/
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- requirements/kv_connectors_rocm.txt
env:
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (H200) [Nightly] Kimi-Linear-48B-A3B Disaggregated DP EP"
# key: kimi-linear-disaggregated
# timeout_in_minutes: 60
# device: h200
# optional: true
# working_dir: "/vllm-workspace/tests"
# num_devices: 4
# env:
# ENABLE_HMA_FLAG: "1"
# DP_EP: "1"
# GPU_MEMORY_UTILIZATION: "0.9"
# VLLM_SSM_CONV_STATE_LAYOUT: "DS"
# PREFILLER_TP_SIZE: "2"
# DECODER_TP_SIZE: "2"
# MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct"
# # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima.
# PREFILL_BLOCK_SIZE: "2048"
# DECODE_BLOCK_SIZE: "2048"
# VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768"
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP"
# dind: false
# device: mi355_4
# num_devices: 4
# timeout_in_minutes: 105
# working_dir: /vllm-workspace/tests
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/platforms/rocm.py
# - vllm/models/kimi_k3/
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - tests/v1/kv_connector/nixl_integration/
# - requirements/kv_connectors_rocm.txt
# env:
# VLLM_WORKER_MULTIPROC_METHOD: "spawn"
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
.buildkite/test_areas/disaggregated_mooncake.yaml
+45 -44
View File
@@ -2,50 +2,51 @@ group: Disaggregated Mooncake
depends_on:
- image-build
steps:
- label: ":nvidia: (B200) [Nightly] Distributed MooncakeConnector PD accuracy"
key: distributed-mooncakeconnector-pd-accuracy-4-gpus
# Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
# Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake
# transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to
# pin memory, and the only ones using host networking. Elsewhere topology
# discovery finds no usable HCA and the tcp fallback then advertises the
# docker bridge address, so every transfer fails to connect.
device: b200-k8s
num_devices: 4
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
- tests/v1/kv_connector/mooncake_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
mirror:
amd:
label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
optional: true
dind: false
device: mi355_4
num_devices: 4
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
- tests/v1/kv_connector/mooncake_integration/
- tests/v1/kv_connector/rocm_pd_accuracy_utils.py
- tests/utils.py
- examples/disaggregated/mooncake_connector/
- vllm/v1/attention/backends/rocm_aiter_unified_attn.py
- vllm/platforms/rocm.py
- requirements/common.txt
- requirements/kv_connectors_rocm.txt
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# Same full accuracy/TP sweep, using same-host HIP GPU transport.
- pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (B200) [Nightly] Distributed MooncakeConnector PD accuracy"
# key: distributed-mooncakeconnector-pd-accuracy-4-gpus
# # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 10
# working_dir: "/vllm-workspace/tests"
# # Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake
# # transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to
# # pin memory, and the only ones using host networking. Elsewhere topology
# # discovery finds no usable HCA and the tcp fallback then advertises the
# # docker bridge address, so every transfer fails to connect.
# device: b200-k8s
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
# - tests/v1/kv_connector/mooncake_integration/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
# mirror:
# amd:
# label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
# optional: true
# dind: false
# device: mi355_4
# num_devices: 4
# timeout_in_minutes: 30
# working_dir: "/vllm-workspace/tests"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
# - tests/v1/kv_connector/mooncake_integration/
# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
# - tests/utils.py
# - examples/disaggregated/mooncake_connector/
# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
# - vllm/platforms/rocm.py
# - requirements/common.txt
# - requirements/kv_connectors_rocm.txt
# commands:
# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
# # Same full accuracy/TP sweep, using same-host HIP GPU transport.
# - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
- label: ":nvidia: (L4) Mooncake EC TCP E2E"
key: mooncake-ec-tcp-e2e-2-gpus
.buildkite/test_areas/distributed.yaml
+87 -82
View File
@@ -13,6 +13,7 @@ steps:
- "!vllm/distributed/kv_transfer/"
- tests/distributed
commands:
- pytest -v -s distributed/test_localization.py
- pytest -v -s distributed/test_comm_ops.py
- pytest -v -s distributed/test_cpu_all_reduce.py
- pytest -v -s distributed/test_shm_broadcast.py
@@ -223,88 +224,89 @@ steps:
- HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: ":nvidia: (L4) [Nightly] Distributed Torchrun + Examples"
key: distributed-torchrun-examples-4-gpus
# Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 45
working_dir: "/vllm-workspace"
device: l4
num_devices: 4
source_file_dependencies:
- vllm/distributed/
- "!vllm/distributed/kv_transfer/"
- tests/distributed/test_torchrun_example.py
- tests/distributed/test_torchrun_example_moe.py
- examples/rl/
- examples/features/data_parallel/data_parallel_offline.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
# test with torchrun tp=2 and external_dp=2
- torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# test with torchrun tp=2 and pp=2
- PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# test with torchrun tp=4 and dp=1
- TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=2, pp=2 and dp=1
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=1 and dp=4 with ep
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with torchrun tp=2 and dp=2 with ep
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# test with internal dp
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
# rlhf examples (each launches its own `vllm serve` and tears it down)
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
# sharded RDT: 2 FSDP2 trainer ranks -> 2 vLLM DP ranks with EP. Its data
# plane is NIXL, which ships only in the kv-connector set, and the script
# also drops the nixl-cu wheel variant that does not match this image. Kept
# last so the commands above run in an unmodified environment. RDT_MODEL swaps
# the script's reference MoE for a tiny one: no weights are downloaded either
# way, but every trainer rank builds the model before sharding it, and the
# default does not fit an L4.
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
optional: false
label: ":amd: (MI355) Distributed Torchrun + Examples"
dind: false
device: mi355_4
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/
- "!vllm/distributed/kv_transfer/"
- tests/distributed/test_torchrun_example.py
- tests/distributed/test_torchrun_example_moe.py
- examples/rl/
- examples/features/data_parallel/data_parallel_offline.py
- .buildkite/scripts/install-kv-connectors.sh
- requirements/kv_connectors_rocm.txt
- requirements/test/rocm.in
- requirements/test/rocm.txt
- vllm/platforms/rocm.py
commands:
# Sharded RDT needs Ray >= 2.56; keep the ROCm image's dependencies.
- uv pip install --system --no-deps 'ray==2.56.1'
- export NCCL_CUMEM_HOST_ENABLE=0
- torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
- PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
- TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# Loading weights from the shared HF cache can exceed the default 300s --timeout.
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager --timeout 900
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (L4) [Nightly] Distributed Torchrun + Examples"
# key: distributed-torchrun-examples-4-gpus
# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 45
# working_dir: "/vllm-workspace"
# device: l4
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/
# - "!vllm/distributed/kv_transfer/"
# - tests/distributed/test_torchrun_example.py
# - tests/distributed/test_torchrun_example_moe.py
# - examples/rl/
# - examples/features/data_parallel/data_parallel_offline.py
# commands:
# # https://github.com/NVIDIA/nccl/issues/1838
# - export NCCL_CUMEM_HOST_ENABLE=0
# # test with torchrun tp=2 and external_dp=2
# - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# # test with torchrun tp=2 and pp=2
# - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# # test with torchrun tp=4 and dp=1
# - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# # test with torchrun tp=2, pp=2 and dp=1
# - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# # test with torchrun tp=1 and dp=4 with ep
# - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# # test with torchrun tp=2 and dp=2 with ep
# - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# # test with internal dp
# - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
# # rlhf examples (each launches its own `vllm serve` and tears it down)
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
# # sharded RDT: 2 FSDP2 trainer ranks -> 2 vLLM DP ranks with EP. Its data
# # plane is NIXL, which ships only in the kv-connector set, and the script
# # also drops the nixl-cu wheel variant that does not match this image. Kept
# # last so the commands above run in an unmodified environment. RDT_MODEL swaps
# # the script's reference MoE for a tiny one: no weights are downloaded either
# # way, but every trainer rank builds the model before sharding it, and the
# # default does not fit an L4.
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
# mirror:
# amd:
# # Run on PRs even when the L4 step is optional (nightly only).
# optional: false
# label: ":amd: (MI355) Distributed Torchrun + Examples"
# dind: false
# device: mi355_4
# timeout_in_minutes: 55
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/
# - "!vllm/distributed/kv_transfer/"
# - tests/distributed/test_torchrun_example.py
# - tests/distributed/test_torchrun_example_moe.py
# - examples/rl/
# - examples/features/data_parallel/data_parallel_offline.py
# - .buildkite/scripts/install-kv-connectors.sh
# - requirements/kv_connectors_rocm.txt
# - requirements/test/rocm.in
# - requirements/test/rocm.txt
# - vllm/platforms/rocm.py
# commands:
# # Sharded RDT needs Ray >= 2.56; keep the ROCm image's dependencies.
# - uv pip install --system --no-deps 'ray==2.56.1'
# - export NCCL_CUMEM_HOST_ENABLE=0
# - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
# - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
# # Loading weights from the shared HF cache can exceed the default 300s --timeout.
# - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager --timeout 900
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py
- label: ":nvidia: (L4) [Nightly] Distributed DP Extended"
key: distributed-dp-tests-4-gpus
@@ -681,10 +683,13 @@ steps:
- vllm/v1/worker/gpu_worker.py
- tests/distributed/
- tests/v1/distributed/test_pp_dp_v2.py
- vllm/benchmarks/forward/
- tests/benchmarks/test_forward_cli.py
commands:
- pytest -v -s distributed/test_pp_cudagraph.py
- pytest -v -s distributed/test_pipeline_parallel.py
- pytest -v -s v1/distributed/test_pp_dp_v2.py
- pytest -v -s benchmarks/test_forward_cli.py::test_bench_forward_pipeline
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
.buildkite/test_areas/docker.yaml
+8
View File
@@ -9,12 +9,17 @@ steps:
source_file_dependencies:
- .buildkite/release-pipeline.yaml
- .buildkite/scripts/docker-build-metadata-args.sh
- .buildkite/scripts/rocm/build-config.sh
- docker/Dockerfile
- docker/Dockerfile.cpu
- docker/docker-bake.hcl
- tests/tools/test_docker_build_metadata_args.py
- tests/tools/test_kernel_symbol_map.py
- tools/ci/kernel_symbol_map.py
- tools/ci/rocm_kernel_symbol_map.py
commands:
- pytest -v -s tools/test_docker_build_metadata_args.py
- pytest -v -s tools/test_kernel_symbol_map.py
mirror:
amd:
label: ":amd: (MI250) Docker Build Metadata"
@@ -32,3 +37,6 @@ steps:
- docker/ci-rocm.hcl
- docker/docker-bake-rocm.hcl
- tests/tools/test_docker_build_metadata_args.py
- tests/tools/test_kernel_symbol_map.py
- tools/ci/kernel_symbol_map.py
- tools/ci/rocm_kernel_symbol_map.py
.buildkite/test_areas/expert_parallelism.yaml
+19 -18
View File
@@ -60,21 +60,22 @@ steps:
- tests/distributed/test_eplb_spec_decode.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H100) [Nightly] Elastic EP Scaling"
key: elastic-ep-scaling-test
# Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs.
optional: true
timeout_in_minutes: 40
device: h100
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/distributed/
- "!vllm/distributed/kv_transfer/"
- vllm/engine/
- vllm/v1/executor/
- vllm/compilation/
- tests/distributed/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- pytest -v -s distributed/test_elastic_ep.py
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (H100) [Nightly] Elastic EP Scaling"
# key: elastic-ep-scaling-test
# # Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs.
# optional: true
# timeout_in_minutes: 40
# device: h100
# working_dir: "/vllm-workspace/tests"
# num_devices: 4
# source_file_dependencies:
# - vllm/distributed/
# - "!vllm/distributed/kv_transfer/"
# - vllm/engine/
# - vllm/v1/executor/
# - vllm/compilation/
# - tests/distributed/
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - pytest -v -s distributed/test_elastic_ep.py
.buildkite/test_areas/fault_tolerance.yaml
+82 -42
View File
@@ -2,13 +2,85 @@ group: Fault Tolerance
depends_on:
- image-build
steps:
- label: ":nvidia: (H100) Fault Tolerance E2E"
key: fault-tolerance-e2e-2xh100
timeout_in_minutes: 35
device: h100
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (H100) Fault Tolerance E2E"
# key: fault-tolerance-e2e-2xh100
# timeout_in_minutes: 35
# device: h100
# num_devices: 2
# working_dir: "/vllm-workspace/tests"
# source_file_dependencies:
# - vllm/v1/fault_tolerance/
# - vllm/v1/worker/sentinel/
# - vllm/entrypoints/serve/fault_tolerance/
# - vllm/distributed/elastic_ep/
# - vllm/distributed/device_communicators/
# - vllm/distributed/kv_transfer/
# - vllm/v1/engine/
# - vllm/v1/executor/
# - vllm/v1/worker/
# - tests/v1/fault_tolerance/
# - tests/v1/distributed/test_external_lb_dp.py
# commands:
# # Base image has no nixl; install it or has_nixl_ep() skips the tests.
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# # https://github.com/NVIDIA/nccl/issues/1838
# - export NCCL_CUMEM_HOST_ENABLE=0
# - pytest -v -s v1/fault_tolerance/test_engine_core_sentinel.py
# - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
# mirror:
# amd:
# label: ":amd: (MI355) MoRI EP Healthy Baseline"
# dind: false
# device: mi355_2
# num_devices: 2
# timeout_in_minutes: 10
# working_dir: "/vllm-workspace/tests"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - csrc/moe/
# - csrc/rocm/
# - tests/kernels/moe/
# - vllm/model_executor/layers/fused_moe/
# - vllm/model_executor/layers/quantization/
# - vllm/distributed/
# - vllm/config/
# - vllm/forward_context.py
# - vllm/v1/worker/workspace.py
# - vllm/utils/import_utils.py
# - vllm/utils/math_utils.py
# - vllm/utils/torch_utils.py
# - vllm/platforms/
# - vllm/_aiter_ops.py
# # TODO: Add fault injection, detection and recovery once MoRI supports them.
# # This mirror currently checks healthy EP execution only.
# commands:
# - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
# - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
# Keep the AMD job: its commands do not run this installation.
- label: ':amd: (MI355) MoRI EP Healthy Baseline'
key: amd-fault-tolerance-e2e-2xh100
timeout_in_minutes: 10
device: mi355_2
num_devices: 2
working_dir: "/vllm-workspace/tests"
working_dir: /vllm-workspace/tests
source_file_dependencies:
- csrc/moe/
- csrc/rocm/
- tests/kernels/moe/
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/
- vllm/config/
- vllm/forward_context.py
- vllm/v1/worker/workspace.py
- vllm/utils/import_utils.py
- vllm/utils/math_utils.py
- vllm/utils/torch_utils.py
- vllm/platforms/
- vllm/_aiter_ops.py
- vllm/v1/fault_tolerance/
- vllm/v1/worker/sentinel/
- vllm/entrypoints/serve/fault_tolerance/
@@ -20,41 +92,9 @@ steps:
- vllm/v1/worker/
- tests/v1/fault_tolerance/
- tests/v1/distributed/test_external_lb_dp.py
- requirements/kv_connectors.txt
commands:
# Base image has no nixl; install it or has_nixl_ep() skips the tests.
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
- pytest -v -s v1/fault_tolerance/test_engine_core_sentinel.py
- pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
mirror:
amd:
label: ":amd: (MI355) MoRI EP Healthy Baseline"
dind: false
device: mi355_2
num_devices: 2
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/moe/
- csrc/rocm/
- tests/kernels/moe/
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/
- vllm/config/
- vllm/forward_context.py
- vllm/v1/worker/workspace.py
- vllm/utils/import_utils.py
- vllm/utils/math_utils.py
- vllm/utils/torch_utils.py
- vllm/platforms/
- vllm/_aiter_ops.py
# TODO: Add fault injection, detection and recovery once MoRI supports them.
# This mirror currently checks healthy EP execution only.
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
- VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
- VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
dind: false
depends_on:
- image-build-amd
.buildkite/test_areas/kernels.yaml
+65 -63
View File
@@ -48,7 +48,7 @@ steps:
dind: false
automatic_shard: true
device: mi300_1
timeout_in_minutes: 45
timeout_in_minutes: 90
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
@@ -744,69 +744,71 @@ steps:
commands:
- pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py
- label: ":nvidia: (DGX Spark) [Nightly] b12x Linear Kernels"
key: b12x-linear-kernels-dgx-spark-nightly
timeout_in_minutes: 60
device: dgx-spark
optional: true
num_devices: 1
depends_on:
- arm64-image-build
source_file_dependencies:
- setup.py
- vllm/config/kernel.py
- vllm/envs.py
- vllm/model_executor/kernels/linear/
- vllm/model_executor/layers/fused_moe/b12x.py
- vllm/model_executor/layers/fused_moe/config.py
- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
- vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/
- vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
- vllm/model_executor/layers/quantization/modelopt.py
- vllm/model_executor/layers/quantization/mxfp4.py
- vllm/model_executor/layers/quantization/utils/b12x_moe.py
- vllm/model_executor/warmup/b12x_warmup.py
- vllm/utils/b12x.py
- tests/kernels/moe/test_b12x.py
- tests/kernels/quantization/nvfp4_utils.py
- tests/model_executor/kernels/test_b12x_linear.py
- tests/model_executor/test_b12x_warmup.py
- tests/kernels/quantization/test_block_fp8.py
- tests/quantization/reference_mxfp4.py
- tests/quantization/test_auto_round.py
commands:
- uv pip install --system b12x==1.5.1
- pytest -v -s model_executor/kernels/test_b12x_linear.py
model_executor/test_b12x_warmup.py
kernels/quantization/test_block_fp8.py
kernels/moe/test_b12x.py
quantization/test_auto_round.py -k b12x
# Disabled until b12x transitive dependencies are locked.
# - label: ":nvidia: (DGX Spark) [Nightly] b12x Linear Kernels"
# key: b12x-linear-kernels-dgx-spark-nightly
# timeout_in_minutes: 60
# device: dgx-spark
# optional: true
# num_devices: 1
# depends_on:
# - arm64-image-build
# source_file_dependencies:
# - setup.py
# - vllm/config/kernel.py
# - vllm/envs.py
# - vllm/model_executor/kernels/linear/
# - vllm/model_executor/layers/fused_moe/b12x.py
# - vllm/model_executor/layers/fused_moe/config.py
# - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
# - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
# - vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/
# - vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
# - vllm/model_executor/layers/quantization/modelopt.py
# - vllm/model_executor/layers/quantization/mxfp4.py
# - vllm/model_executor/layers/quantization/utils/b12x_moe.py
# - vllm/model_executor/warmup/b12x_warmup.py
# - vllm/utils/b12x.py
# - tests/kernels/moe/test_b12x.py
# - tests/kernels/quantization/nvfp4_utils.py
# - tests/model_executor/kernels/test_b12x_linear.py
# - tests/model_executor/test_b12x_warmup.py
# - tests/kernels/quantization/test_block_fp8.py
# - tests/quantization/reference_mxfp4.py
# - tests/quantization/test_auto_round.py
# commands:
# - uv pip install --system b12x==1.5.1
# - pytest -v -s model_executor/kernels/test_b12x_linear.py
# model_executor/test_b12x_warmup.py
# kernels/quantization/test_block_fp8.py
# kernels/moe/test_b12x.py
# quantization/test_auto_round.py -k b12x
- label: ":nvidia: (H100) Helion Kernels Shard %N"
key: kernels-helion-test
timeout_in_minutes: 115
device: h100
source_file_dependencies:
- vllm/utils/import_utils.py
- tests/kernels/helion/
commands:
- pip install helion==1.4.0
- pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 5
mirror:
amd:
label: ":amd: (MI250) Helion Kernels Shard %N"
dind: false
device: mi250_1
timeout_in_minutes: 45
# Explicit so the base step's count does not multiply the mirror's replicas.
parallelism: 2
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/platforms/rocm.py
- tools/benchmark_helion_kernels.py
# Disabled until helion transitive dependencies are locked.
# - label: ":nvidia: (H100) Helion Kernels Shard %N"
# key: kernels-helion-test
# timeout_in_minutes: 115
# device: h100
# source_file_dependencies:
# - vllm/utils/import_utils.py
# - tests/kernels/helion/
# commands:
# - pip install helion==1.4.0
# - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
# parallelism: 5
# mirror:
# amd:
# label: ":amd: (MI250) Helion Kernels Shard %N"
# dind: false
# device: mi250_1
# timeout_in_minutes: 45
# # Explicit so the base step's count does not multiply the mirror's replicas.
# parallelism: 2
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/platforms/rocm.py
# - tools/benchmark_helion_kernels.py
- label: ":nvidia: (H100) [Nightly] FP8 MoE Kernels"
.buildkite/test_areas/lm_eval.yaml
+77 -54
View File
@@ -1121,6 +1121,27 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt
# Octave kernels are built only for gfx942 and gfx1100, so this runs on MI300.
- label: ":amd: (MI300) LM Eval Octave KV Cache"
key: amd-lm-eval-octave-mi300
dind: false
device: mi300_1
num_devices: 1
timeout_in_minutes: 45
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- tests/evals/gsm8k/
- tests/utils.py
- tests/conftest.py
- vllm/v1/attention/backends/rocm_octave_attn.py
- vllm/v1/attention/ops/rocm_octave.py
- vllm/v1/attention/selector.py
- vllm/platforms/rocm.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt
- label: ":nvidia: (H100) [Nightly] GPQA Eval (GPT-OSS)"
key: gpqa-eval-gpt-oss-2xh100
timeout_in_minutes: 35
@@ -1191,59 +1212,61 @@ steps:
commands:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
- label: ":nvidia: (B200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval"
key: accuracy-eval-qwen3-8-flash-next-fp8-b200
timeout_in_minutes: 60
device: b200-k8s
optional: true
num_devices: 4
source_file_dependencies:
- csrc/gdn/
- tests/evals/qwen4_exp/
- vllm/model_executor/layers/mamba/
- vllm/models/qwen4_exp/
- vllm/v1/attention/backends/short_conv_attn.py
commands:
- uv pip install --system 'evalscope==1.10.0'
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-b200.txt
# Disabled until evalscope transitive dependencies are locked.
# - label: ":nvidia: (B200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval"
# key: accuracy-eval-qwen3-8-flash-next-fp8-b200
# timeout_in_minutes: 60
# device: b200-k8s
# optional: true
# num_devices: 4
# source_file_dependencies:
# - csrc/gdn/
# - tests/evals/qwen4_exp/
# - vllm/model_executor/layers/mamba/
# - vllm/models/qwen4_exp/
# - vllm/v1/attention/backends/short_conv_attn.py
# commands:
# - uv pip install --system 'evalscope==1.10.0'
# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-b200.txt
- label: ":nvidia: (H200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval"
key: accuracy-eval-qwen3-8-flash-next-fp8-h200
timeout_in_minutes: 60
device: h200
optional: true
num_devices: 4
source_file_dependencies:
- csrc/gdn/
- tests/evals/qwen4_exp/
- vllm/model_executor/layers/mamba/
- vllm/models/qwen4_exp/
- vllm/v1/attention/backends/short_conv_attn.py
commands:
- uv pip install --system 'evalscope==1.10.0'
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt
mirror:
amd:
label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
dind: false
device: mi355_4
num_devices: 4
timeout_in_minutes: 105
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/gdn/
- tests/evals/qwen4_exp/
- vllm/model_executor/layers/mamba/
- vllm/models/qwen4_exp/
- vllm/v1/attention/backends/short_conv_attn.py
- vllm/platforms/rocm.py
env:
VLLM_WORKER_MULTIPROC_METHOD: "spawn"
commands:
- uv pip install --system 'evalscope==1.10.0'
- pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
# Disabled until evalscope transitive dependencies are locked.
# - label: ":nvidia: (H200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval"
# key: accuracy-eval-qwen3-8-flash-next-fp8-h200
# timeout_in_minutes: 60
# device: h200
# optional: true
# num_devices: 4
# source_file_dependencies:
# - csrc/gdn/
# - tests/evals/qwen4_exp/
# - vllm/model_executor/layers/mamba/
# - vllm/models/qwen4_exp/
# - vllm/v1/attention/backends/short_conv_attn.py
# commands:
# - uv pip install --system 'evalscope==1.10.0'
# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt
# mirror:
# amd:
# label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
# dind: false
# device: mi355_4
# num_devices: 4
# timeout_in_minutes: 105
# working_dir: /vllm-workspace/tests
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - csrc/gdn/
# - tests/evals/qwen4_exp/
# - vllm/model_executor/layers/mamba/
# - vllm/models/qwen4_exp/
# - vllm/v1/attention/backends/short_conv_attn.py
# - vllm/platforms/rocm.py
# env:
# VLLM_WORKER_MULTIPROC_METHOD: "spawn"
# commands:
# - uv pip install --system 'evalscope==1.10.0'
# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt
- label: ":nvidia: (DGX Spark) [Nightly] GPQA Eval (GPT-OSS)"
key: gpqa-eval-gpt-oss-spark
@@ -1326,9 +1349,9 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
mirror:
amd:
label: ":amd: (MI300) LM Eval KV-Offload Medium"
label: ":amd: (MI355) LM Eval KV-Offload Medium"
dind: false
device: mi300_2
device: mi355_2
timeout_in_minutes: 70
depends_on:
- image-build-amd
.buildkite/test_areas/misc.yaml
+431 -337
View File
@@ -144,51 +144,53 @@ steps:
depends_on:
- image-build-amd
- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors"
device: h200_35gb
key: v1-kv-connectors
automatic_shard: true
timeout_in_minutes: 45
source_file_dependencies: *v1-core-kv-metrics-deps
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
mirror:
amd:
label: ":amd: (MI355 DPX) V1 KV Connectors"
automatic_shard: true
dind: false
device: mi355_dpx
timeout_in_minutes: 60
depends_on:
- image-build-amd
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors"
# device: h200_35gb
# key: v1-kv-connectors
# automatic_shard: true
# timeout_in_minutes: 45
# source_file_dependencies: *v1-core-kv-metrics-deps
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
# - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
# mirror:
# amd:
# label: ":amd: (MI355 DPX) V1 KV Connectors"
# automatic_shard: true
# dind: false
# device: mi355_dpx
# timeout_in_minutes: 60
# depends_on:
# - image-build-amd
# The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from
# their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole
# MoRIIO suite on ROCm whenever the connector or its tests change.
- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests"
key: amd-moriio-connector-unit-mi355
dind: false
device: mi355_dpx
num_devices: 1
timeout_in_minutes: 30
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
- tests/v1/kv_connector/unit/test_moriio_
- tests/v1/kv_connector/unit/utils.py
- requirements/kv_connectors_rocm.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py
# Disabled until the KV connectors installation is locked.
# - label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests"
# key: amd-moriio-connector-unit-mi355
# dind: false
# device: mi355_dpx
# num_devices: 1
# timeout_in_minutes: 30
# working_dir: /vllm-workspace/tests
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/
# - vllm/distributed/kv_transfer/kv_connector/v1/base.py
# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
# - tests/v1/kv_connector/unit/test_moriio_
# - tests/v1/kv_connector/unit/utils.py
# - requirements/kv_connectors_rocm.txt
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py
# UMBP is built on mori, so its tests skip on CUDA. Run them on MI355 whenever
# the connector or its tests change.
@@ -211,56 +213,59 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py
# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate
# GPUs over xGMI, compared exactly against a direct serve.
- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N"
key: amd-moriio-write-dummy-e2e-mi355
dind: false
device: mi355_2
num_devices: 2
# ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the
# largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports
# for its servers and proxy and its own compile-cache dir.
parallelism: 6
timeout_in_minutes: 45
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
- vllm/model_executor/model_loader/dummy_loader.py
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
- tests/v1/kv_connector/moriio_integration/
- requirements/kv_connectors_rocm.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
# Single-node MoRIIO P/D on dummy weights, WRITE and READ mode: prefill and
# decode on separate GPUs over xGMI, compared exactly against a direct serve.
# Disabled until the KV connectors installation is locked.
# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N"
# key: amd-moriio-write-dummy-e2e-mi355
# dind: false
# device: mi355_2
# num_devices: 2
# # ~3.5-5 min per case; pytest-shard hashes test ids, so with 28 cases the
# # largest of 6 shards holds 7 (~32 min on MI355). Every case picks free ports
# # for its servers and proxy and its own compile-cache dir.
# parallelism: 6
# timeout_in_minutes: 45
# working_dir: /vllm-workspace/tests
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/
# - vllm/distributed/kv_transfer/kv_connector/v1/base.py
# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
# - vllm/model_executor/model_loader/dummy_loader.py
# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
# - tests/v1/kv_connector/moriio_integration/
# - requirements/kv_connectors_rocm.txt
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
# The same P/D with prefill TP2 writing into decode TP1, which needs 3 GPUs.
- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP"
key: amd-moriio-write-dummy-hetero-tp-mi355
dind: false
device: mi355_4
num_devices: 4
timeout_in_minutes: 45
working_dir: /vllm-workspace/tests
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/moriio/
- vllm/distributed/kv_transfer/kv_connector/v1/base.py
- vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
- vllm/model_executor/model_loader/dummy_loader.py
- examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
- tests/v1/kv_connector/moriio_integration/
- requirements/kv_connectors_rocm.txt
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp
# The same P/D between different TP sizes (WRITE TP2 -> TP1, READ of MLA
# TP1 <-> TP2), which needs 3 GPUs: 8 cases, ~36 min on MI355.
# Disabled until the KV connectors installation is locked.
# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP"
# key: amd-moriio-write-dummy-hetero-tp-mi355
# dind: false
# device: mi355_4
# num_devices: 4
# timeout_in_minutes: 45
# working_dir: /vllm-workspace/tests
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/
# - vllm/distributed/kv_transfer/kv_connector/v1/base.py
# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
# - vllm/model_executor/model_loader/dummy_loader.py
# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
# - tests/v1/kv_connector/moriio_integration/
# - requirements/kv_connectors_rocm.txt
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp
- label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval"
device: h200_35gb
@@ -319,6 +324,7 @@ steps:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s v1/worker/test_gpu_worker.py::test_cpu_worker_without_triton
- pytest -v -s v1/worker/test_cpu_worker_kv_budget.py
- pytest -v -s v1/test_kv_cache_spec_registry.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
@@ -407,37 +413,38 @@ steps:
- tests/v1/kv_connector/extract_hidden_states_integration/
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) Regression"
key: regression
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- "!vllm/distributed/kv_transfer/"
- vllm/engine/
- vllm/inputs/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/test_regression
commands:
- pip install 'modelscope<1.38'
- pytest -v -s test_regression.py
working_dir: "/vllm-workspace/tests" # optional
mirror:
amd:
label: ":amd: (MI355 DPX) Regression"
dind: false
device: mi355_dpx
timeout_in_minutes: 25
depends_on:
- image-build-amd
working_dir: /vllm-workspace/tests
# Disabled until the modelscope installation is locked.
# - label: ":nvidia: (H200 MIG 18GB) Regression"
# key: regression
# timeout_in_minutes: 30
# device: h200_18gb
# source_file_dependencies:
# - vllm/config/
# - vllm/distributed/
# - "!vllm/distributed/kv_transfer/"
# - vllm/engine/
# - vllm/inputs/
# - vllm/model_executor/
# - vllm/multimodal/
# - vllm/platforms/
# - vllm/sampling_params.py
# - vllm/transformers_utils/
# - vllm/utils/
# - vllm/v1/
# - tests/test_regression
# commands:
# - pip install 'modelscope<1.38'
# - pytest -v -s test_regression.py
# working_dir: "/vllm-workspace/tests" # optional
# mirror:
# amd:
# label: ":amd: (MI355 DPX) Regression"
# dind: false
# device: mi355_dpx
# timeout_in_minutes: 25
# depends_on:
# - image-build-amd
# working_dir: /vllm-workspace/tests
- label: ":nvidia: (H200 MIG 35GB) Examples"
device: h200_35gb
@@ -485,43 +492,44 @@ steps:
depends_on:
- image-build-amd
- label: ":nvidia: (L4) Metrics, Tracing"
key: metrics-tracing-2-gpus
timeout_in_minutes: 25
device: l4
num_devices: 2
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- "!vllm/distributed/kv_transfer/"
- vllm/engine/
- vllm/inputs/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/tracing/
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/tracing
- tests/tracing/
commands:
- "pip install \
'opentelemetry-sdk>=1.26.0' \
'opentelemetry-api>=1.26.0' \
'opentelemetry-exporter-otlp>=1.26.0' \
'opentelemetry-semantic-conventions-ai>=0.4.1'"
- pytest -v -s v1/tracing
- pytest -v -s tracing
mirror:
amd:
label: ":amd: (MI300) Metrics, Tracing"
dind: false
device: mi300_2
timeout_in_minutes: 30
depends_on:
- image-build-amd
# Disabled until the tracing packages installation is locked.
# - label: ":nvidia: (L4) Metrics, Tracing"
# key: metrics-tracing-2-gpus
# timeout_in_minutes: 25
# device: l4
# num_devices: 2
# source_file_dependencies:
# - vllm/config/
# - vllm/distributed/
# - "!vllm/distributed/kv_transfer/"
# - vllm/engine/
# - vllm/inputs/
# - vllm/model_executor/
# - vllm/multimodal/
# - vllm/platforms/
# - vllm/sampling_params.py
# - vllm/tracing/
# - vllm/transformers_utils/
# - vllm/utils/
# - vllm/v1/
# - tests/v1/tracing
# - tests/tracing/
# commands:
# - "pip install \
# 'opentelemetry-sdk>=1.26.0' \
# 'opentelemetry-api>=1.26.0' \
# 'opentelemetry-exporter-otlp>=1.26.0' \
# 'opentelemetry-semantic-conventions-ai>=0.4.1'"
# - pytest -v -s v1/tracing
# - pytest -v -s tracing
# mirror:
# amd:
# label: ":amd: (MI300) Metrics, Tracing"
# dind: false
# device: mi300_2
# timeout_in_minutes: 30
# depends_on:
# - image-build-amd
- label: ":nvidia: (H200 MIG 18GB) [Nightly] Python-only Installation"
key: python-only-installation
@@ -599,8 +607,129 @@ steps:
- pytest -v -s multimodal/test_cache.py::test_sleep_wake_preserves_mm_cache_consistency
- pytest -v -s utils_
- label: ":computer: (CPU) Params, Env, Tokenizers, Parser"
key: cpu-params-env-tokenizers-parser
# Disabled until the cmake installation is locked.
# - label: ":computer: (CPU) Params, Env, Tokenizers, Parser"
# key: cpu-params-env-tokenizers-parser
# depends_on:
# - image-build-cpu
# timeout_in_minutes: 35
# source_file_dependencies: &cpu-async-engine-deps
# - vllm/assets/
# - vllm/config/
# - vllm/device_allocator/
# - vllm/connections.py
# - vllm/engine/arg_utils.py
# - vllm/renderers/chat_utils/
# - vllm/entrypoints/mcp/
# - vllm/entrypoints/openai/chat_completion/protocol.py
# - vllm/entrypoints/openai/completion/protocol.py
# - vllm/entrypoints/openai/responses/protocol.py
# - vllm/entrypoints/generate/base/protocol.py
# - vllm/env_override.py
# - vllm/envs.py
# - vllm/exceptions.py
# - vllm/inputs/
# - vllm/logging_utils/
# - vllm/logprobs.py
# - vllm/model_executor/layers/quantization/quark/
# - vllm/model_executor/model_loader/tensorizer.py
# - vllm/models/muse_spark/
# - vllm/multimodal/
# - vllm/outputs.py
# - vllm/parser/
# - vllm/platforms/
# - vllm/pooling_params.py
# - vllm/ray/
# - vllm/reasoning/
# - vllm/renderers/
# - vllm/sampling_params.py
# - vllm/scalar_type.py
# - vllm/tokenizers/
# - vllm/tool_parsers/
# - vllm/transformers_utils/
# - vllm/triton_utils/
# - vllm/utils/
# - vllm/v1/
# - vllm/version.py
# - cmake/utils.cmake
# - cmake/hipify.py
# - tests/test_uvicorn_logging.py
# - tests/test_attention_backend_registry.py
# - tests/test_audio_media_size_precheck.py
# - tests/test_cmake_utils.py
# - tests/test_connections.py
# - tests/test_embedded_commit.py
# - tests/test_envs.py
# - tests/test_force_first_config.py
# - tests/test_fxgraphcache_pickle_patch.py
# - tests/test_logprobs.py
# - tests/test_outputs.py
# - tests/test_pcp_dp.py
# - tests/test_pooling_params.py
# - tests/test_quantization_revision_pin.py
# - tests/test_ray_env.py
# - tests/test_ray_env_utils.py
# - tests/test_request_input_bounds.py
# - tests/test_sampling_params.py
# - tests/test_scalartype.py
# - tests/test_tensorizer_revision_pin.py
# - tests/test_triton_utils.py
# - tests/test_version.py
# - tests/test_zen_cpu_platform_detection.py
# - tests/multimodal
# - tests/renderers
# - tests/standalone_tests/lazy_imports.py
# - tests/standalone_tests/test_precompiled_wheel.py
# - setup.py
# - tests/reasoning
# - tools/pre_commit/check_test_tethering.py
# - tools/pre_commit/test_tethering_allowlist.txt
# - tests/tools/test_check_test_tethering.py
# - tools/pre_commit/check_agent_files.py
# - tests/tools/test_check_agent_files.py
# - tools/pre_commit/check_trust_remote_code.py
# - tests/tools/test_check_trust_remote_code.py
# - tests/tool_parsers
# - tests/tokenizers_
# - tests/parser
# - tests/transformers_utils
# - tests/config
# - tests/device_allocator
# device: cpu-small
# commands:
# - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none
# - python3 standalone_tests/lazy_imports.py
# - python3 standalone_tests/test_precompiled_wheel.py
# - pytest -v -s device_allocator
# - pytest -v -s test_uvicorn_logging.py
# - pytest -v -s test_attention_backend_registry.py
# - pytest -v -s test_audio_media_size_precheck.py
# - pytest -v -s test_cmake_utils.py
# - pytest -v -s test_connections.py
# - pytest -v -s test_embedded_commit.py
# - pytest -v -s test_envs.py
# - pytest -v -s test_force_first_config.py
# - pytest -v -s test_fxgraphcache_pickle_patch.py
# - pytest -v -s test_logprobs.py
# - pytest -v -s test_outputs.py
# - pytest -v -s test_pcp_dp.py
# - pytest -v -s test_pooling_params.py
# - pytest -v -s test_quantization_revision_pin.py
# - pytest -v -s test_ray_env.py
# - pytest -v -s test_ray_env_utils.py
# - pytest -v -s test_request_input_bounds.py
# - pytest -v -s test_sampling_params.py
# - pytest -v -s test_scalartype.py
# - pytest -v -s test_tensorizer_revision_pin.py
# - pytest -v -s test_triton_utils.py
# - pytest -v -s test_version.py
# - pytest -v -s test_zen_cpu_platform_detection.py
# - pytest -v -s tokenizers_
# - pytest -v -s parser
# - pytest -v -s transformers_utils
- label: ":computer: (CPU) Multimodal + Config"
key: cpu-multimodal-config
depends_on:
- image-build-cpu
timeout_in_minutes: 35
@@ -643,6 +772,7 @@ steps:
- vllm/v1/
- vllm/version.py
- cmake/utils.cmake
- cmake/hipify.py
- tests/test_uvicorn_logging.py
- tests/test_attention_backend_registry.py
- tests/test_audio_media_size_precheck.py
@@ -687,45 +817,6 @@ steps:
- tests/device_allocator
device: cpu-small
commands:
- pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none
- python3 standalone_tests/lazy_imports.py
- python3 standalone_tests/test_precompiled_wheel.py
- pytest -v -s device_allocator
- pytest -v -s test_uvicorn_logging.py
- pytest -v -s test_attention_backend_registry.py
- pytest -v -s test_audio_media_size_precheck.py
- pytest -v -s test_cmake_utils.py
- pytest -v -s test_connections.py
- pytest -v -s test_embedded_commit.py
- pytest -v -s test_envs.py
- pytest -v -s test_force_first_config.py
- pytest -v -s test_fxgraphcache_pickle_patch.py
- pytest -v -s test_logprobs.py
- pytest -v -s test_outputs.py
- pytest -v -s test_pcp_dp.py
- pytest -v -s test_pooling_params.py
- pytest -v -s test_quantization_revision_pin.py
- pytest -v -s test_ray_env.py
- pytest -v -s test_ray_env_utils.py
- pytest -v -s test_request_input_bounds.py
- pytest -v -s test_sampling_params.py
- pytest -v -s test_scalartype.py
- pytest -v -s test_tensorizer_revision_pin.py
- pytest -v -s test_triton_utils.py
- pytest -v -s test_version.py
- pytest -v -s test_zen_cpu_platform_detection.py
- pytest -v -s tokenizers_
- pytest -v -s parser
- pytest -v -s transformers_utils
- label: ":computer: (CPU) Multimodal + Config"
key: cpu-multimodal-config
depends_on:
- image-build-cpu
timeout_in_minutes: 35
source_file_dependencies: *cpu-async-engine-deps
device: cpu-small
commands:
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s config
- pytest -v -s tools/test_check_test_tethering.py
@@ -753,146 +844,149 @@ steps:
commands:
- pytest -v -s tool_parsers
- label: ":nvidia: (A100) Batch Invariance"
key: batch-invariance-a100
timeout_in_minutes: 60
device: a100
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
mirror:
amd:
label: ":amd: (MI250) Batch Invariance"
dind: false
device: mi250_1
timeout_in_minutes: 60
env:
VLLM_GPU_MEMORY_UTILIZATION: "0.8"
VLLM_TEST_CLEAN_GPU_MEMORY: "1"
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
# Disabled until the pytest plugins installation is locked.
# - label: ":nvidia: (A100) Batch Invariance"
# key: batch-invariance-a100
# timeout_in_minutes: 60
# device: a100
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pip install pytest-timeout pytest-forked
# - pytest -v -s v1/determinism/test_batch_invariance.py
# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
# mirror:
# amd:
# label: ":amd: (MI250) Batch Invariance"
# dind: false
# device: mi250_1
# timeout_in_minutes: 60
# env:
# VLLM_GPU_MEMORY_UTILIZATION: "0.8"
# VLLM_TEST_CLEAN_GPU_MEMORY: "1"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
- label: ":nvidia: (H100) Batch Invariance"
key: batch-invariance-h100
automatic_shard: true
timeout_in_minutes: 60
device: h100
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/lora/ops/triton_ops
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- pytest -v -s v1/determinism/test_lora_shrink_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 -k 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 -k FLASH_ATTN
mirror:
amd:
label: ":amd: (MI300) Batch Invariance"
dind: false
automatic_shard: true
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/lora/ops/triton_ops
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
- pytest -v -s v1/determinism/test_matmul_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 -k 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 -k TRITON_ATTN
# Disabled until the pytest plugins installation is locked.
# - label: ":nvidia: (H100) Batch Invariance"
# key: batch-invariance-h100
# automatic_shard: true
# timeout_in_minutes: 60
# device: h100
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/lora/ops/triton_ops
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pip install pytest-timeout pytest-forked
# - pytest -v -s v1/determinism/test_batch_invariance.py
# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
# - pytest -v -s v1/determinism/test_lora_shrink_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 -k 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 -k FLASH_ATTN
# mirror:
# amd:
# label: ":amd: (MI300) Batch Invariance"
# dind: false
# automatic_shard: true
# device: mi300_1
# timeout_in_minutes: 60
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/lora/ops/triton_ops
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pip install pytest-timeout pytest-forked
# - pytest -v -s v1/determinism/test_batch_invariance.py
# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
# - pytest -v -s v1/determinism/test_matmul_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 -k 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 -k TRITON_ATTN
- label: ":nvidia: (B200) Batch Invariance"
key: batch-invariance-b200
automatic_shard: true
timeout_in_minutes: 45
device: b200-k8s
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/lora/ops/triton_ops
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
- pytest -v -s v1/determinism/test_batch_invariance_vlm.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 -k 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 -k FLASH_ATTN
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- pytest -v -s v1/determinism/test_matmul_batch_invariant.py
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
- pytest -v -s v1/determinism/test_online_batch_invariance.py
mirror:
amd:
label: ":amd: (MI355) Batch Invariance"
dind: false
device: mi355_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
- vllm/lora/ops/triton_ops
- vllm/model_executor/kernels/linear
- vllm/model_executor/determinism
- tests/v1/determinism/
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pip install pytest-timeout pytest-forked
- pytest -v -s v1/determinism/test_batch_invariance.py
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
- pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
- pytest -v -s v1/determinism/test_batch_invariance_vlm.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 -k 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 -k TRITON_ATTN
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- pytest -v -s v1/determinism/test_matmul_batch_invariant.py
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
- pytest -v -s v1/determinism/test_online_batch_invariance.py
- pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
# Disabled until the pytest plugins installation is locked.
# - label: ":nvidia: (B200) Batch Invariance"
# key: batch-invariance-b200
# automatic_shard: true
# timeout_in_minutes: 45
# device: b200-k8s
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/lora/ops/triton_ops
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pip install pytest-timeout pytest-forked
# - pytest -v -s v1/determinism/test_batch_invariance.py
# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
# - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 -k 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 -k FLASH_ATTN
# - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
# - pytest -v -s v1/determinism/test_online_batch_invariance.py
# mirror:
# amd:
# label: ":amd: (MI355) Batch Invariance"
# dind: false
# device: mi355_1
# timeout_in_minutes: 60
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/v1/attention
# - vllm/model_executor/layers
# - vllm/lora/ops/triton_ops
# - vllm/model_executor/kernels/linear
# - vllm/model_executor/determinism
# - tests/v1/determinism/
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pip install pytest-timeout pytest-forked
# - pytest -v -s v1/determinism/test_batch_invariance.py
# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py
# - pytest -v -s v1/determinism/test_batch_invariance_vlm.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 -k 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 -k TRITON_ATTN
# - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py
# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
# - pytest -v -s v1/determinism/test_online_batch_invariance.py
# - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py
- label: ":nvidia: (H100) [Nightly] Batch Invariance (Distributed)"
key: batch-invariance-distributed-4xh100
.buildkite/test_areas/model_executor.yaml
+69 -40
View File
@@ -2,55 +2,84 @@ group: Model Executor
depends_on:
- image-build
steps:
- label: ":nvidia: (H200 MIG 35GB) Model Executor"
device: h200_35gb
key: model-executor
# Disabled until the apt packages installation is locked.
# - label: ":nvidia: (H200 MIG 35GB) Model Executor"
# device: h200_35gb
# key: model-executor
# automatic_shard: true
# timeout_in_minutes: 60
# source_file_dependencies:
# - vllm/engine/arg_utils.py
# - vllm/config/model.py
# - vllm/model_executor
# - vllm/model_executor/warmup
# - tests/model_executor
# - tests/model_executor/test_jit_warmup.py
# - tests/model_executor/test_jit_warmup_cutedsl_launcher.py
# - tests/model_executor/test_jit_warmup_triton_launcher.py
# - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
# commands:
# - apt-get update && apt-get install -y curl libsodium23
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
# # init surfaces a stack instead of silently stalling.
# - export PYTHONFAULTHANDLER=1
# # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
# # init) fails fast with a traceback instead of running until the global
# # build timeout. The `thread` method also handles hangs inside C/CUDA
# # calls that the signal method cannot interrupt.
# - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread
# - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
# mirror:
# amd:
# label: ":amd: (MI300) Model Executor"
# dind: false
# device: mi300_1
# timeout_in_minutes: 75
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/engine/arg_utils.py
# - vllm/config/model.py
# - vllm/model_executor
# - vllm/model_executor/warmup
# - tests/model_executor
# - tests/model_executor/test_jit_warmup.py
# - tests/model_executor/test_jit_warmup_cutedsl_launcher.py
# - tests/model_executor/test_jit_warmup_triton_launcher.py
# - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export PYTHONFAULTHANDLER=1
# - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread
# - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
# Keep the AMD job: its commands do not run this installation.
- label: ':amd: (MI300) Model Executor'
device: mi300_1
key: amd-model-executor
automatic_shard: true
timeout_in_minutes: 60
timeout_in_minutes: 75
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- vllm/models/deepseek_v32/
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/model_executor/test_jit_warmup_cutedsl_launcher.py
- tests/model_executor/test_jit_warmup_triton_launcher.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
# init surfaces a stack instead of silently stalling.
- export PYTHONFAULTHANDLER=1
# Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
# init) fails fast with a traceback instead of running until the global
# build timeout. The `thread` method also handles hangs inside C/CUDA
# calls that the signal method cannot interrupt.
- pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
mirror:
amd:
label: ":amd: (MI300) Model Executor"
dind: false
device: mi300_1
timeout_in_minutes: 75
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/model_executor/test_jit_warmup_cutedsl_launcher.py
- tests/model_executor/test_jit_warmup_triton_launcher.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTHONFAULTHANDLER=1
- pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTHONFAULTHANDLER=1
- pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
dind: false
depends_on:
- image-build-amd
.buildkite/test_areas/models_basic.yaml
+68 -37
View File
@@ -78,8 +78,10 @@ steps:
- tests/models/test_deepseek_v41_replay_start.py
- tests/models/test_deepseek_v4_vl_stream_weights.py
- tests/models/test_muse_spark_rocm_loading.py
- tests/models/test_muse_spark_logits_local.py
- tests/models/test_muse_spark_rocm_shared_rope.py
commands:
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py
mirror:
amd:
label: ":amd: (MI355) Basic Models (Other)"
@@ -103,17 +105,22 @@ steps:
- tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
- tests/models/test_dsv41_mono_decode.py
- tests/models/test_muse_spark_rocm_loading.py
- tests/models/test_dsv4_mono_decode.py
- tests/models/test_deepseek_v4_vl_stream_weights.py
- tests/models/hy_v4/test_triton_ihc.py
- vllm/models/muse_spark/
- tests/models/test_muse_spark_logits_local.py
- tests/models/test_muse_spark_rocm_shared_rope.py
commands:
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py
- VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
- pytest -v -s models/hy_v4/test_triton_ihc.py
- pytest -v -s models/test_dsv4_mono_decode.py
- label: ":amd: (MI355) MiniMax M3 MonoKernel"
key: minimax-m3-monokernel-mi355
dind: false
timeout_in_minutes: 20
timeout_in_minutes: 45
device: mi355_4
num_devices: 4
working_dir: "/vllm-workspace/tests"
@@ -122,8 +129,11 @@ steps:
source_file_dependencies:
- vllm/models/minimax_m3/amd/
- tests/models/minimax_m3/amd_specialized_decode/
- tests/models/minimax_m3/test_index_cache_overrides.py
- vllm/model_executor/models/config.py
commands:
- pytest -v -s models/minimax_m3/amd_specialized_decode
- pytest -v -s models/minimax_m3/test_index_cache_overrides.py
- label: ":nvidia: (B200) Inkling"
key: inkling-unit-tests-b200
@@ -155,16 +165,17 @@ steps:
commands:
- pytest -v -s models/inkling/rocm
- label: ":nvidia: (B200) Flint"
key: flint-unit-tests-b200
timeout_in_minutes: 30
device: b200-k8s
source_file_dependencies:
- vllm/models/flint/
- tests/models/flint/
commands:
- uv pip install --system 'gguf>=0.19.0'
- pytest -v -s models/flint
# Disabled until the gguf installation is locked.
# - label: ":nvidia: (B200) Flint"
# key: flint-unit-tests-b200
# timeout_in_minutes: 30
# device: b200-k8s
# source_file_dependencies:
# - vllm/models/flint/
# - tests/models/flint/
# commands:
# - uv pip install --system 'gguf>=0.19.0'
# - pytest -v -s models/flint
- label: ":nvidia: (B200) Kimi K3"
key: kimi-k3-unit-tests-b200
@@ -196,29 +207,30 @@ steps:
# cuda.bindings.driver), so it is deliberately left out here.
- pytest -v -s models/kimi_k3
- label: ":nvidia: (4xB200) [Nightly] Kimi K3 Prefix Cache"
key: kimi-k3-prefix-cache-4xb200
# Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs.
optional: true
automatic_shard: true
timeout_in_minutes: 75
device: b200-k8s
num_devices: 4
source_file_dependencies:
- vllm/models/kimi_k3/
- vllm/v1/core/
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/
- vllm/distributed/mooncake_store.py
- vllm/v1/simple_kv_offload/
- vllm/v1/spec_decode/
- vllm/model_executor/layers/mamba/
- tests/models/kimi_k3/test_prefix_cache.py
- .buildkite/scripts/install-kv-connectors.sh
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- pytest -v -s models/kimi_k3/test_prefix_cache.py
# Disabled until the KV connectors installation is locked.
# - label: ":nvidia: (4xB200) [Nightly] Kimi K3 Prefix Cache"
# key: kimi-k3-prefix-cache-4xb200
# # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs.
# optional: true
# automatic_shard: true
# timeout_in_minutes: 75
# device: b200-k8s
# num_devices: 4
# source_file_dependencies:
# - vllm/models/kimi_k3/
# - vllm/v1/core/
# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
# - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/
# - vllm/distributed/mooncake_store.py
# - vllm/v1/simple_kv_offload/
# - vllm/v1/spec_decode/
# - vllm/model_executor/layers/mamba/
# - tests/models/kimi_k3/test_prefix_cache.py
# - .buildkite/scripts/install-kv-connectors.sh
# commands:
# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
# - pytest -v -s models/kimi_k3/test_prefix_cache.py
- label: ":nvidia: (H200) Step5"
key: step5-unit-tests
@@ -303,6 +315,7 @@ steps:
- tests/models/test_vision.py
- tests/models/test_adapters.py
- tests/models/test_qwen3_5_mtp_config.py
- tests/models/test_dory.py
- tests/models/test_alice_ai.py
- tests/models/test_alice_ai_config.py
- tests/models/test_beam_weight_mapping.py
@@ -318,7 +331,7 @@ steps:
- tests/models/test_cohere2_moe_parallel_block.py
device: cpu-small
commands:
- pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py models/test_cohere2_moe_parallel_block.py
- pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_dory.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py models/test_cohere2_moe_parallel_block.py
- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics'
key: dsv41-mono-decode-numerics-amd
@@ -339,6 +352,24 @@ steps:
- pytest -v -s tests/models/test_dsv41_mono_numerics.py
- pytest -v -s tests/models/test_dsv41_mono_layers.py
- label: ':amd: (MI355) DeepSeek-V4-Pro Mono Decode Numerics'
key: dsv4-mono-decode-numerics-amd
dind: false
device: mi355_8
num_devices: 8
timeout_in_minutes: 60
working_dir: /vllm-workspace/
depends_on:
- image-build-amd
source_file_dependencies:
- docker/Dockerfile.rocm_base
- vllm/models/deepseek_v4/amd/
- tests/models/test_dsv4_mono_numerics.py
- tests/models/test_dsv4_mono_layers.py
commands:
- pytest -v -s tests/models/test_dsv4_mono_numerics.py
- pytest -v -s tests/models/test_dsv4_mono_layers.py
- label: ':amd: (MI300) DeepSeek-V4.1 gfx942 Mono Decode Layers'
key: dsv41-mono-decode-layers-gfx942-amd
dind: false
.buildkite/test_areas/plugins.yaml
+74 -58
View File
@@ -11,6 +11,13 @@ steps:
source_file_dependencies:
- vllm/plugins/
- tests/plugins/
# Import-time platform resolution (see test_import_vllm_does_not_resolve_platform)
- vllm/__init__.py
- vllm/env_override.py
- vllm/platforms/__init__.py
- vllm/utils/torch_utils.py
- vllm/utils/torch_version.py
- tests/plugins_tests/test_platform_plugins.py
commands:
# begin platform plugin and general plugin tests, all the code in-between runs on dummy platform
- pip install -e ./plugins/vllm_add_dummy_platform
@@ -62,65 +69,74 @@ steps:
- vllm/plugins/
- tests/plugins/
- vllm/platforms/rocm.py
- vllm/__init__.py
- vllm/env_override.py
- vllm/platforms/__init__.py
- vllm/utils/torch_utils.py
- vllm/utils/torch_version.py
- tests/plugins_tests/test_platform_plugins.py
- label: ":nvidia: (H200 MIG 18GB) [Nightly] GGUF Plugin"
key: gguf-plugin
device: h200_18gb
timeout_in_minutes: 30
soft_fail: true
optional: true
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/gguf/
commands:
- pip install "vllm-gguf-plugin >= 0.0.2"
- pytest -v -s plugins_tests/gguf
mirror:
amd:
label: ":amd: (MI355 DPX) GGUF Plugin"
dind: false
device: mi355_dpx
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/gguf
- vllm/platforms/rocm.py
# Disabled until the GGUF plugin installation is locked.
# - label: ":nvidia: (H200 MIG 18GB) [Nightly] GGUF Plugin"
# key: gguf-plugin
# device: h200_18gb
# timeout_in_minutes: 30
# soft_fail: true
# optional: true
# source_file_dependencies:
# - vllm/model_executor/layers/quantization
# - tests/plugins_tests/gguf/
# commands:
# - pip install "vllm-gguf-plugin >= 0.0.2"
# - pytest -v -s plugins_tests/gguf
# mirror:
# amd:
# label: ":amd: (MI355 DPX) GGUF Plugin"
# dind: false
# device: mi355_dpx
# timeout_in_minutes: 45
# working_dir: "/vllm-workspace/tests"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - vllm/model_executor/layers/quantization
# - tests/plugins_tests/gguf
# - vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) [Nightly] BitsAndBytes Plugin"
key: bitsandbytes-plugin
device: h200_18gb
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
soft_fail: true
optional: true
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/bitsandbytes
commands:
# bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so
# there is nothing to wrap on the vLLM side.
- unset VLLM_GPU_SYNC_CHECK
- pip install "vllm-bnb-plugin >= 0.0.1"
- pytest -v -s plugins_tests/bitsandbytes -m 'not distributed'
# Disabled until the BitsAndBytes plugin installation is locked.
# - label: ":nvidia: (H200 MIG 18GB) [Nightly] BitsAndBytes Plugin"
# key: bitsandbytes-plugin
# device: h200_18gb
# timeout_in_minutes: 30
# working_dir: "/vllm-workspace/tests"
# soft_fail: true
# optional: true
# source_file_dependencies:
# - vllm/model_executor/layers/quantization
# - tests/plugins_tests/bitsandbytes
# commands:
# # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so
# # there is nothing to wrap on the vLLM side.
# - unset VLLM_GPU_SYNC_CHECK
# - pip install "vllm-bnb-plugin >= 0.0.1"
# - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed'
- label: ":nvidia: (L4) [Nightly] BitsAndBytes Plugin"
key: bitsandbytes-plugin-2-gpus
timeout_in_minutes: 15
working_dir: "/vllm-workspace/tests"
device: l4
num_devices: 2
soft_fail: true
optional: true
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/bitsandbytes
commands:
# bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so
# there is nothing to wrap on the vLLM side.
- unset VLLM_GPU_SYNC_CHECK
- pip install "vllm-bnb-plugin >= 0.0.1"
- pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)'
# Disabled until the BitsAndBytes plugin installation is locked.
# - label: ":nvidia: (L4) [Nightly] BitsAndBytes Plugin"
# key: bitsandbytes-plugin-2-gpus
# timeout_in_minutes: 15
# working_dir: "/vllm-workspace/tests"
# device: l4
# num_devices: 2
# soft_fail: true
# optional: true
# source_file_dependencies:
# - vllm/model_executor/layers/quantization
# - tests/plugins_tests/bitsandbytes
# commands:
# # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so
# # there is nothing to wrap on the vLLM side.
# - unset VLLM_GPU_SYNC_CHECK
# - pip install "vllm-bnb-plugin >= 0.0.1"
# - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)'
.buildkite/test_areas/quantization.yaml
+50 -44
View File
@@ -2,50 +2,56 @@ group: Quantization
depends_on:
- image-build
steps:
- label: ":nvidia: (H200 MIG 35GB) Quantization"
device: h200_35gb
key: quantization
automatic_shard: true
timeout_in_minutes: 40
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- vllm/models/deepseek_v41/quant_config.py
- vllm/utils/humming.py
- tests/quantization
commands:
# Pin torchao to the version compatible with the CI PyTorch and CUDA stack.
- uv pip install --system torchao==0.18.0 --index-url https://download.pytorch.org/whl/cu130
- uv pip install --system conch-triton-kernels
# The SM90-only checkpoint currently contains a removed weight_chan_scale
# parameter. It was not exercised by the previous L4 job.
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
mirror:
amd:
label: ":amd: (MI300) Quantization"
automatic_shard: true
dind: false
device: mi300_1
timeout_in_minutes: 115
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
- tests/quantization
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- vllm/model_executor/layers/fused_moe/oracle/unquantized.py
- vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
- tests/rocm/test_moe_weight_replay.py
- tests/rocm/test_moe_padding.py
commands:
- uv pip install --system torchao==0.17.0
- uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py
- pytest -v -s rocm/test_moe_weight_replay.py
- VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py
# Disabled until torchao transitive dependencies are locked.
# - label: ":nvidia: (H200 MIG 35GB) Quantization"
# device: h200_35gb
# key: quantization
# automatic_shard: true
# timeout_in_minutes: 40
# source_file_dependencies:
# - csrc/
# - vllm/model_executor/layers/quantization
# - vllm/model_executor/models/qwen3_dflash.py
# - vllm/model_executor/models/utils.py
# - vllm/transformers_utils/model_arch_config_convertor.py
# - vllm/models/deepseek_v41/quant_config.py
# - vllm/utils/humming.py
# - tests/quantization
# commands:
# # Pin torchao to the version compatible with the CI PyTorch and CUDA stack.
# - uv pip install --system torchao==0.18.0 --index-url https://download.pytorch.org/whl/cu130
# - uv pip install --system conch-triton-kernels
# # The SM90-only checkpoint currently contains a removed weight_chan_scale
# # parameter. It was not exercised by the previous L4 job.
# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py -k 'not test_compressed_tensors_w4a8_fp8'
# mirror:
# amd:
# label: ":amd: (MI300) Quantization"
# automatic_shard: true
# dind: false
# device: mi300_1
# timeout_in_minutes: 115
# working_dir: "/vllm-workspace/tests"
# depends_on:
# - image-build-amd
# source_file_dependencies:
# - csrc/
# - vllm/model_executor/layers/quantization
# - vllm/model_executor/models/qwen3_dflash.py
# - vllm/model_executor/models/utils.py
# - tests/quantization
# - vllm/_aiter_ops.py
# - vllm/platforms/rocm.py
# - vllm/model_executor/layers/fused_moe/oracle/unquantized.py
# - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
# - tests/rocm/test_moe_weight_replay.py
# - tests/rocm/test_moe_padding.py
# commands:
# - uv pip install --system torchao==0.17.0
# - uv pip install --system conch-triton-kernels
# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py
# - pytest -v -s rocm/test_moe_weight_replay.py
# - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py
- label: ":nvidia: (H200 MIG 35GB) Quantized Fusions"
device: h200_35gb
.buildkite/test_areas/rust_frontend.yaml
+77 -26
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: ":nvidia: (H200 MIG 18GB) Rust Frontend OpenAI Coverage"
key: rust-frontend-openai-coverage
timeout_in_minutes: 30
timeout_in_minutes: 50
device: h200_18gb
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -18,35 +18,53 @@ steps:
- tests/entrypoints/openai/chat_completion/test_chat_completion.py
- tests/entrypoints/openai/chat_completion/test_include_reasoning.py
- tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
- tests/entrypoints/openai/chat_completion/test_serving_chat.py
- tests/entrypoints/openai/chat_completion/test_chat_echo.py
- tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py
- tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
- tests/entrypoints/openai/completion/test_completion.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/test_chunked_prompt.py
# - tests/entrypoints/openai/completion/test_prompt_validation.py
- tests/entrypoints/launchers/test_shutdown.py
- tests/entrypoints/openai/test_return_token_ids.py
- tests/entrypoints/serve/instrumentator/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
- tests/v1/logits_processors/test_custom_online.py
- tests/v1/logits_processors/utils.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
- pytest -v -s benchmarks/test_serve_cli.py
# The Rust frontend re-renders messages instead of reusing kv_transfer_params["prompt_token_ids"].
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
- pytest -v -s entrypoints/openai/chat_completion/test_include_reasoning.py
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
- pytest -v -s entrypoints/openai/chat_completion/test_serving_chat.py -k "test_chat_per_request_metrics_follow_server_flag or test_streaming_reasoning_usage_counts_across_deltas or test_completion_tokens_details"
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
- pytest -v -s entrypoints/openai/chat_completion/test_chat_echo.py
- pytest -v -s entrypoints/openai/chat_completion/test_logprob_token_ids.py
# The Qwen3.5-35B FP8 MTP case needs two GPUs.
- pytest -v -s entrypoints/openai/chat_completion/test_thinking_token_budget.py -k "not qwen35"
# The Rust frontend drops prompt_tokens_details when nothing was cached, rejects batched
# prompts, and supports only n=1 (so no parallel sampling or beam search).
- pytest -v -s entrypoints/openai/completion/test_completion.py -k "not test_single_completion and not test_prompt_logprobs_completion and not test_parallel and not test_batch_completions"
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
- pytest -v -s entrypoints/openai/test_chunked_prompt.py
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
- pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
# test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and
# finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk.
- pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
# test_shutdown_on_engine_failure launches the Python API server directly, not the Rust frontend.
- pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure"
- pytest -v -s entrypoints/openai/test_return_token_ids.py
- pytest -v -s entrypoints/serve/instrumentator/test_uds.py
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
# The Rust frontend cannot run the logits processor's validate_params, so an invalid argument
# reaches the engine instead of returning 400.
- pytest -v -s v1/logits_processors/test_custom_online.py -k "not test_invalid_custom_logitsproc_arg"
mirror:
amd:
label: ":amd: (MI355 DPX) Rust Frontend OpenAI Coverage"
dind: false
device: mi355_dpx
timeout_in_minutes: 35
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
@@ -60,16 +78,23 @@ steps:
- tests/entrypoints/openai/chat_completion/test_chat_completion.py
- tests/entrypoints/openai/chat_completion/test_include_reasoning.py
- tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
- tests/entrypoints/openai/chat_completion/test_serving_chat.py
- tests/entrypoints/openai/chat_completion/test_chat_echo.py
- tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py
- tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py
- tests/entrypoints/openai/completion/test_completion.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/test_chunked_prompt.py
- tests/entrypoints/launchers/test_shutdown.py
- tests/entrypoints/openai/test_return_token_ids.py
- tests/entrypoints/serve/instrumentator/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
- tests/v1/logits_processors/test_custom_online.py
- tests/v1/logits_processors/utils.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) Rust Frontend Serve/Admin Coverage"
key: rust-frontend-serve-admin-coverage
timeout_in_minutes: 25
timeout_in_minutes: 45
device: h200_18gb
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -83,23 +108,32 @@ steps:
- tests/entrypoints/serve/instrumentator/test_basic.py
- tests/entrypoints/serve/instrumentator/test_metrics.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
- tests/entrypoints/serve/middleware/test_optional_middleware.py
- tests/entrypoints/launchers/api_server/test_multi_api_servers.py
- tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
# server_load can be flaky under the Rust frontend; keep it excluded for now.
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not server_load"
# test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)).
- pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
# /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op).
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py
# The Rust LoRA request still carries the removed tensorizer_config_dict field, so --lora-modules fails at startup.
- pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not lora"
# run_batch drives the Python batch runner, not a server. test_metrics_exist needs
# vllm:admission_rejections (text) and image inputs (multimodal), which the Rust frontend lacks.
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "not run_batch and not test_metrics_exist"
# /tokenizer_info returns 501 in the Rust frontend.
- pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"
# test_v2_endpoint_accepts_valid_api_token calls /v2/embed; the Rust frontend serves no pooling routes.
- pytest -v -s entrypoints/serve/middleware/test_optional_middleware.py -k "not test_v2_endpoint_accepts_valid_api_token"
# The Rust frontend ignores --api-server-count and runs one server, so this only checks the flag is accepted.
- pytest -v -s entrypoints/launchers/api_server/test_multi_api_servers.py
- pytest -v -s entrypoints/rl/return_routed_experts/test_scale_out_api.py
mirror:
amd:
label: ":amd: (MI355 DPX) Rust Frontend Serve/Admin Coverage"
dind: false
device: mi355_dpx
timeout_in_minutes: 30
timeout_in_minutes: 50
depends_on:
- image-build-amd
source_file_dependencies:
@@ -113,6 +147,9 @@ steps:
- tests/entrypoints/serve/instrumentator/test_basic.py
- tests/entrypoints/serve/instrumentator/test_metrics.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
- tests/entrypoints/serve/middleware/test_optional_middleware.py
- tests/entrypoints/launchers/api_server/test_multi_api_servers.py
- tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) Rust Frontend Core Correctness"
@@ -147,7 +184,7 @@ steps:
- label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use"
key: rust-frontend-tool-use
device: h200_35gb
timeout_in_minutes: 25
timeout_in_minutes: 35
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
@@ -155,16 +192,25 @@ steps:
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
- tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py
- tests/entrypoints/tool_parsers/test_granite4_tool_parser.py
- tests/entrypoints/tool_parsers/test_hermes_tool_parser.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
# The Rust frontend has no `mistral` tokenizer mode or renderer.
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2
# The Rust renderer rejects empty tool-call arguments in the history (Python treats them as {}).
- pytest -v -s entrypoints/openai/chat_completion/test_completion_with_function_calling.py -k "not test_no_args_tool_call"
- pytest -v -s entrypoints/tool_parsers/test_granite4_tool_parser.py
# The llama cases use --lora-modules, which fails at startup until the Rust LoRA request drops tensorizer_config_dict.
- pytest -v -s entrypoints/tool_parsers/test_hermes_tool_parser.py -k granite4
mirror:
amd:
label: ":amd: (MI355 DPX) Rust Frontend Tool Use"
dind: false
device: mi355_dpx
timeout_in_minutes: 30
timeout_in_minutes: 40
depends_on:
- image-build-amd
source_file_dependencies:
@@ -173,6 +219,9 @@ steps:
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
- tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py
- tests/entrypoints/tool_parsers/test_granite4_tool_parser.py
- tests/entrypoints/tool_parsers/test_hermes_tool_parser.py
- vllm/platforms/rocm.py
- label: ":nvidia: (L4) [Nightly] Rust Frontend Distributed"
@@ -201,9 +250,11 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
# /server_info reports the Rust frontend's own config, which has no parallel_config. The Rust frontend
# ignores --api-server-count, so the api-server-count=4 cases run against one multithreaded server.
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not server_info"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not server_info"
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not server_info"
mirror:
amd:
# Run on PRs even when the L4 step is optional (nightly only).
.buildkite/test_areas/spec_decode.yaml
+2 -2
View File
@@ -418,9 +418,9 @@ steps:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
mirror:
amd:
label: ":amd: (MI300) Spec Decode DeepSeek MTP Parallel Load"
label: ":amd: (MI355) Spec Decode DeepSeek MTP Parallel Load"
dind: false
device: mi300_2
device: mi355_2
timeout_in_minutes: 45
depends_on:
- image-build-amd
.buildkite/test_areas/torch_abi.yaml
+20 -19
View File
@@ -12,22 +12,23 @@ steps:
- setup.py
commands:
- python3 /vllm-workspace/.buildkite/check-torch-abi.py
mirror:
amd:
label: ":amd: (MI300) Torch Stable ABI Audit"
dind: false
device: mi300_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
no_gpu: true
working_dir: /vllm-workspace/
source_file_dependencies:
- .buildkite/check-torch-abi.py
- csrc/
- cmake/
- setup.py
- vllm/platforms/rocm.py
commands:
- uv pip install --system torch-abi-audit==0.0.1
- python3 /vllm-workspace/.buildkite/check-torch-abi.py
# Disabled until torch-abi-audit transitive dependencies are locked.
# mirror:
# amd:
# label: ":amd: (MI300) Torch Stable ABI Audit"
# dind: false
# device: mi300_1
# timeout_in_minutes: 35
# depends_on:
# - image-build-amd
# no_gpu: true
# working_dir: /vllm-workspace/
# source_file_dependencies:
# - .buildkite/check-torch-abi.py
# - csrc/
# - cmake/
# - setup.py
# - vllm/platforms/rocm.py
# commands:
# - uv pip install --system torch-abi-audit==0.0.1
# - python3 /vllm-workspace/.buildkite/check-torch-abi.py
.github/mergify.yml
+35
View File
@@ -14,6 +14,19 @@ pull_request_rules:
comment:
message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"
- name: label-recipe
description: Automatically apply recipe label
conditions:
- label != stale
- -closed
- files~=^recipes/
actions:
label:
add:
- recipe
comment:
message: "Recipe preview: https://vllm-recipes-git-preview-vllm-pr-{{number}}-inferact-inc.vercel.app"
- name: comment-pre-commit-failure
description: Comment on PR when pre-commit check fails
conditions:
@@ -334,6 +347,28 @@ pull_request_rules:
add:
- minimax
- name: label-beam
description: Automatically apply beam label for Reflection AI's model
conditions:
- label != stale
- or:
- files~=^vllm/models/beam/
- files=vllm/transformers_utils/configs/beam.py
- files=vllm/parser/beam.py
- files=vllm/reasoning/beam_reasoning_parser.py
- files=vllm/tool_parsers/beam_tool_parser.py
- files~=^rust/src/parser/src/unified/beam(?:\.rs|/)
- files~=^tests/models/test_beam_.*\.py$
- files=tests/parser/test_beam_parser.py
- files=tests/transformers_utils/test_beam_config.py
- files~=^recipes/models/reflection/Beam-
# Bare "beam" also describes beam search; require a model identifier.
- title~=(?i)(?:\bbeam[- ]\d|\breflection(?:\s+ai)?[ /-]+beam\b|\bBeamForCausalLM\b|\bBeamConfig\b|\[beam\])
actions:
label:
add:
- beam
- name: label-inkling
description: Automatically apply inkling label
conditions:
.github/workflows/issue_autolabel.yml
+8
View File
@@ -197,6 +197,14 @@ jobs:
deepseek: {
keywords: [{ term: "DeepSeek", searchIn: "title" }],
},
beam: {
// Bare "beam" also describes beam search; require a model identifier.
regexPatterns: [{
pattern: "(?:\\bbeam[- ]\\d|\\breflection(?:\\s+ai)?[ /-]+beam\\b|\\bBeamForCausalLM\\b|\\bBeamConfig\\b|\\[beam\\])",
flags: "gi",
searchIn: "title",
}],
},
// DSv4 and DSv4.1 are distinct architectures; these keep them exclusive.
DSv4: {
regexPatterns: [{ pattern: "(?:\\bDSv4\\b|deepseek[-\\s_]?v4)(?![._]?1)", flags: "gi", searchIn: "title" }],
.github/workflows/recipes-preview.yml
+69
View File
@@ -0,0 +1,69 @@
name: Recipes Preview
# Tells vllm-project/recipes when recipes/ changes, instead of waiting for its
# scheduled polling, which GitHub runs hours late:
# - a recipe PR is opened, pushed or reopened: rebuild its preview branch
# `preview/vllm-pr-<N>`, which Vercel deploys at the environment URL below
# (recipes debounces bursts of pushes);
# - a recipe PR is closed: delete that branch;
# - recipe changes reach main: sync recipes.vllm.ai (`recipes-updated`).
#
# Runs on pull_request_target so fork PRs can reach the dispatch token. This
# is safe only because no step checks out or runs PR code, and only the PR
# number leaves this workflow. Keep it that way. The token comes from a GitHub
# App installed only on vllm-project/recipes with contents: write; its key
# lives in the `recipes-preview` environment, which only main may use.
on:
pull_request_target:
types: [opened, synchronize, reopened, closed]
paths:
- 'recipes/**'
push:
branches: [main]
paths:
- 'recipes/**'
permissions:
contents: read
concurrency:
group: recipes-preview-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
dispatch:
if: github.repository == 'vllm-project/vllm'
runs-on: ubuntu-latest
timeout-minutes: 5
environment:
name: recipes-preview
url: >-
${{ github.event_name == 'push' && 'https://recipes.vllm.ai' ||
(github.event.action != 'closed' &&
format('https://vllm-recipes-git-preview-vllm-pr-{0}-inferact-inc.vercel.app', github.event.pull_request.number) ||
'') }}
steps:
- name: Create a token for vllm-project/recipes
id: token
uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
with:
app-id: ${{ vars.RECIPES_APP_ID }}
private-key: ${{ secrets.RECIPES_APP_PRIVATE_KEY }}
owner: vllm-project
repositories: recipes
permission-contents: write
- name: Dispatch to vllm-project/recipes
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
EVENT: ${{ github.event_name }}
PR: ${{ github.event.pull_request.number }}
run: |
if [ "$EVENT" = push ]; then
gh api repos/vllm-project/recipes/dispatches -f event_type=recipes-updated
exit 0
fi
[[ "$PR" =~ ^[0-9]+$ ]] || { echo "Invalid PR number: $PR"; exit 1; }
gh api repos/vllm-project/recipes/dispatches \
-f event_type=preview-pr -F "client_payload[pr]=${PR}"
CMakeLists.txt
+35 -1
View File
@@ -31,6 +31,9 @@ endif()
# CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py)
set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM")
if (VLLM_TARGET_DEVICE STREQUAL "rocm" OR "$ENV{VLLM_KERNEL_SYMBOL_MAP}" STREQUAL "1")
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
endif()
message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
message(STATUS "Target device: ${VLLM_TARGET_DEVICE}")
@@ -365,6 +368,20 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/rocm/quickreduce/torch_bindings.cpp"
"csrc/rocm/quickreduce/custom_quickreduce.cu")
# Octave KV cache (store, decode, rotation, prefill): only built for the
# validated architectures, gfx1100 (RDNA3) and gfx942 (MI300).
set(VLLM_OCTAVE_KV OFF)
if(VLLM_GPU_ARCHES MATCHES "gfx942|gfx1100")
set(VLLM_OCTAVE_KV ON)
list(APPEND VLLM_EXT_SRC
"csrc/attention/octave_torch_bindings.cpp"
"csrc/attention/octave_attn.cu"
"csrc/attention/octave_decode_wmma_rdna3.cu"
"csrc/attention/octave_prefill_rdna3.cu"
"csrc/attention/octave_prefill_cdna.cu"
"csrc/attention/octave_decode_mfma_cdna.cu"
)
endif()
message(STATUS "Enabling C extension.")
define_extension_target(
@@ -376,6 +393,13 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
ARCHITECTURES ${VLLM_GPU_ARCHES}
USE_SABI 3
WITH_SOABI)
if(VLLM_OCTAVE_KV)
# MFMA accumulators in VGPRs: in AGPRs every online-softmax rescale of the
# output copies it through 64 extra VGPRs and halves the occupancy.
set_source_files_properties(
"${CMAKE_CURRENT_BINARY_DIR}/csrc/attention/octave_decode_mfma_cdna.hip"
PROPERTIES COMPILE_FLAGS "-mllvm -amdgpu-mfma-vgpr-form")
endif()
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
@@ -406,6 +430,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
WORKING_DIRECTORY "${TORCH_SHADOW_INCLUDE_DIR}"
COMMAND_ERROR_IS_FATAL ANY)
execute_process(
COMMAND ${Python_EXECUTABLE} -c
"import json,sys; from pathlib import Path; Path(sys.argv[1]).write_text(json.dumps({sys.argv[2]: sys.argv[3:]}))"
"${CMAKE_CURRENT_BINARY_DIR}/generated-source-deps.json"
"${TORCH_SHADOW_INCLUDE_DIR}/${_stable_hdr_rel}"
"${TORCH_INCLUDE_DIR}/${_stable_hdr_rel}"
"${CMAKE_CURRENT_LIST_DIR}/cmake/patches/pytorch_stable_string.patch"
COMMAND_ERROR_IS_FATAL ANY)
include_directories(BEFORE "${TORCH_SHADOW_INCLUDE_DIR}")
endif()
endif()
@@ -1618,7 +1651,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
list(APPEND VLLM_ROCM_EXT_SRC
"csrc/rocm/q_gemm_rdna3.cu"
"csrc/rocm/q_gemm_rdna3_wmma.cu"
"csrc/rocm/moe_q_gemm_rdna3.cu")
"csrc/rocm/moe_q_gemm_rdna3.cu"
"csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu")
endif()
define_extension_target(
benchmarks/attention_benchmarks/benchmark.py
+18 -2
View File
@@ -32,6 +32,7 @@ import subprocess
import sys
from dataclasses import replace
from pathlib import Path
from typing import get_args
import yaml
from rich.console import Console
@@ -50,6 +51,7 @@ from common import (
is_mla_backend,
)
from vllm.config.cache import CacheDType
from vllm.v1.worker.workspace import init_workspace_manager
@@ -548,6 +550,11 @@ def main():
# Benchmark settings
parser.add_argument("--device", default="cuda:0", help="Device")
parser.add_argument(
"--model",
default="meta-llama/Meta-Llama-3-8B",
help="Model whose config the backend is built with",
)
parser.add_argument(
"--warmup-ms",
type=int,
@@ -561,8 +568,8 @@ def main():
parser.add_argument(
"--kv-cache-dtype",
default="auto",
choices=["auto", "fp8"],
help="KV cache dtype: auto or fp8",
choices=get_args(CacheDType),
help="KV cache dtype",
)
parser.add_argument(
"--cuda-graphs",
@@ -894,6 +901,7 @@ def main():
device=args.device,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
model=args.model,
use_cuda_graphs=args.cuda_graphs,
prefill_backend="fa4",
)
@@ -964,6 +972,7 @@ def main():
device=args.device,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
model=args.model,
use_cuda_graphs=args.cuda_graphs,
ncu_profile=args.ncu_profile,
warmup_ms=args.warmup_ms,
@@ -1036,6 +1045,8 @@ def main():
pbar.update(1)
if args.ncu_profile:
if any(not result.success for result in all_results):
raise SystemExit(1)
return
# Display decode vs prefill results
@@ -1208,6 +1219,7 @@ def main():
device=args.device,
max_model_len=getattr(args, "max_model_len", None),
kv_cache_dtype=args.kv_cache_dtype,
model=args.model,
profile_memory=args.profile_memory,
use_cuda_graphs=args.cuda_graphs,
ncu_profile=args.ncu_profile,
@@ -1353,6 +1365,7 @@ def main():
device=args.device,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
model=args.model,
use_cuda_graphs=args.cuda_graphs,
ncu_profile=args.ncu_profile,
warmup_ms=args.warmup_ms,
@@ -1435,6 +1448,9 @@ def main():
if args.output_json:
formatter.save_json(all_results, args.output_json)
if any(not result.success for result in all_results):
raise SystemExit(1)
if __name__ == "__main__":
main()
benchmarks/attention_benchmarks/common.py
+3 -1
View File
@@ -284,8 +284,10 @@ class BenchmarkConfig:
torch_profile_iters: int = 3
warmup_ms: int | None = None
# "auto" or "fp8"
# Any CacheDType: "auto", "fp8", "int8_per_token_head", ...
kv_cache_dtype: str = "auto"
# Model whose config the backend sees (dtype, rotary dims, ...)
model: str = "meta-llama/Meta-Llama-3-8B"
# MLA-specific
prefill_backend: str | None = None
benchmarks/attention_benchmarks/runner.py
+26 -21
View File
@@ -36,6 +36,7 @@ from vllm.config import (
set_current_vllm_config,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
from vllm.v1.attention.backends.utils import (
CommonAttentionMetadata,
get_supported_kv_cache_layouts,
@@ -48,6 +49,7 @@ from vllm.v1.kv_cache_interface import (
compute_layer_kv_cache_shape_bytes,
compute_layout_strides,
create_kv_cache_views,
get_kv_quant_mode,
)
# ============================================================================
@@ -142,15 +144,16 @@ def _build_common_attn_metadata(
def _create_vllm_config(
config: BenchmarkConfig,
max_num_blocks: int,
max_model_len: int = 0,
) -> VllmConfig:
"""Create a VllmConfig for benchmarking with mock model methods."""
model_config = ModelConfig(
model="meta-llama/Meta-Llama-3-8B",
tokenizer="meta-llama/Meta-Llama-3-8B",
model=config.model,
tokenizer=config.model,
trust_remote_code=False,
dtype="auto", # Use model's native dtype
seed=0,
max_model_len=1024,
max_model_len=max(1024, max_model_len),
)
cache_config = CacheConfig(
@@ -164,7 +167,7 @@ def _create_vllm_config(
scheduler_config = SchedulerConfig(
max_num_seqs=256,
max_num_batched_tokens=8192,
max_model_len=8192,
max_model_len=max(8192, max_model_len),
is_encoder_decoder=False,
enable_chunked_prefill=True,
)
@@ -342,6 +345,22 @@ def _create_input_tensors(
return q_list, k_list, v_list
def _kv_cache_spec(config: BenchmarkConfig, dtype: torch.dtype) -> FullAttentionSpec:
if config.kv_cache_dtype == "auto":
cache_dtype = dtype
elif config.kv_cache_dtype.startswith("fp8"):
cache_dtype = current_platform.fp8_dtype()
else:
cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[config.kv_cache_dtype]
return FullAttentionSpec(
block_size=config.block_size,
num_kv_heads=config.num_kv_heads,
head_size=config.head_dim,
dtype=cache_dtype,
kv_quant_mode=get_kv_quant_mode(config.kv_cache_dtype),
)
def _create_kv_cache(
config: BenchmarkConfig,
max_num_blocks: int,
@@ -351,16 +370,7 @@ def _create_kv_cache(
layout: KVCacheLayout,
) -> list:
"""Create KV cache tensors for all layers using the standard allocator."""
if config.kv_cache_dtype.startswith("fp8"):
cache_dtype = current_platform.fp8_dtype()
else:
cache_dtype = dtype
spec = FullAttentionSpec(
block_size=config.block_size,
num_kv_heads=config.num_kv_heads,
head_size=config.head_dim,
dtype=cache_dtype,
)
spec = _kv_cache_spec(config, dtype)
# Apply the backend's page customization, as the worker does for the real spec.
spec = backend_class.customize_spec(spec)
total_bytes = (
@@ -492,7 +502,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
# Suppress vLLM logs during setup to reduce spam
with log_warnings_and_errors_only():
# Create vllm_config first - uses model's native dtype via "auto"
vllm_config = _create_vllm_config(config, max_num_blocks)
vllm_config = _create_vllm_config(config, max_num_blocks, max_kv)
dtype = vllm_config.model_config.dtype
# Wrap everything in set_current_vllm_config context
@@ -510,12 +520,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
q_lens, kv_lens, config.block_size, device
)
kv_cache_spec = FullAttentionSpec(
block_size=config.block_size,
num_kv_heads=config.num_kv_heads,
head_size=config.head_dim,
dtype=dtype,
)
kv_cache_spec = backend_class.customize_spec(_kv_cache_spec(config, dtype))
builder = _create_metadata_builder(
backend_class, kv_cache_spec, vllm_config, device, config.backend
benchmarks/kernels/benchmark_moe.py
+42 -2
View File
@@ -123,6 +123,7 @@ def benchmark_config(
num_iters: int = 100,
block_quant_shape: list[int] = None,
use_deep_gemm: bool = False,
per_act_token_quant: bool = False,
) -> float:
init_dtype = torch.float16 if use_fp8_w8a8 else dtype
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
@@ -224,12 +225,19 @@ def benchmark_config(
torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
* factor_for_scale
)
elif per_act_token_quant:
# Per-channel weights; activations are quantized per token at run time.
w1_scale = torch.rand(
(num_experts, shard_intermediate_size, 1), dtype=torch.float32
)
w2_scale = torch.rand((num_experts, hidden_size, 1), dtype=torch.float32)
else:
w1_scale = torch.randn(num_experts, dtype=torch.float32)
w2_scale = torch.randn(num_experts, dtype=torch.float32)
a1_scale = torch.randn(1, dtype=torch.float32)
a2_scale = torch.randn(1, dtype=torch.float32)
if not per_act_token_quant:
a1_scale = torch.randn(1, dtype=torch.float32)
a2_scale = torch.randn(1, dtype=torch.float32)
w1 = w1.to(FP8_DTYPE)
w2 = w2.to(FP8_DTYPE)
@@ -255,6 +263,8 @@ def benchmark_config(
w2_scale=w2_scale,
a1_scale=a1_scale,
a2_scale=a2_scale,
per_act_token_quant=per_act_token_quant,
per_out_ch_quant=per_act_token_quant,
block_shape=block_quant_shape,
weight_dtype="int4" if use_int4_w4a16 else None,
)
@@ -606,6 +616,7 @@ class BenchmarkWorker:
use_int4_w4a16: bool = False,
block_quant_shape: list[int] = None,
use_deep_gemm: bool = False,
per_act_token_quant: bool = False,
) -> tuple[dict[str, int], float]:
# local import to allow serialization by ray
@@ -649,6 +660,7 @@ class BenchmarkWorker:
num_iters=100,
block_quant_shape=block_quant_shape,
use_deep_gemm=use_deep_gemm,
per_act_token_quant=per_act_token_quant,
)
return config, kernel_time
@@ -666,6 +678,7 @@ class BenchmarkWorker:
search_space: list[dict[str, int]],
block_quant_shape: list[int],
use_deep_gemm: bool,
per_act_token_quant: bool = False,
) -> dict[str, int]:
# local import to allow serialization by ray
from vllm.platforms import current_platform
@@ -713,6 +726,7 @@ class BenchmarkWorker:
num_iters=20,
block_quant_shape=block_quant_shape,
use_deep_gemm=use_deep_gemm,
per_act_token_quant=per_act_token_quant,
)
except triton.runtime.autotuner.OutOfResources:
# Some configurations may be invalid and fail to compile.
@@ -819,6 +833,24 @@ def get_weight_block_size_safety(config, default_value=None):
return default_value
def uses_per_token_act_quant(config) -> bool:
"""Whether a compressed-tensors checkpoint quantizes activations per token.
vLLM then runs the fused MoE with per-token activation and per-channel
weight scales, which is a different kernel specialization from per-tensor
FP8 but reads the same config file.
"""
quantization_config = getattr(config, "quantization_config", {})
if not isinstance(quantization_config, dict):
return False
config_groups = quantization_config.get("config_groups", {})
if len(config_groups) != 1:
return False
group = next(iter(config_groups.values()))
input_activations = group.get("input_activations") or {}
return input_activations.get("strategy") == "token"
def get_model_params(config):
architectures = getattr(config, "architectures", None) or [type(config).__name__]
architecture = architectures[0]
@@ -850,6 +882,7 @@ def get_model_params(config):
hidden_size = config.hidden_size
elif architecture in (
"BailingMoeV3ForCausalLM",
"BeamForCausalLM",
"Qwen2MoeForCausalLM",
"Qwen3MoeForCausalLM",
"Qwen3NextForCausalLM",
@@ -990,6 +1023,11 @@ def main(args: argparse.Namespace):
use_int8_w8a16 = args.dtype == "int8_w8a16"
use_int4_w4a16 = args.dtype == "int4_w4a16"
block_quant_shape = get_weight_block_size_safety(config)
per_act_token_quant = (
use_fp8_w8a8 and block_quant_shape is None and uses_per_token_act_quant(config)
)
if per_act_token_quant:
print("Benchmarking per-token activation and per-channel weight FP8.")
if use_int4_w4a16:
group_size = get_quantization_group_size(config)
if group_size is None:
@@ -1092,6 +1130,7 @@ def main(args: argparse.Namespace):
search_space,
block_quant_shape,
use_deep_gemm,
per_act_token_quant,
)
for batch_size in batch_sizes
],
@@ -1130,6 +1169,7 @@ def main(args: argparse.Namespace):
use_int4_w4a16,
block_quant_shape,
use_deep_gemm,
per_act_token_quant,
)
for batch_size in batch_sizes
],
benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py
+21 -2
View File
@@ -16,6 +16,16 @@ import vllm._custom_ops as ops
HEAD_DIM = 64
DEFAULT_TOKENS = [1, 8, 64, 256, 1024, 4096, 8192, 16384, 32768, 65536]
try:
from vllm.models.muse_spark.nvidia.ops.cute_dsl.qk_norm import (
qk_norm,
supports_qk_norm,
)
HAS_CUTEDSL = supports_qk_norm(torch.bfloat16, HEAD_DIM)
except ImportError:
HAS_CUTEDSL = False
def make_cos_sin_cache(max_pos, rotary_dim, base, device):
inv_freq = 1.0 / (
@@ -135,6 +145,12 @@ class Case:
)
self._insert()
def unfused_nope(self):
q, k, _ = self.views()
qk_norm(q, HEAD_DIM, self.eps)
qk_norm(k, HEAD_DIM, self.eps)
self._insert()
def fused_bytes(self):
el = self.qkv.element_size()
cache_el = self.kv_cache.element_size()
@@ -260,7 +276,7 @@ def main():
f"dtype={args.dtype} kv_cache={args.kv_cache_dtype}\n"
)
modes = ["rope"]
modes = ["rope", "nope"]
l2 = _l2_bytes(device)
print(
f"rows marked * have a working set inside the {l2 / 2**20:.0f} MiB L2, "
@@ -274,6 +290,9 @@ def main():
for mode, use_graph in itertools.product(modes, [False, True]):
if use_graph and args.no_graph:
continue
if mode == "nope" and not HAS_CUTEDSL:
print("skipping the NoPE baseline: CuTe-DSL qk_norm unavailable")
continue
label = "graph" if use_graph else "eager"
print(f"── {mode} layer, {label} ──")
print(header)
@@ -282,7 +301,7 @@ def main():
case = Case(nt, args.nq, args.nkv, dtype, args.kv_cache_dtype, device)
use_rope = mode == "rope"
fused_us = timer(lambda c=case, r=use_rope: c.fused(r))
baseline = case.unfused_rope
baseline = case.unfused_rope if use_rope else case.unfused_nope
unfused_us = timer(baseline)
tuned_us = (
min(timer(lambda c=case, t=t: c.unfused_rope(t)) for t in (4, 8))
benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py
+538
View File
@@ -0,0 +1,538 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Benchmark MXFP4 MoE activation quant at GLM-5.3-shaped sizes.
Sweeps expert count at a fixed hidden width (default 6144). Two timed regions
(``--timed-region``):
* ``quant``: ``mxfp4_experts_quant`` / ``silu_and_mul_mxfp4_experts_quant`` only,
with synthetic routing offsets (kernel isolation).
* ``metadata_and_quant``: ``get_cutlass_moe_mm_data`` → ``shuffle_rows`` → quant.
* ``both``: run ``quant`` then ``metadata_and_quant``.
Routing uses 32 rows (4 tokens, topk 8, two picks per expert 0..15). GEMMs are
outside the timed region.
Run from a CUDA-built vLLM checkout, for example:
.venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py
.venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py \\
--timed-region both --experts 16,32,64,128,256
"""
from __future__ import annotations
import argparse
import sys
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
import torch
REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))
import vllm._custom_ops # noqa: E402, F401 -- registers torch.ops._C kernels
from vllm import _custom_ops as ops # noqa: E402
from vllm.triton_utils import triton # noqa: E402
ROUTED_ROWS = 32
TOPK = 8
DEFAULT_HIDDEN_SIZE = 6144
DEFAULT_MOE_INTERMEDIATE_SIZE = 2048
DEFAULT_TP_SIZE = 4
DEFAULT_EXPERT_COUNTS = (16, 32, 64, 128, 256)
class TimedRegion(str, Enum):
QUANT = "quant"
METADATA_AND_QUANT = "metadata_and_quant"
BOTH = "both"
@dataclass(frozen=True)
class Case:
kernel: str
output_width: int
fused_silu_mul: bool
experts: int
mm_n: int
mm_k: int
is_gated: bool
def routing_offsets(
num_experts: int,
rows: int = ROUTED_ROWS,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""32 rows, two assignments on experts 0..15 (production-like small decode)."""
if rows != ROUTED_ROWS:
raise ValueError(f"expected {ROUTED_ROWS} routed rows, got {rows}")
if num_experts < 16:
raise ValueError("expert sweep must use at least 16 experts")
counts = [0] * num_experts
counts[:16] = [2] * 16
row_offsets = [0]
scale_offsets = [0]
expert_ids: list[int] = []
for expert, count in enumerate(counts):
expert_ids.extend([expert] * count)
row_offsets.append(row_offsets[-1] + count)
scale_offsets.append(scale_offsets[-1] + ((count + 127) // 128) * 128)
assert row_offsets[-1] == rows
return (
torch.tensor(row_offsets, dtype=torch.int32, device="cuda"),
torch.tensor(scale_offsets, dtype=torch.int32, device="cuda"),
torch.tensor(expert_ids, dtype=torch.int32, device="cuda"),
scale_offsets[-1],
)
def decode_topk_ids(rows: int = ROUTED_ROWS) -> torch.Tensor:
"""Two picks per expert 0..15 across 4 tokens (topk=8); shuffled (fixed seed)."""
if rows % TOPK != 0:
raise ValueError(f"routed rows {rows} must be divisible by topk {TOPK}")
expert_ids: list[int] = []
for expert in range(16):
expert_ids.extend([expert, expert])
assert len(expert_ids) == rows
gen = torch.Generator(device="cuda")
gen.manual_seed(0)
perm = torch.randperm(rows, generator=gen, device="cuda")
ordered = torch.tensor(expert_ids, device="cuda", dtype=torch.int32)
return ordered[perm].view(rows // TOPK, TOPK)
def benchmark_invoke(invoke, rep: int) -> float:
"""Median latency in microseconds (CUDA graph replay via Triton)."""
quantiles = [0.5, 0.2, 0.8]
ms, _, _ = triton.testing.do_bench_cudagraph(invoke, rep=rep, quantiles=quantiles)
return ms * 1_000
def _quant_tensors(
case: Case,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
scale_rows = int(blockscale_offsets[-1].item())
scale_cols_int32 = (case.output_width // 32 + 3) // 4
output = torch.empty(
(ROUTED_ROWS, case.output_width // 2), dtype=torch.uint8, device="cuda"
)
output_scales = torch.zeros(
(scale_rows, scale_cols_int32), dtype=torch.int32, device="cuda"
)
return output, output_scales
def benchmark_mxfp4_quant_only(
case: Case,
inputs: torch.Tensor,
rep: int,
sorted_first: bool,
) -> tuple[float, float]:
expert_offsets, blockscale_offsets, expert_ids, _ = routing_offsets(case.experts)
output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets)
op = (
torch.ops._C.silu_and_mul_mxfp4_experts_quant
if case.fused_silu_mul
else torch.ops._C.mxfp4_experts_quant
)
def invoke_offset() -> None:
op(
output,
output_scales,
inputs,
expert_offsets,
blockscale_offsets,
case.experts,
None,
)
def invoke_sorted() -> None:
op(
output,
output_scales,
inputs,
expert_offsets,
blockscale_offsets,
case.experts,
expert_ids,
)
if sorted_first:
direct_id_us = benchmark_invoke(invoke_sorted, rep)
baseline_us = benchmark_invoke(invoke_offset, rep)
else:
baseline_us = benchmark_invoke(invoke_offset, rep)
baseline_output = output.clone()
baseline_scales = output_scales.clone()
direct_id_us = benchmark_invoke(invoke_sorted, rep)
if not (
torch.equal(output, baseline_output)
and torch.equal(output_scales, baseline_scales)
):
raise AssertionError(
f"MXFP4 output mismatch for {case.kernel} E={case.experts}"
)
return baseline_us, direct_id_us
def benchmark_mxfp4_metadata_and_quant(
case: Case,
inputs: torch.Tensor,
rep: int,
sorted_first: bool,
) -> tuple[float, float]:
num_experts = case.experts
topk_ids = decode_topk_ids()
expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda")
blockscale_offsets = torch.empty_like(expert_offsets)
problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda")
problem_sizes2 = torch.empty_like(problem_sizes1)
input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda")
output_permutation = torch.empty_like(input_permutation)
sorted_expert_ids = torch.empty_like(input_permutation)
torch.ops._C.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
input_permutation,
output_permutation,
num_experts,
case.mm_n,
case.mm_k,
blockscale_offsets,
case.is_gated,
sorted_expert_ids,
)
output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets)
op = (
torch.ops._C.silu_and_mul_mxfp4_experts_quant
if case.fused_silu_mul
else torch.ops._C.mxfp4_experts_quant
)
def invoke_offset() -> None:
torch.ops._C.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
input_permutation,
output_permutation,
num_experts,
case.mm_n,
case.mm_k,
blockscale_offsets,
case.is_gated,
None,
)
sorted_inputs = ops.shuffle_rows(inputs, input_permutation)
op(
output,
output_scales,
sorted_inputs,
expert_offsets,
blockscale_offsets,
num_experts,
None,
)
def invoke_sorted() -> None:
torch.ops._C.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
input_permutation,
output_permutation,
num_experts,
case.mm_n,
case.mm_k,
blockscale_offsets,
case.is_gated,
sorted_expert_ids,
)
sorted_inputs = ops.shuffle_rows(inputs, input_permutation)
op(
output,
output_scales,
sorted_inputs,
expert_offsets,
blockscale_offsets,
num_experts,
sorted_expert_ids,
)
if sorted_first:
direct_id_us = benchmark_invoke(invoke_sorted, rep)
baseline_us = benchmark_invoke(invoke_offset, rep)
else:
baseline_us = benchmark_invoke(invoke_offset, rep)
baseline_output = output.clone()
baseline_scales = output_scales.clone()
direct_id_us = benchmark_invoke(invoke_sorted, rep)
if not (
torch.equal(output, baseline_output)
and torch.equal(output_scales, baseline_scales)
):
raise AssertionError(
f"MXFP4 output mismatch for {case.kernel} E={case.experts}"
)
return baseline_us, direct_id_us
def verify_sorted_id_writer(num_experts: int) -> None:
"""Check the routing sort writes the expected sorted expert ID at each row."""
topk_ids = torch.arange(ROUTED_ROWS, device="cuda", dtype=torch.int32) % num_experts
topk_ids = topk_ids.view(4, 8)
expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda")
blockscale_offsets = torch.empty_like(expert_offsets)
problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda")
problem_sizes2 = torch.empty_like(problem_sizes1)
input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda")
output_permutation = torch.empty_like(input_permutation)
sorted_expert_ids = torch.empty_like(input_permutation)
torch.ops._C.get_cutlass_moe_mm_data(
topk_ids,
expert_offsets,
problem_sizes1,
problem_sizes2,
input_permutation,
output_permutation,
num_experts,
64,
64,
blockscale_offsets,
True,
sorted_expert_ids,
)
expected = torch.sort(topk_ids.flatten()).values
torch.accelerator.synchronize()
if not torch.equal(sorted_expert_ids, expected):
raise AssertionError("routing sort produced incorrect sorted expert IDs")
def parse_expert_counts(raw: str) -> tuple[int, ...]:
counts = tuple(int(x.strip()) for x in raw.split(",") if x.strip())
if not counts or any(c < 16 for c in counts):
raise ValueError(
"expert counts must be a comma-separated list of integers >= 16"
)
return counts
def parse_timed_region(raw: str) -> TimedRegion:
try:
return TimedRegion(raw)
except ValueError:
choices = ", ".join(r.value for r in TimedRegion)
raise ValueError(f"timed-region must be one of: {choices}") from None
def intermediate_dims(moe_intermediate_size: int, tp_size: int) -> tuple[int, int]:
if moe_intermediate_size % tp_size != 0:
raise ValueError(
f"moe intermediate size {moe_intermediate_size} must be divisible "
f"by tp size {tp_size}"
)
n_per_rank = moe_intermediate_size // tp_size
return 2 * n_per_rank, n_per_rank
def build_cases(
experts: tuple[int, ...],
hidden_size: int,
moe_intermediate_size: int,
tp_size: int,
) -> list[tuple[Case, torch.Tensor]]:
intermediate_in, intermediate_out = intermediate_dims(
moe_intermediate_size, tp_size
)
cases: list[tuple[Case, torch.Tensor]] = []
hidden_input = torch.randn(
(ROUTED_ROWS, hidden_size), device="cuda", dtype=torch.bfloat16
)
intermediate_input = torch.randn(
(ROUTED_ROWS, intermediate_in), device="cuda", dtype=torch.bfloat16
)
w1_n = 2 * intermediate_out
for expert_count in experts:
cases.append(
(
Case(
"mxfp4_experts_quant",
hidden_size,
False,
expert_count,
mm_n=w1_n,
mm_k=hidden_size,
is_gated=True,
),
hidden_input,
)
)
cases.append(
(
Case(
"silu_and_mul_mxfp4_experts_quant",
intermediate_out,
True,
expert_count,
mm_n=hidden_size,
mm_k=intermediate_out,
is_gated=True,
),
intermediate_input,
)
)
return cases
def run_region(
region: TimedRegion,
cases: list[tuple[Case, torch.Tensor]],
rep: int,
sorted_first: bool,
) -> None:
print(f"timed_region={region.value}")
print("kernel,experts,variant,median_us,speedup_vs_offset_search")
bench_fn = (
benchmark_mxfp4_metadata_and_quant
if region == TimedRegion.METADATA_AND_QUANT
else benchmark_mxfp4_quant_only
)
for case, inputs in cases:
baseline_us, direct_id_us = bench_fn(case, inputs, rep, sorted_first)
print(f"{case.kernel},{case.experts},offset_search,{baseline_us:.4f},1.0000")
print(
f"{case.kernel},{case.experts},sorted_expert_id,"
f"{direct_id_us:.4f},{baseline_us / direct_id_us:.4f}"
)
print()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--rep",
type=int,
default=100,
help="Repeat count (default: 100)",
)
parser.add_argument(
"--hidden-size",
type=int,
default=DEFAULT_HIDDEN_SIZE,
help=f"hidden activation width K (default: {DEFAULT_HIDDEN_SIZE})",
)
parser.add_argument(
"--tp-size",
type=int,
default=DEFAULT_TP_SIZE,
help=(
"tensor-parallel size; shards moe intermediate per rank "
f"(default: {DEFAULT_TP_SIZE})"
),
)
parser.add_argument(
"--moe-intermediate-size",
type=int,
default=DEFAULT_MOE_INTERMEDIATE_SIZE,
help=(
"total MoE intermediate size before TP; per-rank N is this divided "
f"by --tp-size (default: {DEFAULT_MOE_INTERMEDIATE_SIZE})"
),
)
parser.add_argument(
"--experts",
type=parse_expert_counts,
default=DEFAULT_EXPERT_COUNTS,
help=(
"comma-separated expert counts (default: "
f"{','.join(map(str, DEFAULT_EXPERT_COUNTS))})"
),
)
parser.add_argument(
"--timed-region",
type=parse_timed_region,
default=TimedRegion.QUANT,
help=(
"quant: activation quant only; metadata_and_quant: mm_data + shuffle + "
"quant; both: run each region"
),
)
parser.add_argument(
"--sorted-id-first",
action="store_true",
help="time sorted-expert-ID path before offset search (default: offset first)",
)
parser.add_argument("--skip-metadata-check", action="store_true")
args = parser.parse_args()
if args.rep < 1:
parser.error("--rep must be positive")
if args.hidden_size < 128 or args.hidden_size % 128 != 0:
parser.error("--hidden-size must be a multiple of 128 and >= 128")
if args.tp_size < 1:
parser.error("--tp-size must be positive")
if args.moe_intermediate_size < args.tp_size:
parser.error("--moe-intermediate-size must be >= --tp-size")
if args.moe_intermediate_size % args.tp_size != 0:
parser.error("--moe-intermediate-size must be divisible by --tp-size")
try:
intermediate_in, intermediate_out = intermediate_dims(
args.moe_intermediate_size, args.tp_size
)
except ValueError as exc:
parser.error(str(exc))
if intermediate_out < 32 or intermediate_out % 32 != 0:
parser.error(
"per-rank MoE intermediate (moe-intermediate-size / tp-size) "
"must be a multiple of 32 and >= 32"
)
if not torch.cuda.is_available():
parser.error("requires a CUDA GPU")
capability = torch.cuda.get_device_capability()
if capability[0] not in (10, 11):
parser.error(f"requires an SM10x/11x GPU, got SM{capability[0]}{capability[1]}")
torch.manual_seed(0)
if not args.skip_metadata_check:
verify_sorted_id_writer(num_experts=256)
cases = build_cases(
args.experts,
args.hidden_size,
args.moe_intermediate_size,
args.tp_size,
)
print(f"GPU: {torch.cuda.get_device_name()} (SM{capability[0]}{capability[1]})")
print(
f"hidden={ROUTED_ROWS}x{args.hidden_size}, "
f"intermediate={ROUTED_ROWS}x{intermediate_in}->{intermediate_out} "
f"(moe_intermediate={args.moe_intermediate_size}, tp={args.tp_size}), "
f"topk={TOPK}"
)
regions: list[TimedRegion]
if args.timed_region == TimedRegion.BOTH:
regions = [TimedRegion.QUANT, TimedRegion.METADATA_AND_QUANT]
else:
regions = [args.timed_region]
for region in regions:
run_region(region, cases, args.rep, args.sorted_id_first)
if __name__ == "__main__":
main()
benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py
+12 -6
View File
@@ -54,7 +54,7 @@ WEIGHT_SHAPES = {
# ---------------------------------------------------------------------------
# Weight packing
# ---------------------------------------------------------------------------
def prepare_hybrid_weights(K, N, group_size, device="cuda"):
def prepare_hybrid_weights(K, N, group_size, dtype=torch.float16, device="cuda"):
"""Create random weights for benchmarking.
Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
@@ -69,7 +69,7 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"):
unpacked = torch.randint(0, 16, (N, K), dtype=torch.int32, device=device)
w_q_skinny = pack_skinny_int4(unpacked)
del unpacked
w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
w_s_skinny = torch.randn(N, num_groups, dtype=dtype, device=device) * 0.01
# Per-group zero-points for asymmetric benchmarks, in the packed layout the
# kernels read: [N//8, num_groups] int32, row n's nibble at bits 4*(n%8).
@@ -80,7 +80,7 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"):
).contiguous()
# FP16 baseline for F.linear
w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
w_fp16 = torch.randn(N, K, dtype=dtype, device=device) * 0.01
return w_q_skinny, w_s_skinny, w_fp16, w_zp
@@ -104,7 +104,7 @@ PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
args={},
)
)
def benchmark(batch_size, provider, N, K, group_size, weights):
def benchmark(batch_size, provider, N, K, group_size, dtype, weights):
"""Time one GEMM.
``do_bench`` rather than ``do_bench_cudagraph``: it zeroes a 256 MB buffer
@@ -114,7 +114,6 @@ def benchmark(batch_size, provider, N, K, group_size, weights):
"""
M = batch_size
device = "cuda"
dtype = torch.float16
a = torch.randn((M, K), device=device, dtype=dtype)
quantiles = [0.5, 0.2, 0.8]
@@ -182,8 +181,12 @@ if __name__ == "__main__":
)
parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
parser.add_argument("--group-size", type=int, default=128)
parser.add_argument(
"--dtype", type=str, default="float16", choices=["float16", "bfloat16"]
)
parser.add_argument("--save-path", type=str, default=None)
args = parser.parse_args()
dtype = getattr(torch, args.dtype)
for K, N, model in prepare_shapes(args):
group_size = args.group_size
@@ -191,7 +194,9 @@ if __name__ == "__main__":
print(f"{model}, N={N} K={K}, group_size={group_size}")
print(f"{'=' * 70}")
w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(
K, N, group_size, dtype
)
weights = {
"w_q_skinny": w_q_skinny,
@@ -209,6 +214,7 @@ if __name__ == "__main__":
N=N,
K=K,
group_size=group_size,
dtype=dtype,
weights=weights,
)
cmake/cpu_extension.cmake
+3
View File
@@ -502,6 +502,7 @@ set(VLLM_EXT_SRC
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/cpu_isa.cpp"
"csrc/cpu/sampling_kernels.cpp"
"csrc/cpu/fp8_quant.cpp"
"csrc/cpu/torch_bindings.cpp")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
@@ -593,6 +594,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/sampling_kernels.cpp"
"csrc/cpu/fp8_quant.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
"csrc/cpu/activation.cpp"
@@ -611,6 +613,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/sampling_kernels.cpp"
"csrc/cpu/fp8_quant.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
"csrc/cpu/activation.cpp"
cmake/external_projects/deepgemm.cmake
+5 -2
View File
@@ -32,9 +32,12 @@ else()
# plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port,
# configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4
# Mega MoE, plus swiglustep for BF16/FP8/NVFP4 Mega MoE and F8_128x4
# SFs for the FP8 einsum output (vllm-project/DeepGEMM#40).
# SFs for the FP8 einsum output (vllm-project/DeepGEMM#40), plus
# fp4_fp4_mega_moe and fused post-expert RMSNorm for Muse Spark Open
# (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix
# (vllm-project/DeepGEMM#43).
# Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API.
set(_DEEPGEMM_UPSTREAM_TAG "727f847082d2e7946f70dc02b94949db206310d3")
set(_DEEPGEMM_UPSTREAM_TAG "0820030756945b2daec208e8e8255268360f337d")
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _deepgemm_fc_root)
cmake/hipify.py
+40 -6
View File
@@ -11,8 +11,10 @@
#
import argparse
import json
import os
import shutil
from pathlib import Path
from torch.utils.hipify.hipify_python import get_hip_file_path, hipify
@@ -48,14 +50,29 @@ if __name__ == "__main__":
parser.add_argument(
"sources", help="Source files to hipify.", nargs="*", default=[]
)
parser.add_argument(
"--manifest", help="Write generated-to-original source/header provenance."
)
args = parser.parse_args()
# Limit include scope to project_dir only
includes = [os.path.join(args.project_dir, "*")]
# Get absolute path for all source files.
extra_files = [os.path.abspath(s) for s in args.sources]
project_dir = Path(args.project_dir).resolve()
output_dir = Path(args.output_dir).resolve()
provenance = {
str(output_dir / source.relative_to(project_dir)): str(source)
for source in project_dir.rglob("*")
if source.is_file()
}
source_copies = {
os.path.abspath(source): str(
output_dir / Path(source).resolve().relative_to(project_dir)
)
for source in args.sources
}
extra_files = list(source_copies.values())
# Copy sources from project directory to output directory.
# The directory might already exist to hold object files so we ignore that.
@@ -75,16 +92,24 @@ if __name__ == "__main__":
is_pytorch_extension=True,
hipify_extra_files_only=True,
)
copied_sources = dict(provenance)
for source, result in hipify_result.items():
if result.hipified_path is not None:
original = copied_sources.get(
os.path.abspath(source), os.path.abspath(source)
)
provenance[os.path.abspath(result.hipified_path)] = original
hipified_sources = []
for source in args.sources:
s_abs = os.path.abspath(source)
if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
path = hipify_result[s_abs].hipified_path
copied = source_copies[s_abs]
if copied in hipify_result and hipify_result[copied].hipified_path is not None:
path = hipify_result[copied].hipified_path
# PyTorch skips writing when is_pytorch_extension and text unchanged;
# hipified_path then stays *.cu. CMake expects *.hip under output_dir.
if s_abs.endswith(".cu") and path.endswith(".cu"):
dest = _expected_hip_build_path(s_abs, args.output_dir)
dest = _expected_hip_build_path(copied, args.output_dir)
if os.path.normpath(path) != os.path.normpath(dest):
os.makedirs(os.path.dirname(dest), exist_ok=True)
shutil.copy2(path, dest)
@@ -92,10 +117,19 @@ if __name__ == "__main__":
else:
hipified_s_abs = path
else:
hipified_s_abs = s_abs
hipified_s_abs = copied
hipified_sources.append(hipified_s_abs)
provenance[os.path.abspath(hipified_s_abs)] = s_abs
assert len(hipified_sources) == len(args.sources)
manifest = (
Path(args.manifest) if args.manifest else output_dir / "hipify-source-map.json"
)
manifest.parent.mkdir(parents=True, exist_ok=True)
temporary = manifest.with_suffix(manifest.suffix + ".tmp")
temporary.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n")
temporary.replace(manifest)
# Print hipified source files.
print("\n".join(hipified_sources))
cmake/utils.cmake
+5 -2
View File
@@ -113,11 +113,14 @@ function (vllm_finalize_hipify_target)
list(REMOVE_DUPLICATES ALL_BYPRODUCTS)
set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
set(HIPIFY_MANIFEST ${CMAKE_CURRENT_BINARY_DIR}/hipify-source-map.json)
add_custom_target(
hipify_all
COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS}
COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} --manifest ${HIPIFY_MANIFEST} ${ALL_SRCS}
DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS}
BYPRODUCTS ${ALL_BYPRODUCTS}
BYPRODUCTS ${ALL_BYPRODUCTS} ${HIPIFY_MANIFEST}
WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
VERBATIM
COMMENT "Running hipify on all extension source files.")
endfunction()
csrc/attention/octave_arch.cuh
+36
View File
@@ -0,0 +1,36 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Per-architecture primitives of the Octave kernels. Every AMD target vLLM
// builds compiles the portable tier; faster tiers are selected at runtime by
// the host code (see docs/design/rocm_octave.md, "Architecture contract").
//
// int8 dot4: v_dot4_i32_iu8 on gfx11/gfx12, v_dot4_i32_i8 on gfx9/gfx10.3
// decode, fast tier: octave_decode_wmma_rdna3.cu (WMMA, gfx11)
// prefill: octave_prefill_rdna3.cu (WMMA, gfx11)
// decode, portable: octave_attn.cu decode_kernel (dot4, every target)
//
// The portable kernels are written for 32-lane logical waves: shuffles never
// cross 32 lanes, so they are correct on wave64 (CDNA) as well.
#pragma once
#include <cstdint>
namespace octave {
// sum_i a.i8[i] * b.i8[i] + c, signed bytes on both sides.
__device__ __forceinline__ int dot4_i8(int a, int b, int c) {
#if defined(__GFX11__) || defined(__GFX12__)
return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
#elif defined(__gfx942__)
return __builtin_amdgcn_sdot4(a, b, c, false);
#else
#pragma unroll
for (int i = 0; i < 4; ++i)
c += (int)(int8_t)(a >> (8 * i)) * (int)(int8_t)(b >> (8 * i));
return c;
#endif
}
} // namespace octave
csrc/attention/octave_attn.cu
+1180
View File
File diff suppressed because it is too large. Load diff
csrc/attention/octave_decode_mfma_cdna.cu
+719
View File
@@ -0,0 +1,719 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Octave decode on MFMA for CDNA3 (gfx942). Same scheme as the WMMA decode
// (octave_decode_wmma_rdna3.cu): independent waves, one (split, 16-query row
// tile) per wave, transposed orientation, codebook bytes fed to fp16 as
// x + 1152 and debiased afterwards. With 16x16x16 MFMA on wave64 (lane l
// holds column l % 16 and rows 4 (l / 16) .. + 3 of the result):
//
// QK S^T[token][q] = K x Q^T. A lane owns one token (l % 16) and, in the
// four k-steps of each 64-dim block, the 16 contiguous dims
// 16 (l / 16) .. + 15 of that block, so its codes are one contiguous
// run of the slot. One accumulator per K scale.
// PV O^T[dim][q] = V^T x P^T. Output row r of fragment c is dim 16 r + c,
// so a lane reads the 16 contiguous dims 16 (l % 16) .. + 15 of tokens
// 4 (l / 16) .. + 3. P^T is the S^T accumulator, one query per lane.
#include <algorithm>
#include <cstdint>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#if defined(USE_ROCM)
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include "octave_format.cuh"
#if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__)
#define OCTAVE_MFMA_STUB
#endif
namespace octave_mfma {
typedef _Float16 h4v __attribute__((ext_vector_type(4)));
typedef _Float16 h8v __attribute__((ext_vector_type(8)));
typedef _Float16 h2v __attribute__((ext_vector_type(2)));
typedef float f4v __attribute__((ext_vector_type(4)));
typedef int i4v __attribute__((ext_vector_type(4)));
typedef const int __attribute__((address_space(4))) cint;
using octave::D;
using octave::Format;
using octave::LUT_ONE;
[[maybe_unused]] constexpr int QROW = D + 8;
[[maybe_unused]] constexpr float PSCALE = 256.0f;
[[maybe_unused]] constexpr float BIAS = 1152.0f; // byte x fed as x + 1152
[[maybe_unused]] constexpr uint32_t X = 0x80808080u;
// Bytes of a 16-token tile (slot <= 256).
[[maybe_unused]] constexpr int TILE_MAX = 4096;
// Bytes b of a dword -> fp16 1024 + b, exact.
__device__ __forceinline__ h4v bytes_to_h4(uint32_t w) {
const uint32_t lo = __builtin_amdgcn_perm(0x64646464u, w, 0x07010700u);
const uint32_t hi = __builtin_amdgcn_perm(0x64646464u, w, 0x07030702u);
return __builtin_bit_cast(h4v, (uint2){lo, hi});
}
template <bool BIASED>
__device__ __forceinline__ uint32_t lut4_fast(uint32_t codes) {
constexpr uint32_t B = BIASED ? X : 0u;
const uint32_t lo = __builtin_amdgcn_perm(
octave::LUT4_1 ^ B, octave::LUT4_0 ^ B, codes & 0x07070707u);
const uint32_t hi = __builtin_amdgcn_perm(
octave::LUT4_3 ^ B, octave::LUT4_2 ^ B, codes & 0x07070707u);
const uint32_t t = codes & 0x08080808u;
const uint32_t m = (t << 5) - (t >> 3);
return (hi & m) | (lo & ~m);
}
// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each.
__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
uint32_t a3, uint32_t o[4]) {
const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
}
template <typename T>
__device__ __forceinline__ float to_f(T x) {
return (float)x;
}
template <>
__device__ __forceinline__ float to_f<half>(half x) {
return __half2float(x);
}
__device__ __forceinline__ float sign_of(const int* bits, int i) {
return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f;
}
__device__ __forceinline__ float lo16(uint32_t x) {
return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF)));
}
__device__ __forceinline__ float hi16(uint32_t x) {
return __half2float(__ushort_as_half((unsigned short)(x >> 16)));
}
// Over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48).
__device__ __forceinline__ float qmax4(float v) {
v = fmaxf(v, __shfl_xor(v, 16));
return fmaxf(v, __shfl_xor(v, 32));
}
__device__ __forceinline__ float qsum4(float v) {
v += __shfl_xor(v, 16);
return v + __shfl_xor(v, 32);
}
// QG: query tokens per block (4 for multi-token verification, else 1).
template <int VB, bool KC, int NW, int QG, typename QT>
__global__ __launch_bounds__(64 * NW) void decode_mfma(
const QT* __restrict__ Q, const uint8_t* __restrict__ cache,
const int* __restrict__ block_table, const int* __restrict__ q_to_req,
const int* __restrict__ q_to_klen, float* __restrict__ mid_o,
const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads,
int num_kv_heads, int block_size, int max_blocks, int num_reqs,
int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles,
int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct,
int64_t smo, int64_t smh, int64_t sms) {
#ifndef OCTAVE_MFMA_STUB
using F = Format<VB, KC>;
constexpr int NG = KC ? 1 : 4; // QK accumulators (one per K scale)
const int grp = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int w = __builtin_amdgcn_readfirstlane(tid >> 6), lane = tid & 63;
const int col = lane & 15, gq = lane >> 4;
const int rt = num_row_tiles == 2 ? (w & 1) : 0;
const int splits_per_block = NW / num_row_tiles;
const int si =
blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w);
const int hpk = num_q_heads / num_kv_heads;
const int max_kv = max_blocks * block_size;
const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE;
// The staged queries, then (once they are in registers) one 4 KB tile
// buffer per wave.
constexpr int SMEM =
32 * QROW * 2 > NW * TILE_MAX ? 32 * QROW * 2 : NW * TILE_MAX;
__shared__ __attribute__((aligned(16))) uint8_t smem[SMEM];
auto sQ = reinterpret_cast<_Float16 (*)[QROW]>(smem);
// Stage the group's queries rotated like K: each 32-lane half owns one
// query vector, 8 dims per lane.
{
const int l32 = lane & 31;
for (int v = 2 * w + (lane >> 5); v < 32; v += 2 * NW) {
const int r = v / hpk, g = v % hpk;
const int qi = grp * QG + r;
float x[8];
const bool live = v < QG * hpk && qi < num_q;
#pragma unroll
for (int i = 0; i < 8; ++i)
x[i] = live ? to_f<QT>(Q[(int64_t)qi * sq0 +
(int64_t)(kvh * hpk + g) * sq1 + l32 * 8 + i])
: 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, l32 * 8 + i);
#pragma unroll
for (int hh = 1; hh < 8; hh <<= 1) {
#pragma unroll
for (int i = 0; i < 8; ++i)
if ((i & hh) == 0) {
const float a = x[i], b = x[i + hh];
x[i] = a + b;
x[i + hh] = a - b;
}
}
#pragma unroll
for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) {
const bool hi = l32 & m;
#pragma unroll
for (int i = 0; i < 8; ++i) {
const float o = __shfl_xor(x[i], m);
x[i] = hi ? o - x[i] : x[i] + o;
}
}
#pragma unroll
for (int i = 0; i < 8; ++i)
sQ[v][l32 * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f));
}
}
__syncthreads();
int rreq[QG], rlen[QG];
#pragma unroll
for (int r = 0; r < QG; ++r) {
const int qi = grp * QG + r;
int rq = -1, ln = 0;
if (qi < num_q) {
rq = q_to_req[qi];
ln = q_to_klen[qi];
if (rq < 0 || rq >= num_reqs) {
rq = 0;
ln = 0;
}
ln = min(max(ln, 0), max_kv);
}
rreq[r] = __builtin_amdgcn_readfirstlane(rq);
rlen[r] = __builtin_amdgcn_readfirstlane(ln);
}
// This lane's query (column) as int8 B fragments for the i8 MFMA: step
// 2 b + h holds dims 64 b + 16 gq + 8 h .. + 7 (the dot product does not
// care how dims map to k, only that K uses the same map). Two int8 halves,
// q ~ qsc (hi + lo / 256), keep ~16 bits of the query.
const int qv = rt * 16 + col;
long Qh[8], Ql[8];
float qsc;
{
float x[64];
float qmax = 1e-20f;
#pragma unroll
for (int i = 0; i < 64; ++i) {
x[i] = (float)sQ[qv][64 * (i >> 4) + 16 * gq + (i & 15)];
qmax = fmaxf(qmax, fabsf(x[i]));
}
qmax = qmax4(qmax);
const float inv = 127.0f / qmax;
qsc = qmax * (1.0f / 127.0f);
#pragma unroll
for (int st = 0; st < 8; ++st) {
uint64_t wh = 0, wl = 0;
#pragma unroll
for (int e = 0; e < 8; ++e) {
const float y = x[8 * st + e] * inv;
const float h = rintf(y);
const int lo = max(-127, min(127, (int)rintf((y - h) * 256.0f)));
wh |= (uint64_t)(uint8_t)(int8_t)(int)h << (8 * e);
wl |= (uint64_t)(uint8_t)(int8_t)lo << (8 * e);
}
Qh[st] = (long)wh;
Ql[st] = (long)wl;
}
}
const float sm_q = sm_log2 * qsc * (1.0f / 256.0f);
__syncthreads();
uint8_t* lt = smem + w * TILE_MAX;
// rreq/rlen are only ever indexed by constants (unrolled selections): a
// dynamic index would put them in scratch memory.
auto pick = [&](const int* a, int i) {
int x = a[0];
#pragma unroll
for (int r = 1; r < QG; ++r) x = i == r ? a[r] : x;
return x;
};
int seg0 = 0;
while (seg0 < QG && pick(rreq, seg0) >= 0) {
const int req = pick(rreq, seg0);
int seg1 = seg0 + 1;
while (seg1 < QG && pick(rreq, seg1) == req) ++seg1;
int Lmax = 0;
#pragma unroll
for (int r = 0; r < QG; ++r)
if (r >= seg0 && r < seg1) Lmax = max(Lmax, rlen[r]);
const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk;
int tps = (Lmax + num_splits - 1) / num_splits;
if (tps < min_tps) tps = min_tps;
tps = (tps + 15) & ~15;
const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16;
const bool qlive = qv >= qv_lo && qv < qv_hi;
if (blockIdx.z * splits_per_block * tps >= Lmax) {
if (tile_live && si < num_splits && gq == 0 && qlive) {
float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo +
(int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms;
op[D] = -INFINITY;
op[D + 1] = 0.0f;
}
seg0 = seg1;
continue;
}
if (tile_live && si < num_splits) {
int qlen = 0;
#pragma unroll
for (int r = 0; r < QG; ++r)
if (qlive && qv / hpk == r) qlen = rlen[r];
const int start = si * tps;
const int end = min(start + tps, Lmax);
float m = -INFINITY, l = 0.0f, cp = 0.0f;
f4v O[16];
#pragma unroll
for (int c = 0; c < 16; ++c) O[c] = (f4v){0, 0, 0, 0};
const uint8_t* cache_h = cache + (int64_t)kvh * sch;
auto tile_ptr = [&](int pos) {
const int lb = pos / block_size;
int pb = ((cint*)block_table)[req * max_blocks + lb];
if (pb < 0 || pb >= num_phys_blocks) pb = 0;
return cache_h + (int64_t)pb * scb +
(int64_t)(pos - lb * block_size) * sct;
};
// Raw K of token col: per 64-dim block, the codes of dims
// 16 gq .. + 15. 4-bit: 8 bytes. 3-bit: 2-bit plane word gq and the
// 16 bits of the 1-bit plane. Scales of tokens 4 gq .. + 3.
constexpr int NB3 = KC ? 4 : 3; // 3-bit blocks
uint2 kr = {0, 0};
uint2 kn[NB3];
uint32_t sa[4], sb[4], sv[4];
// A 16-token tile is one contiguous run of a block: each lane fetches
// four 16-byte chunks of the next tile into registers, then stores
// them to its wave's LDS buffer, where the fragments are read.
constexpr int tb = 16 * F::SLOT;
uint4 raw[TILE_MAX / 1024], raw2[TILE_MAX / 1024];
auto fetch = [&](uint4* dst, const uint8_t* tile) {
#pragma unroll
for (int j = 0; j < TILE_MAX / 1024; ++j) {
const int o = 16 * (lane + 64 * j);
dst[j] = o < tb ? *(const uint4*)(tile + o) : make_uint4(0, 0, 0, 0);
}
};
auto stash = [&]() {
#pragma unroll
for (int j = 0; j < TILE_MAX / 1024; ++j) {
const int o = 16 * (lane + 64 * j);
if (o < tb) *(uint4*)(lt + o) = raw[j];
}
};
auto load_k = [&](const uint8_t* tile) {
const uint8_t* pk = tile + (int64_t)col * F::SLOT;
if constexpr (!KC) kr = *(const uint2*)(pk + 8 * gq);
#pragma unroll
for (int b = 0; b < NB3; ++b) {
kn[b].x = *(const uint32_t*)(pk + F::OFF_KN + 16 * b + 4 * gq);
kn[b].y = *(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 8 * b +
4 * (gq >> 1)) >>
(4 * (gq & 1));
}
#pragma unroll
for (int e = 0; e < 4; ++e) {
const uint8_t* ps =
tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_SC;
if constexpr (KC) {
sa[e] = *(const uint32_t*)ps;
sv[e] = sa[e] >> 16;
} else {
const uint2 u = *(const uint2*)ps;
sa[e] = u.x;
sb[e] = u.y;
sv[e] = *(const uint16_t*)(ps + 8);
}
}
};
// Raw V of tokens 4 gq .. + 3, dims 16 col .. + 15.
uint2 vr[4];
auto load_v = [&](const uint8_t* tile) {
#pragma unroll
for (int e = 0; e < 4; ++e) {
const uint8_t* p = tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_V;
if constexpr (VB == 4) {
vr[e] = *(const uint2*)(p + 8 * col);
} else {
vr[e].x = *(const uint32_t*)(p + 4 * col);
vr[e].y = *(const uint32_t*)(p + D / 4 + 4 * (col >> 1)) >>
(4 * (col & 1));
}
}
};
// Codebook bytes of dims 4 s .. + 3 of a 16-dim run: int8 for K, biased
// (x + 128, the bias folded into the tables) for V.
auto codes4 = [&](uint2 raw, int s) -> uint32_t {
const uint32_t wd = s < 2 ? raw.x : raw.y;
return (wd >> (4 * (s & 1))) & 0x0F0F0F0Fu;
};
auto codes3 = [&](uint2 raw, int s) -> uint32_t {
return ((raw.x >> (2 * s)) & 0x03030303u) |
(((raw.y >> s) & 0x01010101u) << 2);
};
auto k_bytes = [&](int b, int s) -> uint32_t {
return (!KC && b == 0) ? lut4_fast<false>(codes4(kr, s))
: octave::lut3(codes3(kn[KC ? b : b - 1], s));
};
auto v_bytes = [&](uint2 raw, int s) -> uint32_t {
if constexpr (VB == 4)
return lut4_fast<true>(codes4(raw, s));
else
return __builtin_amdgcn_perm(octave::LUT3_HI ^ X, octave::LUT3_LO ^ X,
codes3(raw, s));
};
if (start < end) fetch(raw, tile_ptr(start));
if (start + 16 < end) fetch(raw2, tile_ptr(start + 16));
for (int base = start; base < end; base += 16) {
stash();
load_k(lt);
#pragma unroll
for (int j = 0; j < TILE_MAX / 1024; ++j) raw[j] = raw2[j];
if (base + 32 < end) fetch(raw2, tile_ptr(base + 32));
i4v Sh[NG], Sl[NG];
#pragma unroll
for (int g = 0; g < NG; ++g) Sh[g] = Sl[g] = (i4v){0, 0, 0, 0};
#pragma unroll
for (int b = 0; b < 4; ++b) {
const int g = KC ? 0 : b;
#pragma unroll
for (int h = 0; h < 2; ++h) {
const long a = (long)(((uint64_t)k_bytes(b, 2 * h + 1) << 32) |
k_bytes(b, 2 * h));
Sh[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Qh[2 * b + h],
Sh[g], 0, 0, 0);
Sl[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Ql[2 * b + h],
Sl[g], 0, 0, 0);
}
}
auto dot = [&](int g, int e) {
return (float)(Sh[g][e] * 256 + Sl[g][e]);
};
load_v(lt);
float sc[4];
float mloc = -INFINITY;
#pragma unroll
for (int e = 0; e < 4; ++e) {
const int t = base + 4 * gq + e;
const bool ok = t < end && t < qlen;
float v;
if constexpr (KC) {
v = lo16(sa[e]) * dot(0, e);
} else {
v = lo16(sa[e]) * dot(0, e) + hi16(sa[e]) * dot(1, e) +
lo16(sb[e]) * dot(2, e) + hi16(sb[e]) * dot(3, e);
}
sc[e] = ok ? v * sm_q : -INFINITY;
mloc = fmaxf(mloc, sc[e]);
}
const float mn = fmaxf(m, qmax4(mloc));
float alpha = 1.0f;
if (mn > m)
alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn);
m = mn;
h4v pb;
float psum = 0.0f, lsum = 0.0f;
#pragma unroll
for (int e = 0; e < 4; ++e) {
const float p =
sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn);
lsum += p;
const _Float16 hp = (_Float16)(p * lo16(sv[e]) * PSCALE);
psum += (float)hp;
pb[e] = hp;
}
l = l * alpha + lsum;
cp = cp * alpha + psum;
if (__ballot(alpha != 1.0f)) {
#pragma unroll
for (int c = 0; c < 16; ++c) O[c] *= alpha;
}
// A of fragment 4 s + d: dim 16 col + 4 s + d of tokens 4 gq .. + 3.
// One perm per token pair transposes the biased bytes and widens
// them to fp16 [b, 0x00] (selector 0x0C), i.e. the exact denormal
// b * 2^-24.
#pragma unroll
for (int s = 0; s < 4; ++s) {
uint32_t a[4];
#pragma unroll
for (int e = 0; e < 4; ++e) a[e] = v_bytes(vr[e], s);
#pragma unroll
for (int d = 0; d < 4; ++d) {
constexpr uint32_t kSel = 0x0C000C00u;
const uint32_t sel = kSel | ((4u + d) << 16) | (uint32_t)d;
const h4v A = __builtin_bit_cast(
h4v, (uint2){__builtin_amdgcn_perm(a[1], a[0], sel),
__builtin_amdgcn_perm(a[3], a[2], sel)});
O[4 * s + d] = __builtin_amdgcn_mfma_f32_16x16x16f16(
A, pb, O[4 * s + d], 0, 0, 0);
}
}
}
l = qsum4(l);
cp = qsum4(cp);
if (qlive) {
float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo +
(int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms;
const bool any = m != -INFINITY;
// O holds sum hp (x + 128) 2^-24.
const float corr = 128.0f * 0x1p-24f * cp;
constexpr float kOut = 0x1p24f / (PSCALE * LUT_ONE);
#pragma unroll
for (int e = 0; e < 4; ++e) {
float* dp = op + 16 * (4 * gq + e);
#pragma unroll
for (int c = 0; c < 16; c += 4) {
float4 v;
v.x = any ? (O[c][e] - corr) * kOut : 0.0f;
v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f;
v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f;
v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f;
*(float4*)(dp + c) = v;
}
}
if (gq == 0) {
op[D] = m;
op[D + 1] = l;
}
}
}
seg0 = seg1;
}
#endif // OCTAVE_MFMA_STUB
}
// Split combine for gfx942: grid (num_q, num_q_heads), 16 waves. Thread t
// weighs split t; wave w sums splits w, w + 16, ... four at a time over
// dims 4 lane .. + 3; wave 0 adds the 16 partials and undoes the 256-dim V
// rotation in registers (two in-lane stages, six shuffle stages).
template <typename OutT>
__global__ __launch_bounds__(1024) void reduce_mfma(
const float* __restrict__ mid_o, OutT* __restrict__ out,
const int* __restrict__ v_signs, int num_splits, int64_t smo, int64_t smh,
int64_t sms, int64_t soo, int64_t soh) {
#ifndef OCTAVE_MFMA_STUB
__shared__ float sw[1024], swm[16], swl[16];
__shared__ __attribute__((aligned(16))) float so[16][D];
const int qi = blockIdx.x, hi = blockIdx.y;
const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63;
const float* base = mid_o + qi * smo + hi * smh;
const float ms = tid < num_splits ? base[tid * sms + D] : -INFINITY;
const float ls = tid < num_splits ? base[tid * sms + D + 1] : 0.0f;
float m = ms;
#pragma unroll
for (int k = 1; k < 64; k <<= 1) m = fmaxf(m, __shfl_xor(m, k));
if (lane == 0) swm[w] = m;
__syncthreads();
float mg = swm[0];
#pragma unroll
for (int i = 1; i < 16; ++i) mg = fmaxf(mg, swm[i]);
const float a = ms == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(ms - mg);
sw[tid] = a;
float l = ls * a;
#pragma unroll
for (int k = 1; k < 64; k <<= 1) l += __shfl_xor(l, k);
if (lane == 0) swl[w] = l;
__syncthreads();
float4 o = {0, 0, 0, 0};
for (int s0 = w; s0 < num_splits; s0 += 64) {
float4 x[4];
float as[4];
#pragma unroll
for (int u = 0; u < 4; ++u) {
const int s = s0 + 16 * u;
as[u] = s < num_splits ? sw[s] : 0.0f;
x[u] = s < num_splits ? *(const float4*)(base + s * sms + 4 * lane)
: make_float4(0, 0, 0, 0);
}
#pragma unroll
for (int u = 0; u < 4; ++u) {
// Splits without tokens weigh zero; their O was never written.
o.x += as[u] == 0.0f ? 0.0f : x[u].x * as[u];
o.y += as[u] == 0.0f ? 0.0f : x[u].y * as[u];
o.z += as[u] == 0.0f ? 0.0f : x[u].z * as[u];
o.w += as[u] == 0.0f ? 0.0f : x[u].w * as[u];
}
}
*(float4*)&so[w][4 * lane] = o;
__syncthreads();
if (w != 0) return;
float x[4] = {0, 0, 0, 0};
#pragma unroll
for (int i = 0; i < 16; ++i) {
const float4 v = *(const float4*)&so[i][4 * lane];
x[0] += v.x;
x[1] += v.y;
x[2] += v.z;
x[3] += v.w;
}
float lg = 0.0f;
#pragma unroll
for (int i = 0; i < 16; ++i) lg += swl[i];
{
const float a0 = x[0] + x[1], a1 = x[0] - x[1];
const float a2 = x[2] + x[3], a3 = x[2] - x[3];
x[0] = a0 + a2;
x[1] = a1 + a3;
x[2] = a0 - a2;
x[3] = a1 - a3;
}
#pragma unroll
for (int mk = 1; mk < 64; mk <<= 1) {
const bool hi_lane = lane & mk;
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float p = __shfl_xor(x[i], mk);
x[i] = hi_lane ? p - x[i] : x[i] + p;
}
}
const float inv_l = lg > 0.f ? 0.0625f / lg : 0.f;
OutT* op = out + qi * soo + hi * soh + 4 * lane;
#pragma unroll
for (int i = 0; i < 4; ++i)
op[i] = (OutT)(x[i] * sign_of(v_signs, 4 * lane + i) * inv_l);
#endif // OCTAVE_MFMA_STUB
}
} // namespace octave_mfma
// Split combine on gfx942; false when the shape is not covered.
bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o,
torch::Tensor v_signs, int num_splits,
cudaStream_t stream) {
using namespace octave_mfma;
// 1024 threads per (query, head) pay off while those blocks alone cannot
// fill the GPU; big batches keep the 256-thread kernel.
static const int max_blocks =
2 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
if (num_splits > 1024 || out.size(0) * out.size(1) > max_blocks ||
out.size(2) != D || out.stride(2) != 1 || mid_o.stride(3) != 1)
return false;
dim3 grid(out.size(0), out.size(1));
if (out.dtype() == at::kBFloat16)
reduce_mfma<__hip_bfloat16><<<grid, 1024, 0, stream>>>(
mid_o.data_ptr<float>(), (__hip_bfloat16*)out.data_ptr(),
v_signs.data_ptr<int>(), num_splits, mid_o.stride(0), mid_o.stride(1),
mid_o.stride(2), out.stride(0), out.stride(1));
else
reduce_mfma<half><<<grid, 1024, 0, stream>>>(
mid_o.data_ptr<float>(), (half*)out.data_ptr(), v_signs.data_ptr<int>(),
num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2),
out.stride(0), out.stride(1));
return true;
}
constexpr int kMfmaSplits = 256;
constexpr int kMfmaMinTps = 32;
constexpr int kMfmaNw = 4;
// Launches the MFMA decode when the shape is covered; returns the number of
// splits it wrote, or 0 when the caller must use its own kernel. Dense only.
int octave_decode_mfma(torch::Tensor query, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor q_to_req,
torch::Tensor q_to_klen, torch::Tensor mid_o,
torch::Tensor k_signs, double sm_scale,
int64_t num_kv_splits, int64_t fmt,
int64_t query_group) {
using namespace octave_mfma;
const int num_q = query.size(0);
const int num_q_heads = query.size(1);
const int num_kv_heads = cache.size(1);
const int hpk = num_q_heads / num_kv_heads;
const int QG = query_group > 1 ? 4 : 1;
const int slot = fmt == 4 ? octave::Format<4, false>::SLOT
: fmt == 3 ? octave::Format<3, false>::SLOT
: octave::Format<3, true>::SLOT;
const bool eligible =
query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 &&
cache.size(2) % 16 == 0 &&
cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 &&
cache.stride(2) == slot && cache.stride(1) % 16 == 0 &&
cache.stride(0) % 16 == 0 &&
reinterpret_cast<uintptr_t>(cache.data_ptr()) % 16 == 0 &&
num_kv_splits >= 1;
if (!eligible) return 0;
const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1;
// Enough splits for ~two resident waves per SIMD, no more: every extra
// split is another partial for the reduction to read.
static const int target_waves =
8 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
const int groups = std::max(1, (num_q + QG - 1) / QG) * num_kv_heads * nrt;
const int ns =
std::min<int>({kMfmaSplits, (int)num_kv_splits,
std::max(1, (target_waves + groups - 1) / groups)});
if (num_q == 0) return ns;
auto stream = at::cuda::getCurrentCUDAStream().stream();
const int spb = kMfmaNw / nrt;
dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb);
#define SQM_G(B, KC, G, T) \
decode_mfma<B, KC, kMfmaNw, G, T><<<grid, 64 * kMfmaNw, 0, stream>>>( \
(const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \
block_table.data_ptr<int>(), q_to_req.data_ptr<int>(), \
q_to_klen.data_ptr<int>(), mid_o.data_ptr<float>(), \
k_signs.data_ptr<int>(), (float)sm_scale, num_q, num_q_heads, \
num_kv_heads, cache.size(2), block_table.size(1), block_table.size(0), \
cache.size(0), ns, kMfmaMinTps, nrt, query.stride(0), query.stride(1), \
cache.stride(0), cache.stride(1), cache.stride(2), mid_o.stride(0), \
mid_o.stride(1), mid_o.stride(2))
#define SQM(B, KC, T) \
do { \
if (QG == 4) \
SQM_G(B, KC, 4, T); \
else \
SQM_G(B, KC, 1, T); \
} while (0)
const bool bf = query.dtype() == at::kBFloat16;
if (fmt == 4) {
if (bf)
SQM(4, false, __hip_bfloat16);
else
SQM(4, false, half);
} else if (fmt == 3) {
if (bf)
SQM(3, false, __hip_bfloat16);
else
SQM(3, false, half);
} else {
if (bf)
SQM(3, true, __hip_bfloat16);
else
SQM(3, true, half);
}
#undef SQM
#undef SQM_G
return ns;
}
#endif // USE_ROCM
csrc/attention/octave_decode_wmma_rdna3.cu
+687
View File
@@ -0,0 +1,687 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Octave decode on WMMA for RDNA3 (gfx11). Independent waves, one (split,
// 16-query row tile) per wave, transposed orientation, and lane pairs sharing
// A fragments through a DPP swap. The products:
//
// QK S^T[token][q] = K x Q^T over 16 k-steps of 16 dims: steps 0-3 the
// RoPE block (4-bit codes), 4-15 the NoPE blocks (3-bit codes). Codes
// are mapped to their int8 codebook values and fed as x + 1152. Each
// 64-dim block has its own per-token scale, so the steps accumulate
// into four tiles that are scaled and debiased in the softmax. Compact
// K: 16 steps of 3-bit codes and one scale, so one tile.
// PV O^T[dim][q] = V^T x P^T with the V codebook values fed as x + 1152;
// the per-token V scale is folded into P, the bias removed with sum(P).
//
// The query is rotated like K while it is staged into LDS; the output stays
// rotated and octave_reduce rotates it back.
#include <cstdint>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#if defined(USE_ROCM)
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include "octave_format.cuh"
#if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__)
#define OCTAVE_WMMA_STUB
#endif
namespace octave_wmma {
typedef _Float16 h16v __attribute__((ext_vector_type(16)));
typedef _Float16 h8v __attribute__((ext_vector_type(8)));
typedef _Float16 h2v __attribute__((ext_vector_type(2)));
typedef float f8v __attribute__((ext_vector_type(8)));
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef const int __attribute__((address_space(4))) cint;
using octave::D;
using octave::Format;
using octave::LUT_ONE;
using octave::R;
constexpr int QG = 4; // query tokens per block
// Exact recent and leading positions (sparse window mode).
[[maybe_unused]] constexpr int kWindow = 64;
[[maybe_unused]] constexpr int kSinks = 4;
[[maybe_unused]] constexpr int QROW = D + 8; // padded LDS row (bank spread)
[[maybe_unused]] constexpr float PSCALE = 256.0f;
[[maybe_unused]] constexpr float BIAS = 1152.0f; // int8 byte x fed as x + 1152
[[maybe_unused]] constexpr uint32_t X = 0x80808080u;
__device__ __forceinline__ uint32_t swap1(uint32_t v) {
return (uint32_t)__builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xF, 0xF,
false);
}
__device__ __forceinline__ uint32_t xhalf(uint32_t v) {
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false,
false);
}
__device__ __forceinline__ float xhalf(float v) {
return __uint_as_float(xhalf(__float_as_uint(v)));
}
// Bytes b -> fp16 1024 + b, exact.
__device__ __forceinline__ h16v bytes_to_h16(uint32_t w0, uint32_t w1,
uint32_t w2, uint32_t w3) {
const uint32_t ws[4] = {w0, w1, w2, w3};
u8v r;
#pragma unroll
for (int i = 0; i < 4; ++i) {
r[2 * i] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07010700u);
r[2 * i + 1] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07030702u);
}
return __builtin_bit_cast(h16v, r);
}
__device__ __forceinline__ h16v ld16(const _Float16* p) {
const h8v lo = *(const h8v*)p;
const h8v hi = *(const h8v*)(p + 8);
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
12, 13, 14, 15);
}
// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each.
__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
uint32_t a3, uint32_t o[4]) {
const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
}
template <typename T>
__device__ __forceinline__ float to_f(T x) {
return (float)x;
}
template <>
__device__ __forceinline__ float to_f<half>(half x) {
return __half2float(x);
}
__device__ __forceinline__ float sign_of(const int* bits, int i) {
return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f;
}
// SP: sparse mode. Row qi attends to the logical positions idx[qi, 0..topk)
// (negative entries are skipped) instead of its first q_to_klen[qi] tokens,
// and every row is its own segment since rows do not share positions.
template <int VB, bool KC, int NSB, typename QT, bool SP>
__global__ __launch_bounds__(64 * NSB) void decode_wmma(
const QT* __restrict__ Q, const uint8_t* __restrict__ cache,
const int* __restrict__ block_table, const int* __restrict__ q_to_req,
const int* __restrict__ q_to_klen, float* __restrict__ mid_o,
const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads,
int num_kv_heads, int block_size, int max_blocks, int num_reqs,
int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles,
int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct,
int64_t smo, int64_t smh, int64_t sms, const int* __restrict__ idx,
int64_t sidx, int topk, const int* __restrict__ qpos,
const int* __restrict__ wtags, int wmask, const int* __restrict__ stags,
int smask) {
#ifndef OCTAVE_WMMA_STUB
using F = Format<VB, KC>;
constexpr int NG = KC ? 1 : 4; // QK tiles (one per K scale)
const int grp = blockIdx.x;
const int kvh = blockIdx.y;
const int tid = threadIdx.x;
const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31;
// Two row tiles when the group can hold more than 16 query vectors;
// otherwise every wave takes its own split.
const int rt = num_row_tiles == 2 ? (w & 1) : 0;
const int splits_per_block = 2 * NSB / num_row_tiles;
const int si =
blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w);
const int j = lane & 15, h = lane >> 4;
const bool useful = (j & 1) == h;
const int rr = useful ? j : (j ^ 1);
const int hpk = num_q_heads / num_kv_heads;
const int max_kv = max_blocks * block_size;
const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE;
__shared__ __attribute__((aligned(16))) _Float16 sQ[32][QROW];
// Stage the group's queries rotated like K: lane owns 8 dims.
for (int v = w; v < 32; v += 2 * NSB) {
const int r = v / hpk, g = v % hpk;
const int qi = grp * QG + r;
float x[8];
const bool live = v < QG * hpk && qi < num_q;
#pragma unroll
for (int i = 0; i < 8; ++i)
x[i] = live ? to_f<QT>(Q[(int64_t)qi * sq0 +
(int64_t)(kvh * hpk + g) * sq1 + lane * 8 + i])
: 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, lane * 8 + i);
#pragma unroll
for (int hh = 1; hh < 8; hh <<= 1) {
#pragma unroll
for (int i = 0; i < 8; ++i)
if ((i & hh) == 0) {
const float a = x[i], b = x[i + hh];
x[i] = a + b;
x[i + hh] = a - b;
}
}
#pragma unroll
for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) {
const bool hi = lane & m;
#pragma unroll
for (int i = 0; i < 8; ++i) {
const float o = __shfl_xor(x[i], m);
x[i] = hi ? o - x[i] : x[i] + o;
}
}
#pragma unroll
for (int i = 0; i < 8; ++i)
sQ[v][lane * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f));
}
__syncthreads();
int rreq[QG], rlen[QG];
#pragma unroll
for (int r = 0; r < QG; ++r) {
const int qi = grp * QG + r;
int rq = -1, ln = 0;
if (qi < num_q) {
rq = q_to_req[qi];
ln = SP ? topk : q_to_klen[qi];
if (rq < 0 || rq >= num_reqs) {
rq = 0;
ln = 0;
}
ln = min(max(ln, 0), SP ? topk : max_kv);
}
rreq[r] = __builtin_amdgcn_readfirstlane(rq);
rlen[r] = __builtin_amdgcn_readfirstlane(ln);
}
// Per-segment sums of this lane's query (debiasing): half h sums dims
// [128 h, 128 h + 128) as two 64-dim segments; the halves swap.
const int qv = rt * 16 + j;
float bq[4];
{
const h2v one = {(_Float16)1.0f, (_Float16)1.0f};
float s0 = 0.f, s1 = 0.f;
#pragma unroll
for (int s = 0; s < 8; ++s) {
const h8v a = *(const h8v*)&sQ[qv][128 * h + 8 * s];
const h8v b = *(const h8v*)&sQ[qv][128 * h + 64 + 8 * s];
#pragma unroll
for (int i = 0; i < 8; i += 2) {
s0 = __builtin_amdgcn_fdot2((h2v){a[i], a[i + 1]}, one, s0, false);
s1 = __builtin_amdgcn_fdot2((h2v){b[i], b[i + 1]}, one, s1, false);
}
}
const float o0 = xhalf(s0), o1 = xhalf(s1);
bq[0] = (h ? o0 : s0) * BIAS;
bq[1] = (h ? o1 : s1) * BIAS;
bq[2] = (h ? s0 : o0) * BIAS;
bq[3] = (h ? s1 : o1) * BIAS;
if constexpr (KC) bq[0] += bq[1] + bq[2] + bq[3];
}
int seg0 = 0;
while (seg0 < QG && rreq[seg0] >= 0) {
int seg1 = seg0 + 1;
if constexpr (!SP)
while (seg1 < QG && rreq[seg1] == rreq[seg0]) ++seg1;
const int req = rreq[seg0];
int Lmax = 0;
for (int r = seg0; r < seg1; ++r) Lmax = max(Lmax, rlen[r]);
const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk;
int tps = (Lmax + num_splits - 1) / num_splits;
if (tps < min_tps) tps = min_tps;
tps = (tps + 15) & ~15;
const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16;
const bool qlive = qv >= qv_lo && qv < qv_hi;
if (blockIdx.z * splits_per_block * tps >= Lmax) {
if (tile_live && si < num_splits && h == 0 && qlive) {
int qx = qv;
asm volatile("" : "+v"(qx));
float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo +
(int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms;
op[D] = -INFINITY;
op[D + 1] = 0.0f;
}
seg0 = seg1;
continue;
}
if (tile_live && si < num_splits) {
int qlen = 0;
#pragma unroll
for (int r = 0; r < QG; ++r)
if (qlive && qv / hpk == r) qlen = rlen[r];
const int start = si * tps;
const int end = min(start + tps, Lmax);
float m = -INFINITY, l = 0.0f, cp = 0.0f;
f8v O[16];
#pragma unroll
for (int c = 0; c < 16; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
const uint8_t* cache_h = cache + (int64_t)kvh * sch;
auto pickw = [&](uint32_t v, int t) {
const uint32_t ev = __builtin_amdgcn_readlane(v, t);
const uint32_t od = __builtin_amdgcn_readlane(v, t + 1);
return h ? od : ev;
};
auto tile_ptr = [&](int pos) {
const int lb = pos / block_size;
int pb = ((cint*)block_table)[req * max_blocks + lb];
if (pb < 0 || pb >= num_phys_blocks) pb = 0;
return cache_h + (int64_t)pb * scb +
(int64_t)(pos - lb * block_size) * sct;
};
// Sparse mode: slot of the pos-th selected position of this row, and
// whether that position maps to a live cache slot. Dead entries read
// slot 0 of block 0 and are masked out of the softmax.
const int* irow = SP ? idx + (int64_t)(grp * QG + seg0) * sidx : nullptr;
auto sel_block = [&](int pos, int& off) {
const int t = irow[min(pos, topk - 1)];
const int lb = max(t, 0) / block_size;
off = max(t, 0) - lb * block_size;
const int pb = (t >= 0 && lb < max_blocks)
? ((cint*)block_table)[req * max_blocks + lb]
: -1;
return (pb >= 0 && pb < num_phys_blocks) ? pb : -1;
};
// Exact-window tokens (attended by octave_window_attn instead): the
// last kWindow positions before the query, and the first kSinks.
const int rpos = SP && wtags ? qpos[grp * QG + seg0] : 0;
auto in_window = [&](int pos, int pb, int off) {
const int t = irow[min(pos, topk - 1)];
const int d = rpos - t;
const int slot = pb * block_size + off;
return (d >= 0 && d < kWindow && wtags[slot & wmask] == slot) ||
(t >= 0 && t < kSinks && stags[slot & smask] == slot);
};
auto sel_ptr = [&](int pos) {
int off;
const int pb = sel_block(pos, off);
return pb < 0 ? cache_h
: cache_h + (int64_t)pb * scb + (int64_t)off * sct;
};
// Raw K loads, one tile ahead. Each lane of a pair owns 8 of the 16
// k-steps of token rr: NR4 4-bit steps (8 bytes of codes each), then
// 3-bit steps (2-bit plane word, 1-bit plane word pre-shifted). Block
// K: the useful lane's slots are steps {0,1,4..9}, its partner's
// {2,3,10..15}. Compact K: {0..7} and {8..15}. The partner's reach the
// useful lane through the DPP swap.
constexpr int NR4 = F::KR / 32; // 4-bit slots per lane
constexpr int NR3 = 8 - NR4; // 3-bit slots per lane
uint2 kr_raw[NR4 > 0 ? NR4 : 1];
uint2 kn_raw[NR3];
// Scales of token j (lane j), raw fp16 pairs. Block K: {k0, k1},
// {k2, k3}, {v}. Compact K: {k, v}.
uint32_t sc_a = 0, sc_b = 0, sc_v = 0;
auto load_k = [&](const uint8_t* tile, int base, float dep) {
uint32_t off = (uint32_t)(rr * (int)sct);
asm volatile("" : "+v"(off) : "v"(dep));
const uint8_t* pk = SP ? sel_ptr(base + rr) : tile + off;
const int chunk0 = useful ? 0 : NR3;
if constexpr (NR4 > 0) {
const int rope0 = useful ? 0 : NR4;
#pragma unroll
for (int k = 0; k < NR4; ++k)
kr_raw[k] = *(const uint2*)(pk + 8 * (rope0 + k));
}
#pragma unroll
for (int k = 0; k < NR3; ++k) {
const int c = chunk0 + k;
kn_raw[k].x = *(const uint32_t*)(pk + F::OFF_KN + 4 * c);
kn_raw[k].y =
*(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 4 * (c >> 1)) >>
(4 * (c & 1));
}
const uint8_t* ps =
(SP ? sel_ptr(base + j) : tile + (int64_t)j * sct) + F::OFF_SC;
if constexpr (KC) {
sc_a = *(const uint32_t*)ps;
sc_v = sc_a >> 16;
} else {
const uint2 u = *(const uint2*)ps;
sc_a = u.x;
sc_b = u.y;
sc_v = *(const uint16_t*)(ps + 8);
}
};
// Raw V loads: lane rr's 16 dims of the even (useful) or odd tokens.
uint2 vraw[8];
auto load_v = [&](const uint8_t* tile, int base, float dep) {
uint32_t off = (uint32_t)((useful ? 0 : 1) * (int)sct);
asm volatile("" : "+v"(off) : "v"(dep));
const uint8_t* pv = tile + off + F::OFF_V;
#pragma unroll
for (int i = 0; i < 8; ++i) {
const uint8_t* p =
SP ? sel_ptr(base + 2 * i + (useful ? 0 : 1)) + F::OFF_V
: pv + (int64_t)(2 * i) * sct;
if constexpr (VB == 4) {
vraw[i] = *(const uint2*)(p + 8 * rr);
} else {
vraw[i].x = *(const uint32_t*)(p + 4 * rr);
vraw[i].y = *(const uint32_t*)(p + D / 4 + 4 * (rr >> 1));
}
}
};
// Biased codebook bytes (x + 128) of a 16-dim step from its raw words.
auto rope_bytes = [&](uint2 raw, uint32_t o[4]) {
o[0] = octave::lut4(raw.x & 0x0F0F0F0Fu) ^ X;
o[1] = octave::lut4((raw.x >> 4) & 0x0F0F0F0Fu) ^ X;
o[2] = octave::lut4(raw.y & 0x0F0F0F0Fu) ^ X;
o[3] = octave::lut4((raw.y >> 4) & 0x0F0F0F0Fu) ^ X;
};
auto nope_bytes = [&](uint2 raw, uint32_t o[4]) {
#pragma unroll
for (int i = 0; i < 4; ++i)
o[i] = octave::lut3(((raw.x >> (2 * i)) & 0x03030303u) |
(((raw.y >> i) & 0x01010101u) << 2)) ^
X;
};
// V of a tile is loaded after its QK (softmax hides it): holding it
// through the QK, next to O, the four S tiles and the raw K, spilled.
const uint8_t* tcur = nullptr;
if (start < end) {
tcur = SP ? nullptr : tile_ptr(start);
load_k(tcur, start, 0.f);
}
for (int base = start; base < end; base += 16) {
auto lo16 = [](uint32_t x) {
return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF)));
};
auto hi16 = [](uint32_t x) {
return __half2float(__ushort_as_half((unsigned short)(x >> 16)));
};
const uint32_t sa_c = sc_a, sb_c = sc_b, sv_c = sc_v;
// QK. Iteration k feeds the useful lane's slot k and, through the
// swap, the partner's slot k.
f8v S[NG];
#pragma unroll
for (int g = 0; g < NG; ++g) S[g] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
int qoff = qv * QROW;
asm volatile("" : "+v"(qoff));
constexpr int kStepOwn[8] = {0, 1, 4, 5, 6, 7, 8, 9};
constexpr int kStepPar[8] = {2, 3, 10, 11, 12, 13, 14, 15};
#pragma unroll
for (int k = 0; k < 8; ++k) {
uint32_t wd[4];
if (k < NR4)
rope_bytes(kr_raw[k < NR4 ? k : 0], wd);
else
nope_bytes(kn_raw[k - NR4], wd);
const int so = KC ? k : kStepOwn[k], sp = KC ? 8 + k : kStepPar[k];
const int go = KC ? 0 : so / 4, gp = KC ? 0 : sp / 4;
// One A and one B fragment live at a time. Each read waits for the
// previous WMMA: without the dependence the compiler hoists all Q
// fragments and spills.
asm volatile("" : "+v"(qoff) : "v"(S[gp][0]));
{
const h16v b = ld16(&sQ[0][0] + qoff + 16 * so);
const h16v a = bytes_to_h16(wd[0], wd[1], wd[2], wd[3]);
S[go] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[go]);
}
__builtin_amdgcn_sched_barrier(0);
asm volatile("" : "+v"(qoff) : "v"(S[go][0]));
{
const h16v b = ld16(&sQ[0][0] + qoff + 16 * sp);
const h16v a = bytes_to_h16(swap1(wd[0]), swap1(wd[1]),
swap1(wd[2]), swap1(wd[3]));
S[gp] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[gp]);
}
__builtin_amdgcn_sched_barrier(0);
}
load_v(tcur, base, S[0][0]);
// Next tile's K and scales now: softmax and PV hide them.
const bool next = base + 16 < end;
const uint8_t* tn = next && !SP ? tile_ptr(base + 16) : nullptr;
if (next) load_k(tn, base + 16, S[0][0]);
float sc[8];
float mloc = -INFINITY;
#pragma unroll
for (int e = 0; e < 8; ++e) {
const int t = base + 2 * e + h;
bool ok = t < end && t < qlen;
if constexpr (SP) {
int off;
const int pb = sel_block(t, off);
ok = ok && pb >= 0 && !(wtags && in_window(t, pb, off));
}
const uint32_t a = pickw(sa_c, 2 * e);
float v;
if constexpr (KC) {
v = lo16(a) * (S[0][e] - bq[0]);
} else {
const uint32_t b = pickw(sb_c, 2 * e);
v = lo16(a) * (S[0][e] - bq[0]) + hi16(a) * (S[1][e] - bq[1]) +
lo16(b) * (S[2][e] - bq[2]) + hi16(b) * (S[3][e] - bq[3]);
}
sc[e] = ok ? v * sm_log2 : -INFINITY;
mloc = fmaxf(mloc, sc[e]);
}
const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc)));
float alpha = 1.0f;
if (mn > m)
alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn);
m = mn;
uint32_t pd[4];
float psum = 0.0f, lsum = 0.0f;
#pragma unroll
for (int e = 0; e < 8; e += 2) {
const float v0 = lo16(pickw(sv_c, 2 * e));
const float v1 = lo16(pickw(sv_c, 2 * e + 2));
const float p0 =
sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn);
const float p1 = sc[e + 1] == -INFINITY
? 0.0f
: __builtin_amdgcn_exp2f(sc[e + 1] - mn);
lsum += p0 + p1;
const _Float16 h0 = (_Float16)(p0 * v0 * PSCALE);
const _Float16 h1 = (_Float16)(p1 * v1 * PSCALE);
psum += (float)h0 + (float)h1;
pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1});
}
l = l * alpha + lsum;
cp = cp * alpha + psum;
if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) {
#pragma unroll
for (int c = 0; c < 16; ++c) O[c] *= alpha;
}
u8v pf;
#pragma unroll
for (int i = 0; i < 4; ++i) {
const uint32_t o = xhalf(pd[i]);
pf[i] = h ? o : pd[i];
pf[4 + i] = h ? pd[i] : o;
}
const h16v pb16 = __builtin_bit_cast(h16v, pf);
// Biased codebook bytes of dims 4q..4q+3 of this lane's chunk,
// token i.
auto vq = [&](int i, int q) -> uint32_t {
if constexpr (VB == 4) {
const uint32_t w = q < 2 ? vraw[i].x : vraw[i].y;
return octave::lut4((w >> (4 * (q & 1))) & 0x0F0F0F0Fu) ^ X;
} else {
const uint32_t hi = vraw[i].y >> (4 * (rr & 1));
return octave::lut3(((vraw[i].x >> (2 * q)) & 0x03030303u) |
(((hi >> q) & 0x01010101u) << 2)) ^
X;
}
};
#pragma unroll
for (int q = 0; q < 4; ++q) {
uint32_t t0[4], t1[4];
tr4(vq(0, q), vq(1, q), vq(2, q), vq(3, q), t0);
tr4(vq(4, q), vq(5, q), vq(6, q), vq(7, q), t1);
#pragma unroll
for (int b = 0; b < 4; ++b) {
const int c = 4 * q + b;
uint32_t o0 = t0[b], o1 = t1[b];
if (c >= 2)
asm volatile("" : "+v"(o0), "+v"(o1) : "v"(O[c - 2][0]));
const h16v a = bytes_to_h16(o0, o1, swap1(o0), swap1(o1));
O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]);
}
}
tcur = tn;
}
l += xhalf(l);
cp += xhalf(cp);
if (qlive) {
int qx = qv;
asm volatile("" : "+v"(qx));
float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo +
(int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms;
const bool any = m != -INFINITY;
const float corr = BIAS * cp;
constexpr float kOut = 1.0f / (PSCALE * LUT_ONE);
#pragma unroll
for (int e = 0; e < 8; ++e) {
float* dp = op + 16 * (2 * e + h);
#pragma unroll
for (int c = 0; c < 16; c += 4) {
float4 v;
v.x = any ? (O[c][e] - corr) * kOut : 0.0f;
v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f;
v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f;
v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f;
*(float4*)(dp + c) = v;
}
}
if (h == 0) {
op[D] = m;
op[D + 1] = l;
}
}
}
seg0 = seg1;
}
#endif // OCTAVE_WMMA_STUB
}
} // namespace octave_wmma
// Splits per row group and minimum tokens per split, as in the int8 kernel.
constexpr int kSqWmmaSplits = 192;
constexpr int kSqWmmaMinTps = 64;
constexpr int kSqWmmaNsb = 2;
int octave_decode_mfma(torch::Tensor query, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor q_to_req,
torch::Tensor q_to_klen, torch::Tensor mid_o,
torch::Tensor k_signs, double sm_scale,
int64_t num_kv_splits, int64_t fmt, int64_t query_group);
// Launches the WMMA decode when the shape is covered; returns the number of
// splits it wrote (for the reduce), or 0 when the caller must use its own
// kernel.
// With `indices` ([num_q, topk] int32 logical positions) the sparse variant
// runs and q_to_klen is ignored.
int octave_decode_wmma(torch::Tensor query, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor q_to_req,
torch::Tensor q_to_klen, torch::Tensor mid_o,
torch::Tensor k_signs, double sm_scale,
int64_t num_kv_splits, int64_t fmt,
const std::optional<torch::Tensor>& indices,
const int* qpos, const int* wtags, int wmask,
const int* stags, int smask, int64_t query_group) {
using namespace octave_wmma;
// 0: gfx1100 (this file), 1: gfx942 (octave_decode_mfma_cdna.cu, dense
// only).
static const int arch = [] {
const std::string name =
at::cuda::getCurrentDeviceProperties()->gcnArchName;
if (name.rfind("gfx1100", 0) == 0) return 0;
if (name.rfind("gfx942", 0) == 0) return 1;
return -1;
}();
if (arch == 1)
return indices.has_value()
? 0
: octave_decode_mfma(query, cache, block_table, q_to_req,
q_to_klen, mid_o, k_signs, sm_scale,
num_kv_splits, fmt, query_group);
const bool arch_ok = arch == 0;
const int num_q = query.size(0);
const int num_q_heads = query.size(1);
const int num_kv_heads = cache.size(1);
const int hpk = num_q_heads / num_kv_heads;
const bool eligible =
arch_ok && query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 &&
cache.size(2) % 16 == 0 &&
cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 &&
num_kv_splits >= 1;
if (!eligible) return 0;
const int ns = std::min<int>(kSqWmmaSplits, (int)num_kv_splits);
if (num_q == 0) return ns;
auto stream = at::cuda::getCurrentCUDAStream().stream();
const bool sp = indices.has_value();
const int* idx = sp ? indices->data_ptr<int>() : nullptr;
const int64_t sidx = sp ? indices->stride(0) : 0;
const int topk = sp ? (int)indices->size(1) : 0;
const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1;
const int spb = 2 * kSqWmmaNsb / nrt;
dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb);
#define SQW_SP(B, KC, T, SP) \
decode_wmma<B, KC, kSqWmmaNsb, T, SP> \
<<<grid, 64 * kSqWmmaNsb, 0, stream>>>( \
(const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \
block_table.data_ptr<int>(), q_to_req.data_ptr<int>(), \
q_to_klen.data_ptr<int>(), mid_o.data_ptr<float>(), \
k_signs.data_ptr<int>(), (float)sm_scale, num_q, num_q_heads, \
num_kv_heads, cache.size(2), block_table.size(1), \
block_table.size(0), cache.size(0), ns, kSqWmmaMinTps, nrt, \
query.stride(0), query.stride(1), cache.stride(0), \
cache.stride(1), cache.stride(2), mid_o.stride(0), \
mid_o.stride(1), mid_o.stride(2), idx, sidx, topk, qpos, wtags, \
wmask, stags, smask)
#define SQW(B, KC, T) \
do { \
if (sp) \
SQW_SP(B, KC, T, true); \
else \
SQW_SP(B, KC, T, false); \
} while (0)
const bool bf = query.dtype() == at::kBFloat16;
if (fmt == 4) {
if (bf)
SQW(4, false, __hip_bfloat16);
else
SQW(4, false, half);
} else if (fmt == 3) {
if (bf)
SQW(3, false, __hip_bfloat16);
else
SQW(3, false, half);
} else {
if (bf)
SQW(3, true, __hip_bfloat16);
else
SQW(3, true, half);
}
#undef SQW
#undef SQW_SP
return ns;
}
#endif // USE_ROCM
csrc/attention/octave_format.cuh
+78
View File
@@ -0,0 +1,78 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Byte layout of one Octave slot (one token, one KV head). The PyTorch
// reference in vllm/v1/attention/ops/rocm_octave.py is the source of truth.
//
// [0, OFF_KN) K RoPE block (dims 0-63), rotated, 4-bit codes
// [OFF_KN, OFF_V) K NoPE blocks (dims 64-255), rotated, 3-bit codes
// [OFF_V, OFF_SC) V, rotated, VBITS-bit codes
// [OFF_SC, ...) fp16 scales: K block 0..3, V
//
// Compact K (KC): K rotated over all 256 dims like V, 3-bit codes for every
// dim (OFF_KN = 0), one K scale; slots pad to 4 bytes instead of 8.
//
// 4-bit codes: dims 8i..8i+7 in one dword, low nibbles = dims 0-3, high
// nibbles = dims 4-7. 3-bit codes: a 2-bit plane (16 dims per dword, dims
// 4k..4k+3 at bit 2k of bytes 0-3) followed by a 1-bit plane (32 dims per
// dword, dims 4k..4k+3 at bit k of bytes 0-3). A code c of a block with
// scale s stands for LUT[c] * s / LUT_ONE, LUT being the int8 Lloyd-Max
// codebook of its width.
#pragma once
#include <cstdint>
namespace octave {
constexpr int D = 256; // head size
constexpr int R = 64; // RoPE dims (K block 0)
constexpr int N = D - R; // NoPE dims (K blocks 1-3)
constexpr int NKB = D / 64; // K Hadamard blocks
constexpr float LUT_ONE = 64.0f;
template <int VBITS, bool KC = false>
struct Format {
static constexpr int KR = KC ? 0 : R; // K dims with 4-bit codes
static constexpr int KN = D - KR; // K dims with 3-bit codes
static constexpr int OFF_KN = KR * 4 / 8;
static constexpr int OFF_V = OFF_KN + KN * 3 / 8;
static constexpr int OFF_SC = OFF_V + D * VBITS / 8;
static constexpr int NUM_KSC = KC ? 1 : NKB;
static constexpr int NUM_SC = NUM_KSC + 1;
static constexpr int ALIGN = KC ? 4 : 8;
static constexpr int SLOT = (OFF_SC + 2 * NUM_SC + ALIGN - 1) / ALIGN * ALIGN;
};
// Format id passed by the host: V bits, plus 16 for compact K.
constexpr int kCompactFlag = 16;
// int8 Lloyd-Max codebooks as perm tables: byte c of the table is LUT[c].
// 3-bit: -127 -79 -45 -14 14 45 79 127
// 4-bit: -126 -95 -74 -58 -43 -30 -18 -6 6 18 30 43 58 74 95 126
constexpr uint32_t LUT3_LO = 0xF2D3B181u, LUT3_HI = 0x7F4F2D0Eu;
constexpr uint32_t LUT4_0 = 0xC6B6A182u, LUT4_1 = 0xFAEEE2D5u;
constexpr uint32_t LUT4_2 = 0x2B1E1206u, LUT4_3 = 0x7E5F4A3Au;
// Four 3-bit codes (one per byte) -> their four int8 codebook values.
__device__ __forceinline__ uint32_t lut3(uint32_t codes) {
return __builtin_amdgcn_perm(LUT3_HI, LUT3_LO, codes);
}
// Four 4-bit codes (one per byte) -> their four int8 codebook values.
__device__ __forceinline__ uint32_t lut4(uint32_t codes) {
const uint32_t lo =
__builtin_amdgcn_perm(LUT4_1, LUT4_0, codes & 0x07070707u);
const uint32_t hi =
__builtin_amdgcn_perm(LUT4_3, LUT4_2, codes & 0x07070707u);
const uint32_t m = ((codes >> 3) & 0x01010101u) * 0xFFu;
return (hi & m) | (lo & ~m);
}
template <int BITS>
__device__ __forceinline__ uint32_t lut(uint32_t codes) {
if constexpr (BITS == 4)
return lut4(codes);
else
return lut3(codes);
}
} // namespace octave
csrc/attention/octave_prefill_cdna.cu
+666
View File
@@ -0,0 +1,666 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Octave prefill for CDNA3 (gfx942), head size 256, fp16. Same contract as
// octave_prefill_rdna3.cu: reads the packed cache directly, everything in the
// rotated space, phase 1 writes (O, m, l) partials over the cached prefix and
// phase 2 runs causal attention over the chunk and merges.
//
// Both phases compute transposed tiles so that every MFMA result is already in
// the layout the next MFMA consumes (16x16 MFMA, wave64: lane l holds column
// l % 16 and rows 4 (l / 16) .. 4 (l / 16) + 3):
// S^T[token][q] = K x Q^T (phase 1: int8 16x16x32, phase 2: fp16)
// O^T[dim][q] += V^T x P^T (fp16 16x16x16)
// so P^T comes straight out of the S^T accumulator, one query per lane.
//
// Phase 1: one block = 128 query vectors (token, head) of one KV head, 8 waves
// of 16 query vectors x 256 dims. 256 threads unpack the next 16-token tile
// into LDS while the current one is consumed.
#include <cstdint>
#include <algorithm>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#if defined(USE_ROCM)
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include "octave_format.cuh"
#if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__)
#define OCTAVE_CDNA_STUB
#endif
namespace octave_cdna {
using octave::D;
using octave::Format;
using octave::LUT_ONE;
typedef _Float16 h4v __attribute__((ext_vector_type(4)));
typedef _Float16 h2v __attribute__((ext_vector_type(2)));
typedef float f4v __attribute__((ext_vector_type(4)));
typedef int i4v __attribute__((ext_vector_type(4)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
constexpr int KT = 16; // tokens per tile
[[maybe_unused]] constexpr int FR = D / 16; // 16-dim output fragments
// Both phases run 8 waves (128 query vectors of one KV head) per block: every
// block reads the whole K/V it attends to, so more query vectors per block
// means less of that work per query.
constexpr int NW1 = 8;
constexpr int NT1 = NW1 * 64;
constexpr int QB1 = NW1 * 16;
constexpr int KT1 = 32; // phase-1 tile: two 16-token sub-tiles
[[maybe_unused]] constexpr int VST1 = KT1 + 4;
[[maybe_unused]] constexpr int KROW = D + 16; // bytes per int8 K row
[[maybe_unused]] constexpr int VST = KT + 4; // halves per V^T row
[[maybe_unused]] constexpr int HROW = D + 8; // halves per fp16 K row
// Phase 2 LDS rows, picked for conflict-free banks: a K row is 2 dwords
// mod 32, and V^T rows of 40 halves with the 4-token group XORed by
// (row >> 3) & 7 spread both the transposed stores and the MFMA reads.
[[maybe_unused]] constexpr int KROW2 = D + 4;
[[maybe_unused]] constexpr int VST2 = KT1 + 8;
[[maybe_unused]] constexpr float PSCALE = 4096.0f;
constexpr int kMinSplitTokens = 2048;
__device__ __forceinline__ int used_splits(int ctx, int num_splits) {
if (ctx <= 0) return 0;
return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens);
}
__device__ __forceinline__ int split_len(int ctx, int used) {
return ((ctx + used - 1) / used + 31) / 32 * 32;
}
// Reductions over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48).
__device__ __forceinline__ float qmax4(float v) {
v = fmaxf(v, __shfl_xor(v, 16));
return fmaxf(v, __shfl_xor(v, 32));
}
__device__ __forceinline__ float qsum4(float v) {
v += __shfl_xor(v, 16);
return v + __shfl_xor(v, 32);
}
// Codes of dims 4k..4k+3 of a 64-dim block (see octave_prefill_rdna3.cu).
template <int BITS>
__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw,
uint32_t hi, int k) {
if constexpr (BITS == 4) {
return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu;
} else {
return ((raw[0] >> (2 * (k & 3))) & 0x03030303u) |
(((hi >> (k & 7)) & 0x01010101u) << 2);
}
}
// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens).
__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
uint32_t a3, uint32_t o[4]) {
const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
}
// ===========================================================================
// Phase 1: prefix
// ===========================================================================
template <int VB, bool KC>
__global__ __launch_bounds__(NT1) void prefix_attn(
const __half* __restrict__ q, const uint8_t* __restrict__ cache,
const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q,
const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale,
int num_q_heads, int num_kv_heads, int block_size, int max_blocks,
int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head,
int64_t scb, int64_t sch, int64_t sct) {
#ifndef OCTAVE_CDNA_STUB
using F = Format<VB, KC>;
const int seq = blockIdx.x, kvh = blockIdx.y;
const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits;
const int tid = threadIdx.x;
const int w = tid >> 6, lane = tid & 63;
const int col = lane & 15, gq = lane >> 4;
const int hpk = num_q_heads / num_kv_heads;
const int q0 = cu_seqlens_q[seq];
const int qlen = cu_seqlens_q[seq + 1] - q0;
int ctx = seq_lens[seq] - qlen;
ctx = max(0, min(ctx, max_blocks * block_size));
const int nqv = qlen * hpk;
if (rowtile * QB1 >= nqv) return;
const int used = used_splits(ctx, num_splits);
if (split >= used) return;
const int k_len = split_len(ctx, used);
const int k_begin = min(ctx, split * k_len);
const int k_stop = min(ctx, k_begin + k_len);
const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT1 - 1) / KT1 : 0;
__shared__ __attribute__((aligned(16))) int8_t sK[2][KT1][KROW];
__shared__ __attribute__((aligned(16))) _Float16 sV[2][D * VST1];
__shared__ float sKs[2][4][KT1], sVs[2][KT1];
// ---- loader mapping. Thread half sub (256 threads each) loads the 16-token
// sub-tile sub of a tile; within it, K: dims 16 kp .. + 15 of block kq of
// token tk; V: dims vd + 4 vh .. + 3 of token quad vq.
const int sub = tid >> 8, lt = tid & 255;
const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3;
const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8;
const bool k3 = KC || kq > 0;
const int b3 = KC ? kq : kq - 1;
const int koff0 = k3 ? F::OFF_KN + 16 * b3 + 4 * kp : 8 * kp;
const int koff1 =
k3 ? F::OFF_KN + F::KN / 4 + 8 * b3 + 4 * (kp >> 1) : 8 * kp + 4;
constexpr int VW = VB == 4 ? 1 : 2;
const int tlast = k_stop - 1;
const uint8_t* const head = cache + (int64_t)kvh * sch;
const int* const bt = block_table + seq * max_blocks;
struct Stg {
uint32_t k0, k1;
uint32_t v[4][VW];
uint32_t s0, s1;
};
// Sub-tiles start 16-aligned and sit in one cache block (block sizes are
// multiples of 16), so all 16 slots exist: tokens past k_stop are read as
// they are and masked (K to -inf, P to 0 by selection, never by product).
// Slot offsets inside a sub-tile are fixed per thread.
auto sub_start = [&](int i) {
return __builtin_amdgcn_readfirstlane(
min(k_begin + i * KT1 + 16 * sub, tlast & ~15));
};
const uint32_t ok0 = tk * (uint32_t)sct + koff0;
const uint32_t ok1 = tk * (uint32_t)sct + koff1;
uint32_t ov[4];
#pragma unroll
for (int r = 0; r < 4; ++r) {
ov[r] = (4 * vq + r) * (uint32_t)sct + F::OFF_V +
(VB == 4 ? vd / 2 : 4 * (vd >> 4));
}
const uint32_t os =
(lt & 15) * (uint32_t)sct + F::OFF_SC + (!KC && lt >= 16 ? 8 : 0);
auto fetch = [&](int i, Stg& S) {
const int start = sub_start(i);
const int pb = __builtin_amdgcn_readfirstlane(bt[start / block_size]);
const uint8_t* base = head + (int64_t)(pb < 0 ? 0 : pb) * scb +
(int64_t)(start % block_size) * sct;
S.k0 = *(const uint32_t*)(base + ok0);
S.k1 = *(const uint32_t*)(base + ok1);
#pragma unroll
for (int r = 0; r < 4; ++r) {
S.v[r][0] = *(const uint32_t*)(base + ov[r]);
if constexpr (VB != 4)
S.v[r][1] = *(const uint32_t*)(base + ov[r] - 4 * (vd >> 4) + D / 4 +
4 * (vd >> 5));
}
if (lt < 32) {
if constexpr (KC) {
S.s0 = *(const uint32_t*)(base + os);
} else {
const u2v x = *(const u2v*)(base + os);
S.s0 = x[0];
S.s1 = x[1];
}
}
};
auto store = [&](const Stg& S, int bf) {
const int t0 = 16 * sub;
{
const uint32_t r4[2] = {S.k0, S.k1};
const uint32_t r3[1] = {S.k0};
const uint32_t hi = S.k1 >> (4 * (kp & 1));
u4v o;
#pragma unroll
for (int x = 0; x < 4; ++x)
o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, x))
: octave::lut4(block_codes<4>(r4, 0, x));
*(u4v*)&sK[bf][t0 + tk][64 * kq + 16 * kp] = o;
}
uint32_t c[4];
#pragma unroll
for (int r = 0; r < 4; ++r) {
if constexpr (VB == 4) {
c[r] = (S.v[r][0] >> (4 * vh)) & 0x0F0F0F0Fu;
} else {
const int sl = 2 * ((vd & 15) >> 2) + 2 * vh;
const int sh = ((vd & 31) >> 2) + vh;
c[r] = ((S.v[r][0] >> sl) & 0x03030303u) |
(((S.v[r][1] >> sh) & 0x01010101u) << 2);
}
c[r] = octave::lut<VB>(c[r]) ^ 0x80808080u;
}
uint32_t t[4];
tr4(c[0], c[1], c[2], c[3], t);
// Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x.
const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f};
#pragma unroll
for (int d = 0; d < 4; ++d) {
const int row = vd + 4 * vh + d;
const h2v p01 = __builtin_bit_cast(
h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u));
const h2v p23 = __builtin_bit_cast(
h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u));
*(u2v*)&sV[bf][row * VST1 + t0 + 4 * vq] =
(u2v){__builtin_bit_cast(uint32_t, p01 - kofs),
__builtin_bit_cast(uint32_t, p23 - kofs)};
}
if (lt < 16) {
const h2v a = __builtin_bit_cast(h2v, S.s0);
sKs[bf][0][t0 + lt] = (float)a[0];
if constexpr (!KC) {
const h2v cc = __builtin_bit_cast(h2v, S.s1);
sKs[bf][1][t0 + lt] = (float)a[1];
sKs[bf][2][t0 + lt] = (float)cc[0];
sKs[bf][3][t0 + lt] = (float)cc[1];
}
} else if (lt < 32) {
sVs[bf][t0 + lt - 16] =
(float)__builtin_bit_cast(_Float16,
(uint16_t)(KC ? S.s0 >> 16 : S.s0)) *
(PSCALE / LUT_ONE);
}
};
// ---- this lane's query vector, int8 over dims 32 s + 8 gq .. + 7
const int qv = rowtile * QB1 + w * 16 + col;
const bool qlive = qv < nqv;
long Q[8];
float qsc;
{
const int qtok = qlive ? qv / hpk : 0;
const int qhead = kvh * hpk + (qlive ? qv % hpk : 0);
const __half* qrow =
q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head;
_Float16 x[8][8];
float qmax = 1e-8f;
#pragma unroll
for (int s = 0; s < 8; ++s) {
const u4v r = qlive ? *(const u4v*)(qrow + 32 * s + 8 * gq) : (u4v){};
__builtin_memcpy(x[s], &r, 16);
#pragma unroll
for (int e = 0; e < 8; ++e) qmax = fmaxf(qmax, fabsf((float)x[s][e]));
}
qmax = qmax4(qmax);
const float qinv = 127.0f / qmax;
qsc = qmax * (1.0f / 127.0f);
#pragma unroll
for (int s = 0; s < 8; ++s) {
uint64_t word = 0;
#pragma unroll
for (int e = 0; e < 8; ++e) {
int v = (int)lrintf((float)x[s][e] * qinv);
v = max(-127, min(127, v));
word |= ((uint64_t)(uint8_t)(int8_t)v) << (8 * e);
}
Q[s] = (long)word;
}
}
f4v O[FR];
#pragma unroll
for (int c = 0; c < FR; ++c) O[c] = (f4v){0, 0, 0, 0};
const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE);
float m = -INFINITY, l = 0.f;
// Pipeline: iteration i computes tile i from LDS while the loads of tile
// i + 1 are in flight, then stores them into the other buffer.
Stg S;
if (ntiles > 0) {
fetch(0, S);
store(S, 0);
}
__syncthreads();
for (int i = 0; i < ntiles; ++i) {
const int bf = i & 1;
if (i + 1 < ntiles) fetch(i + 1, S);
// S^T = K x Q^T over both sub-tiles first (independent MFMA chains),
// then one online-softmax step over the 32 tokens, then PV.
f4v Sf[2];
#pragma unroll
for (int j = 0; j < 2; ++j) {
Sf[j] = (f4v){0, 0, 0, 0};
const int8_t* krow = &sK[bf][16 * j + col][8 * gq];
const int tq = 16 * j + 4 * gq;
if constexpr (KC) {
i4v Si = {0, 0, 0, 0};
#pragma unroll
for (int s = 0; s < 8; ++s)
Si = __builtin_amdgcn_mfma_i32_16x16x32_i8(
*(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0);
#pragma unroll
for (int e = 0; e < 4; ++e)
Sf[j][e] = sKs[bf][0][tq + e] * (float)Si[e];
} else {
#pragma unroll
for (int b = 0; b < 4; ++b) {
i4v Si = {0, 0, 0, 0};
#pragma unroll
for (int s = 2 * b; s < 2 * b + 2; ++s)
Si = __builtin_amdgcn_mfma_i32_16x16x32_i8(
*(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0);
#pragma unroll
for (int e = 0; e < 4; ++e)
Sf[j][e] += sKs[bf][b][tq + e] * (float)Si[e];
}
}
}
// Online softmax in log2 units; this lane holds tokens 16 j + 4 gq + e.
float sc[2][4];
float mloc = -INFINITY;
#pragma unroll
for (int j = 0; j < 2; ++j)
#pragma unroll
for (int e = 0; e < 4; ++e) {
const int t = k_begin + i * KT1 + 16 * j + 4 * gq + e;
sc[j][e] = t < k_stop ? Sf[j][e] * c1 : -INFINITY;
mloc = fmaxf(mloc, sc[j][e]);
}
const float mn = fmaxf(m, qmax4(mloc));
const float alpha = __builtin_amdgcn_exp2f(m - mn);
m = mn;
h4v pb[2];
float lsum = 0.f;
#pragma unroll
for (int j = 0; j < 2; ++j)
#pragma unroll
for (int e = 0; e < 4; ++e) {
const bool live = sc[j][e] != -INFINITY;
const float p = live ? __builtin_amdgcn_exp2f(sc[j][e] - mn) : 0.f;
lsum += p;
pb[j][e] =
live ? (_Float16)(p * sVs[bf][16 * j + 4 * gq + e]) : (_Float16)0.f;
}
l = l * alpha + lsum;
if (__ballot(alpha != 1.0f)) {
#pragma unroll
for (int c = 0; c < FR; ++c) O[c] *= alpha;
}
// O^T += V^T x P^T: A = V^T[dim 16 c + col][tokens 16 j + 4 gq .. + 3].
#pragma unroll
for (int j = 0; j < 2; ++j) {
const _Float16* vrow = &sV[bf][col * VST1 + 16 * j + 4 * gq];
#pragma unroll
for (int c = 0; c < FR; ++c)
O[c] = __builtin_amdgcn_mfma_f32_16x16x16f16(
*(const h4v*)(vrow + 16 * c * VST1), pb[j], O[c], 0, 0, 0);
}
if (i + 1 < ntiles) store(S, bf ^ 1);
__syncthreads();
}
// ---- partials for the phase-2 merge (natural-log max, unnormalized O)
l = qsum4(l);
if (!qlive) return;
const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk;
float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) +
((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2);
#pragma unroll
for (int c = 0; c < FR; ++c)
#pragma unroll
for (int e = 0; e < 4; ++e)
wp[16 * c + 4 * gq + e] = O[c][e] * (1.0f / PSCALE);
if (gq == 0) {
wp[D] = m * 0.6931471805599453f;
wp[D + 1] = l;
}
#endif // OCTAVE_CDNA_STUB
}
// ===========================================================================
// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix
// partials. Like phase 1, one block = 128 query vectors (token, head) of one
// KV head, so the GQA heads share every K/V tile; 8 waves of 16 vectors.
// ===========================================================================
__global__ void __launch_bounds__(NT1)
chunk_attn(__half* __restrict__ out, const float* __restrict__ ws,
const __half* __restrict__ q, const __half* __restrict__ k,
const __half* __restrict__ v,
const int* __restrict__ cu_seqlens_q,
const int* __restrict__ seq_lens, int num_q_heads,
int num_kv_heads, int max_ctx, float sm_scale, int num_splits,
int total_q_tokens, int64_t sq_tok, int64_t sq_head,
int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head,
int64_t so_tok, int64_t so_head) {
#ifndef OCTAVE_CDNA_STUB
constexpr float LOG2E = 1.4426950408889634f;
const int seq = blockIdx.x, kvh = blockIdx.y, rowtile = blockIdx.z;
const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63;
const int col = lane & 15, gq = lane >> 4;
const int hpk = num_q_heads / num_kv_heads;
const int q0 = cu_seqlens_q[seq];
const int qlen = cu_seqlens_q[seq + 1] - q0;
const int nqv = qlen * hpk;
if (rowtile * QB1 >= nqv) return;
const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx));
const int qv = rowtile * QB1 + w * 16 + col;
const bool qlive = qv < nqv;
const int qtok = qlive ? qv / hpk : 0;
const int head = kvh * hpk + (qlive ? qv % hpk : 0);
// Keys this block and this wave need (causal: up to their last token).
const int k_end = min(qlen, (min(nqv, (rowtile + 1) * QB1) - 1) / hpk + 1);
const int wave_first = rowtile * QB1 + w * 16;
const bool wave_live = wave_first < nqv;
const int wave_last_tok = min(nqv - 1, wave_first + 15) / hpk;
// One 32-token tile in LDS (two would not fit in 64 KB with fp16 K and V);
// the next tile waits in registers.
__shared__ __attribute__((aligned(16))) _Float16 sK[KT1 * KROW2];
__shared__ __attribute__((aligned(16))) _Float16 sV[D * VST2];
// Merge the prefix splits (log2 units from here on). O^T fragment c holds
// dims 16 c + 4 gq + e of this lane's vector; l stays partial per lane.
f4v acc[FR];
#pragma unroll
for (int c = 0; c < FR; ++c) acc[c] = (f4v){0, 0, 0, 0};
float m = -INFINITY, l = 0.f;
const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2);
const int used = used_splits(ctx, num_splits);
if (qlive && used > 0) {
const float* w0 =
ws + ((int64_t)(q0 + qtok) * num_q_heads + head) * (D + 2);
for (int s = 0; s < used; ++s)
m = fmaxf(m, w0[s * split_stride + D] * LOG2E);
for (int s = 0; s < used; ++s) {
const float* wp = w0 + s * split_stride;
if (wp[D] == -INFINITY) continue;
const float wgt = __builtin_amdgcn_exp2f(wp[D] * LOG2E - m);
l += wp[D + 1] * wgt;
#pragma unroll
for (int c = 0; c < FR; ++c)
#pragma unroll
for (int e = 0; e < 4; ++e) acc[c][e] += wp[16 * c + 4 * gq + e] * wgt;
}
if (gq != 0) l = 0.f;
}
// Q^T operand: dims 16 s + 4 gq .. + 3 of this lane's vector.
h4v Qh[FR];
{
const __half* qrow =
q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)head * sq_head;
#pragma unroll
for (int s = 0; s < FR; ++s) Qh[s] = *(const h4v*)(qrow + 16 * s + 4 * gq);
}
const float c2 = sm_scale * LOG2E;
// Loaders. K: thread t moves dims 8 (t % 32) .. + 7 of tokens t / 32
// and t / 32 + 16 of the 32-token tile. V: dims 4 (t % 64) .. + 3 of
// tokens 4 (t / 64) .. + 3, transposed in registers into V^T rows. Rows
// past the chunk are zeros.
const int lt_tok = tid >> 5, lt_d = (tid & 31) * 8;
const int vt_d = (tid & 63) * 4, vt_g = tid >> 6;
const __half* kbase = k + (int64_t)kvh * sk_head + lt_d;
const __half* vbase = v + (int64_t)kvh * sv_head + vt_d;
u4v Kr[2];
u2v Vr[4];
auto fetch = [&](int n0) {
#pragma unroll
for (int h = 0; h < 2; ++h) {
const int t = n0 + lt_tok + 16 * h;
const bool ok = t < qlen;
const int64_t r = q0 + (ok ? t : 0);
Kr[h] = ok ? *(const u4v*)(kbase + r * sk_tok) : (u4v){0, 0, 0, 0};
}
#pragma unroll
for (int e = 0; e < 4; ++e) {
const int t = n0 + 4 * vt_g + e;
const bool ok = t < qlen;
const int64_t r = q0 + (ok ? t : 0);
Vr[e] = ok ? *(const u2v*)(vbase + r * sv_tok) : (u2v){0, 0};
}
};
auto store = [&]() {
#pragma unroll
for (int h = 0; h < 2; ++h) {
_Float16* kp = &sK[(lt_tok + 16 * h) * KROW2 + lt_d];
*(u2v*)kp = (u2v){Kr[h][0], Kr[h][1]};
*(u2v*)(kp + 4) = (u2v){Kr[h][2], Kr[h][3]};
}
#pragma unroll
for (int d = 0; d < 4; ++d) {
const uint32_t sel = (d & 1) ? 0x07060302u : 0x05040100u;
const uint32_t lo =
__builtin_amdgcn_perm(Vr[1][d >> 1], Vr[0][d >> 1], sel);
const uint32_t hi =
__builtin_amdgcn_perm(Vr[3][d >> 1], Vr[2][d >> 1], sel);
const int row = vt_d + d;
*(u2v*)&sV[row * VST2 + 4 * (vt_g ^ ((row >> 3) & 7))] = (u2v){lo, hi};
}
};
const int ntiles = (k_end + KT1 - 1) / KT1;
fetch(0);
for (int i = 0; i < ntiles; ++i) {
const int n0 = i * KT1;
store();
__syncthreads();
if (i + 1 < ntiles) fetch(n0 + KT1);
if (wave_live && n0 <= wave_last_tok) {
// S^T for both 16-token halves, one softmax step, then PV.
f4v sacc[2];
#pragma unroll
for (int j = 0; j < 2; ++j) {
sacc[j] = (f4v){0, 0, 0, 0};
const _Float16* krow = &sK[(16 * j + col) * KROW2 + 4 * gq];
#pragma unroll
for (int s = 0; s < FR; ++s)
sacc[j] = __builtin_amdgcn_mfma_f32_16x16x16f16(
*(const h4v*)(krow + 16 * s), Qh[s], sacc[j], 0, 0, 0);
}
float sc[2][4];
float mloc = -INFINITY;
#pragma unroll
for (int j = 0; j < 2; ++j)
#pragma unroll
for (int e = 0; e < 4; ++e) {
const int kabs = n0 + 16 * j + 4 * gq + e;
sc[j][e] = qlive && kabs <= qtok ? sacc[j][e] * c2 : -INFINITY;
mloc = fmaxf(mloc, sc[j][e]);
}
const float mn = fmaxf(m, qmax4(mloc));
const float al = mn == -INFINITY ? 1.f : __builtin_amdgcn_exp2f(m - mn);
h4v pb[2];
float lsum = 0.f;
#pragma unroll
for (int j = 0; j < 2; ++j)
#pragma unroll
for (int e = 0; e < 4; ++e) {
const float p = sc[j][e] == -INFINITY
? 0.f
: __builtin_amdgcn_exp2f(sc[j][e] - mn);
lsum += p;
pb[j][e] = (_Float16)p;
}
m = mn;
l = l * al + lsum;
if (__ballot(al != 1.0f)) {
#pragma unroll
for (int c = 0; c < FR; ++c) acc[c] *= al;
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
// Row 16 c + col; its group 4 j + gq sits at (4 j + gq) ^ (row >> 3).
#pragma unroll
for (int c = 0; c < FR; ++c) {
const int row = 16 * c + col;
const _Float16* vp =
&sV[row * VST2 + 4 * ((4 * j + gq) ^ ((row >> 3) & 7))];
acc[c] = __builtin_amdgcn_mfma_f32_16x16x16f16(*(const h4v*)vp, pb[j],
acc[c], 0, 0, 0);
}
}
}
__syncthreads();
}
l = qsum4(l);
if (!qlive) return;
const float inv = 1.0f / (l + 1e-10f);
__half* orow = out + (int64_t)(q0 + qtok) * so_tok + (int64_t)head * so_head;
#pragma unroll
for (int c = 0; c < FR; ++c)
#pragma unroll
for (int e = 0; e < 4; ++e)
orow[16 * c + 4 * gq + e] = __float2half(acc[c][e] * inv);
#endif // OCTAVE_CDNA_STUB
}
} // namespace octave_cdna
// Called by octave_prefill (octave_prefill_rdna3.cu) on gfx942; same
// arguments, already validated there.
void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k,
torch::Tensor v, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor cu_seqlens_q,
torch::Tensor seq_lens, int64_t max_query_len,
double sm_scale, int64_t fmt, float* ws,
int num_splits) {
using namespace octave_cdna;
auto stream = at::cuda::getCurrentCUDAStream().stream();
const int num_seqs = seq_lens.size(0);
const int hq = q.size(1), hkv = cache.size(1);
const int hpk = hq / hkv;
const int block_size = cache.size(2);
const int max_blocks = block_table.size(1);
const int total_q = q.size(0);
const int rowtiles = (int)((max_query_len * hpk + QB1 - 1) / QB1);
dim3 grid1(num_seqs, hkv, rowtiles * num_splits);
#define PREFIX(B, KC) \
prefix_attn<B, KC><<<grid1, NT1, 0, stream>>>( \
(const __half*)q.data_ptr(), cache.data_ptr<uint8_t>(), \
block_table.data_ptr<int>(), cu_seqlens_q.data_ptr<int>(), \
seq_lens.data_ptr<int>(), ws, (float)sm_scale, hq, hkv, block_size, \
max_blocks, num_splits, total_q, q.stride(0), q.stride(1), \
cache.stride(0), cache.stride(1), cache.stride(2))
if (fmt == 4)
PREFIX(4, false);
else if (fmt == 3)
PREFIX(3, false);
else
PREFIX(3, true);
#undef PREFIX
dim3 grid2(num_seqs, hkv, rowtiles);
chunk_attn<<<grid2, NT1, 0, stream>>>(
(__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(),
(const __half*)k.data_ptr(), (const __half*)v.data_ptr(),
cu_seqlens_q.data_ptr<int>(), seq_lens.data_ptr<int>(), hq, hkv,
max_blocks * block_size, (float)sm_scale, num_splits, total_q,
q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0),
v.stride(1), out.stride(0), out.stride(1));
}
#endif // USE_ROCM
csrc/attention/octave_prefill_rdna3.cu
+809
View File
@@ -0,0 +1,809 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Octave prefill for RDNA3 (gfx11), head size 256, fp16. Reads the packed
// cache directly; no intermediate buffers. Everything is in the rotated space
// (the caller rotates Q/K/V of the chunk and rotates the output back).
//
// Phase 1 (prefix_attn): attention of the chunk's queries over the cached
// prefix, split over the prefix length; writes (O, m, l) partials.
// One block = 128 query vectors (token, head) of one KV head, so the GQA
// heads share every K/V tile. 16 compute waves in pairs: wave g of a pair
// holds Q (int8) and O for head dims [128 g, 128 g + 128). 8 loader waves
// unpack 16-token tiles into double-buffered LDS, one tile ahead.
// QK S^T[token][q] = K x Q^T with int8 WMMA. K rows hold the int8 codebook
// value of every dim. Each wave's 128 dims are two 64-dim blocks with
// their own per-token scale; the wave scales both and the pair
// exchanges the float partial through LDS.
// PV O^T[dim][q] = V^T x P^T with V as its exact fp16 codebook values; the
// V scale is folded into P.
// Phase 2 (chunk_attn): causal attention over the chunk's own K/V (fp16),
// merged with the phase-1 partials; writes the output.
#include <cstdint>
#include <mutex>
#include <unordered_map>
#include <vector>
#include <algorithm>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#if defined(USE_ROCM)
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include "octave_format.cuh"
#if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__)
#define OCTAVE_PREFILL_STUB
#endif
namespace octave_pf {
using octave::D;
using octave::Format;
using octave::LUT_ONE;
typedef _Float16 h16v __attribute__((ext_vector_type(16)));
typedef _Float16 h8v __attribute__((ext_vector_type(8)));
typedef _Float16 h2v __attribute__((ext_vector_type(2)));
typedef float f8v __attribute__((ext_vector_type(8)));
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
typedef int v4i_t __attribute__((ext_vector_type(4)));
typedef int v8i_t __attribute__((ext_vector_type(8)));
constexpr int KT = 16; // tokens per tile
// Splits of the cached prefix used by a sequence: one per kMinSplitTokens, at
// most num_splits. Both phases must agree.
constexpr int kMinSplitTokens = 2048;
__device__ __forceinline__ int used_splits(int ctx, int num_splits) {
if (ctx <= 0) return 0;
return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens);
}
__device__ __forceinline__ int split_len(int ctx, int used) {
return ((ctx + used - 1) / used + KT - 1) / KT * KT;
}
// ===========================================================================
// Phase 1: prefix
// ===========================================================================
constexpr int NW = 8; // compute wave pairs (16 query vectors each)
constexpr int NL = 8; // loader waves
[[maybe_unused]] constexpr int KROW = D + 16; // bytes per K row in LDS
[[maybe_unused]] constexpr float PSCALE = 4096.0f;
// The LDS read of PV step c waits for the WMMA of step c - DEPV, which keeps
// the scheduler from hoisting every read (and spilling).
[[maybe_unused]] constexpr int DEPV = 4;
constexpr int NTHREADS = (2 * NW + NL) * 32;
constexpr int VPAD = 16; // halves after every 8 rows of V^T
constexpr int HD = D / 2;
[[maybe_unused]] constexpr int QSTEPS = HD / 16;
[[maybe_unused]] constexpr int OFRAGS = HD / 16;
__device__ __forceinline__ int vrow(int row) {
return row * KT + (row >> 3) * VPAD;
}
__device__ __forceinline__ uint32_t xhalf(uint32_t v) {
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false,
false);
}
__device__ __forceinline__ float xhalf(float v) {
return __uint_as_float(xhalf(__float_as_uint(v)));
}
// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens).
__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2,
uint32_t a3, uint32_t o[4]) {
const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u);
const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u);
const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u);
const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u);
o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u);
o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u);
o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u);
o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u);
}
// LDS-only barrier: __syncthreads() would also wait for the loaders' global
// prefetch.
__device__ __forceinline__ void tile_barrier() {
asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier" ::: "memory");
}
// Codes of dims 4k..4k+3 of a 64-dim block, one byte each. 4-bit: raw is
// the block's 32 bytes. 3-bit: raw starts at the block's 16 bytes of the
// 2-bit plane; hi0/hi1 are its two dwords of the 1-bit plane.
template <int BITS>
__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw,
uint32_t hi0, uint32_t hi1,
int k) {
if constexpr (BITS == 4) {
return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu;
} else {
const uint32_t hi = k < 8 ? hi0 : hi1;
return ((raw[k >> 2] >> (2 * (k & 3))) & 0x03030303u) |
(((hi >> (k & 7)) & 0x01010101u) << 2);
}
}
template <int VB, bool KC>
__global__ __launch_bounds__(NTHREADS) void prefix_attn(
const __half* __restrict__ q, const uint8_t* __restrict__ cache,
const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q,
const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale,
int num_q_heads, int num_kv_heads, int block_size, int max_blocks,
int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head,
int64_t scb, int64_t sch, int64_t sct) {
#ifndef OCTAVE_PREFILL_STUB
using F = Format<VB, KC>;
const int seq = blockIdx.x, kvh = blockIdx.y;
const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits;
const int tid = threadIdx.x;
const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31;
const int j = lane & 15, h = lane >> 4;
const int pr = w >> 1, g = w & 1;
const int hpk = num_q_heads / num_kv_heads;
const bool loader = w >= 2 * NW;
const int q0 = cu_seqlens_q[seq];
const int qlen = cu_seqlens_q[seq + 1] - q0;
int ctx = seq_lens[seq] - qlen;
ctx = max(0, min(ctx, max_blocks * block_size));
const int nqv = qlen * hpk;
if (rowtile * NW * 16 >= nqv) return;
const int used = used_splits(ctx, num_splits);
if (split >= used) return;
const int k_len = split_len(ctx, used);
const int k_begin = min(ctx, split * k_len);
const int k_stop = min(ctx, k_begin + k_len);
const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT - 1) / KT : 0;
__shared__ __attribute__((aligned(16))) int8_t sK[2][KT][KROW];
__shared__
__attribute__((aligned(16))) _Float16 sV[2][D * KT + (D / 8) * VPAD];
__shared__ __attribute__((aligned(16))) u4v sS[2][2 * NW][2][32];
// K scales travel with K (block 0 = RoPE, 1-3 = NoPE), V scales with V.
__shared__ float sKs[2][4][KT], sVs[2][KT];
// ---- loader waves
if (loader) {
// 256 threads. K: thread lt unpacks dims 16 kp .. 16 kp + 15 of block kq
// of token tk. V: dims 4 vh .. 4 vh + 3 of the 8-dim group vd of token
// quad vq.
const int lt = tid - 2 * NW * 32;
const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3;
const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8;
// Block sizes are multiples of KT, so a tile sits in one cache block.
// Tiles past the end read the last one; tokens past k_stop read the last
// token: their K is masked to -inf and their V meets p = 0.
const int tlast = k_stop - 1;
const uint8_t* const head = cache + (int64_t)kvh * sch;
const int* const bt = block_table + seq * max_blocks;
// Walk the tiles in order: block index and offset inside it, advanced by
// KT per tile and frozen on the last tile.
struct Pos {
int start, blk, off;
};
auto first = [&]() {
return Pos{k_begin, k_begin / block_size, k_begin % block_size};
};
auto next = [&](Pos p) {
if (p.start + KT > tlast) return p;
p.start += KT;
p.off += KT;
if (p.off == block_size) {
p.off = 0;
++p.blk;
}
return p;
};
auto pbase = [&](Pos p, int pb) {
return head + (int64_t)(pb < 0 ? 0 : pb) * scb + (int64_t)p.off * sct;
};
auto tok = [&](int start, int r) {
return (int64_t)min(r, tlast - start) * sct;
};
// 4-bit: dwords 2 kp, 2 kp + 1 of the RoPE codes. 3-bit: dword kp of the
// block's 2-bit plane and dword kp / 2 of its 1-bit plane.
const bool k3 = KC || kq > 0;
const int b = KC ? kq : kq - 1; // 3-bit block index
const int koff0 = k3 ? F::OFF_KN + 16 * b + 4 * kp : 8 * kp;
const int koff1 =
k3 ? F::OFF_KN + F::KN / 4 + 8 * b + 4 * (kp >> 1) : 8 * kp + 4;
constexpr int VW = VB == 4 ? 1 : 2; // dwords of 8 V dims
struct Stg {
uint32_t k0, k1;
uint32_t v[4][VW];
uint32_t s0, s1;
};
auto fetch = [&](int TK, const uint8_t* bK, int TV, const uint8_t* bV,
Stg& S) {
const uint8_t* pk = bK + tok(TK, tk);
S.k0 = *(const uint32_t*)(pk + koff0);
S.k1 = *(const uint32_t*)(pk + koff1);
#pragma unroll
for (int i = 0; i < 4; ++i) {
const uint8_t* pv = bV + tok(TV, 4 * vq + i) + F::OFF_V;
if constexpr (VB == 4) {
S.v[i][0] = *(const uint32_t*)(pv + vd / 2);
} else {
S.v[i][0] = *(const uint32_t*)(pv + 4 * (vd >> 4));
S.v[i][1] = *(const uint32_t*)(pv + D / 4 + 4 * (vd >> 5));
}
}
// Threads 0-15: the K scales of token lt; 16-31: the V scale of token
// lt - KT. Block K: two 8-byte loads inside the slot, V scale in the
// low half of s0. Compact K: one 4-byte load {k, v}.
if (lt < 2 * KT) {
if constexpr (KC) {
S.s0 = *(const uint32_t*)((lt < KT ? bK + tok(TK, lt)
: bV + tok(TV, lt - KT)) +
F::OFF_SC);
} else {
const uint8_t* ps = lt < KT ? bK + tok(TK, lt) + F::OFF_SC
: bV + tok(TV, lt - KT) + F::OFF_SC + 8;
const u2v x = *(const u2v*)ps;
S.s0 = x[0];
S.s1 = x[1];
}
}
};
auto store_k = [&](const Stg& S, int bf) {
const uint32_t r4[2] = {S.k0, S.k1};
const uint32_t r3[1] = {S.k0};
const uint32_t hi = S.k1 >> (4 * (kp & 1));
u4v o;
#pragma unroll
for (int x = 0; x < 4; ++x)
o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, hi, x))
: octave::lut4(block_codes<4>(r4, 0, 0, x));
*(u4v*)&sK[bf][tk][64 * kq + 16 * kp] = o;
if (lt < KT) {
const h2v a = __builtin_bit_cast(h2v, S.s0);
if constexpr (KC) {
sKs[bf][0][lt] = (float)a[0];
} else {
const h2v c = __builtin_bit_cast(h2v, S.s1);
sKs[bf][0][lt] = (float)a[0];
sKs[bf][1][lt] = (float)a[1];
sKs[bf][2][lt] = (float)c[0];
sKs[bf][3][lt] = (float)c[1];
}
}
};
auto store_v = [&](const Stg& S, int bf) {
uint32_t c[4]; // [token] code bytes of dims vd + 4 vh ..
#pragma unroll
for (int i = 0; i < 4; ++i) {
if constexpr (VB == 4) {
c[i] = (S.v[i][0] >> (4 * vh)) & 0x0F0F0F0Fu;
} else {
const int sl = 2 * ((vd & 15) >> 2) + 2 * vh;
const int sh = ((vd & 31) >> 2) + vh;
c[i] = ((S.v[i][0] >> sl) & 0x03030303u) |
(((S.v[i][1] >> sh) & 0x01010101u) << 2);
}
c[i] = octave::lut<VB>(c[i]) ^ 0x80808080u;
}
uint32_t t[4];
tr4(c[0], c[1], c[2], c[3], t);
// Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x.
const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f};
#pragma unroll
for (int d = 0; d < 4; ++d) {
const int row = vd + 4 * vh + d, sw = ((row >> 3) & 1) * 8;
const h2v p01 = __builtin_bit_cast(
h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u));
const h2v p23 = __builtin_bit_cast(
h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u));
*(u2v*)&sV[bf][vrow(row) + ((4 * vq) ^ sw)] =
(u2v){__builtin_bit_cast(uint32_t, p01 - kofs),
__builtin_bit_cast(uint32_t, p23 - kofs)};
}
if (lt >= KT && lt < 2 * KT)
sVs[bf][lt - KT] = (float)__builtin_bit_cast(
_Float16, (uint16_t)(KC ? S.s0 >> 16 : S.s0)) *
(PSCALE / LUT_ONE);
};
// Pipeline: K of tile i + 1 and V of tile i are written in iteration i
// from registers; the block number of a tile is loaded two tiles ahead.
Stg S;
if (ntiles > 0) {
const Pos p0 = first(), p1 = next(p0), p2 = next(p1);
const uint8_t* b0 = pbase(p0, bt[p0.blk]);
const int q1 = bt[p1.blk];
fetch(p0.start, b0, p0.start, b0, S);
const int q2 = bt[p2.blk];
Pos p3 = next(p2);
int pn = bt[p3.blk]; // block of tile i + 3
const uint8_t* bv = pbase(p1, q1);
const uint8_t* bk = pbase(p2, q2);
int sv = p1.start, sk = p2.start;
store_k(S, 0);
fetch(p1.start, bv, p0.start, b0, S);
tile_barrier();
for (int i = 0; i < ntiles; ++i) {
store_k(S, (i + 1) & 1);
store_v(S, i & 1);
fetch(sk, bk, sv, bv, S);
bv = bk;
sv = sk;
bk = pbase(p3, pn);
sk = p3.start;
p3 = next(p3);
pn = bt[p3.blk];
tile_barrier();
}
} else {
tile_barrier();
}
tile_barrier();
return;
}
// ---- compute waves: this lane's query row, int8 over this wave's dims
const int qr = pr * 16 + j;
const int qv = rowtile * NW * 16 + qr;
const bool qlive = qv < nqv;
u4v Q[QSTEPS];
float qsc;
{
const int qtok = qlive ? qv / hpk : 0;
const int qhead = kvh * hpk + (qlive ? qv % hpk : 0);
const __half* qrow =
q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head;
float qmax = 1e-8f;
for (int i = 0; i < D / 8; ++i) {
const h8v x = qlive ? *(const h8v*)(qrow + 8 * i) : (h8v){};
#pragma unroll
for (int k = 0; k < 8; ++k) qmax = fmaxf(qmax, fabsf((float)x[k]));
}
const float qinv = 127.0f / qmax;
qsc = qmax * (1.0f / 127.0f);
#pragma unroll
for (int f = 0; f < QSTEPS; ++f) {
uint32_t wd[4];
#pragma unroll
for (int hh = 0; hh < 2; ++hh) {
const h8v x =
qlive ? *(const h8v*)(qrow + HD * g + 16 * f + 8 * hh) : (h8v){};
#pragma unroll
for (int d4 = 0; d4 < 2; ++d4) {
uint32_t word = 0;
#pragma unroll
for (int b = 0; b < 4; ++b) {
int v = (int)lrintf((float)x[4 * d4 + b] * qinv);
v = max(-127, min(127, v));
word |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b);
}
wd[2 * hh + d4] = word;
}
}
Q[f] = (u4v){wd[0], wd[1], wd[2], wd[3]};
}
}
f8v O[OFRAGS];
#pragma unroll
for (int c = 0; c < OFRAGS; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
// A operands are only read in lanes with (j % 2) == h; the other lane of
// each pair reads its partner's row, so the pair costs one LDS fetch.
const int rr = ((j & 1) == h) ? j : (j ^ 1);
const int sw = ((rr >> 3) & 1) * 8;
const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE);
float m = -INFINITY, l = 0.f;
f8v Sp = {0, 0, 0, 0, 0, 0, 0, 0}; // this wave's partial S, tile i-1
tile_barrier();
for (int i = 0; i <= ntiles; ++i) {
f8v Sn = {0, 0, 0, 0, 0, 0, 0, 0};
if (i < ntiles) {
// QK over this wave's two 64-dim blocks, S^T = K x Q^T, 4 k-steps
// each; block s is scaled into Sn before block s + 1 starts, so one
// int accumulator is live at a time.
uint32_t koff = (uint32_t)(rr * KROW + HD * g);
if constexpr (KC) { // one K scale: one accumulator for both blocks
v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll
for (int f = 0; f < 8; ++f) {
const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f);
Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(
true, __builtin_bit_cast(v4i_t, a), true,
__builtin_bit_cast(v4i_t, Q[f]), Si, false);
}
const float* ks = &sKs[i & 1][0][0];
#pragma unroll
for (int e = 0; e < 8; ++e) Sn[e] = ks[2 * e + h] * (float)Si[e];
}
#pragma unroll
for (int sg = 0; sg < (KC ? 0 : 2); ++sg) {
v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll
for (int f = 4 * sg; f < 4 * sg + 4; ++f) {
const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f);
Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(
true, __builtin_bit_cast(v4i_t, a), true,
__builtin_bit_cast(v4i_t, Q[f]), Si, false);
}
const float* ks = &sKs[i & 1][2 * g + sg][0];
#pragma unroll
for (int e = 0; e < 8; ++e) Sn[e] += ks[2 * e + h] * (float)Si[e];
if (sg == 0) asm volatile("" : "+v"(koff) : "v"(Sn[0]));
}
const u8v su = __builtin_bit_cast(u8v, Sn);
sS[i & 1][w][0][lane] = (u4v){su[0], su[1], su[2], su[3]};
sS[i & 1][w][1][lane] = (u4v){su[4], su[5], su[6], su[7]};
}
if (i > 0) {
const int bf = (i - 1) & 1, base = k_begin + (i - 1) * KT;
const u4v o0 = sS[bf][w ^ 1][0][lane], o1 = sS[bf][w ^ 1][1][lane];
const f8v S =
Sp + __builtin_bit_cast(f8v, (u8v){o0[0], o0[1], o0[2], o0[3], o1[0],
o1[1], o1[2], o1[3]});
// Scores in log2 units are S * c1 (c1 > 0). Only the last tile has
// keys past k_stop; exp2(-inf) = 0, and m = -inf only before the first
// tile, so nothing else needs a guard.
float s[8];
#pragma unroll
for (int e = 0; e < 8; ++e) s[e] = S[e];
if (base + KT > k_stop) {
#pragma unroll
for (int e = 0; e < 8; ++e)
if (base + 2 * e + h >= k_stop) s[e] = -INFINITY;
}
float mloc = fmaxf(fmaxf(fmaxf(s[0], s[1]), fmaxf(s[2], s[3])),
fmaxf(fmaxf(s[4], s[5]), fmaxf(s[6], s[7])));
const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc)) * c1);
const float alpha = __builtin_amdgcn_exp2f(m - mn);
m = mn;
uint32_t pd[4];
float lsum = 0.f;
#pragma unroll
for (int e = 0; e < 8; e += 2) {
const float p0 = __builtin_amdgcn_exp2f(fmaf(s[e], c1, -mn));
const float p1 = __builtin_amdgcn_exp2f(fmaf(s[e + 1], c1, -mn));
lsum += p0 + p1;
const _Float16 h0 = (_Float16)(p0 * sVs[bf][2 * e + h]);
const _Float16 h1 = (_Float16)(p1 * sVs[bf][2 * e + 2 + h]);
pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1});
}
l = l * alpha + lsum;
if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) {
#pragma unroll
for (int c = 0; c < OFRAGS; ++c) O[c] *= alpha;
}
u8v pf;
// B operand in natural token order: lane half 0 holds the even tokens,
// half 1 the odd ones; interleave them per pair.
#pragma unroll
for (int k = 0; k < 4; ++k) {
const uint32_t o = xhalf(pd[k]);
const uint32_t ev = h ? o : pd[k], od = h ? pd[k] : o;
pf[2 * k] = __builtin_amdgcn_perm(od, ev, 0x05040100u);
pf[2 * k + 1] = __builtin_amdgcn_perm(od, ev, 0x07060302u);
}
const h16v pb16 = __builtin_bit_cast(h16v, pf);
#pragma unroll
for (int c = 0; c < OFRAGS; ++c) {
uint32_t voff = (uint32_t)vrow(HD * g + 16 * c + rr);
if (c >= DEPV) asm volatile("" : "+v"(voff) : "v"(O[c - DEPV][0]));
const _Float16* vr = &sV[bf][0] + voff;
const h8v lo = *(const h8v*)(vr + sw);
const h8v hi = *(const h8v*)(vr + (8 ^ sw));
const h16v a = __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7,
8, 9, 10, 11, 12, 13, 14, 15);
O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]);
}
}
Sp = Sn;
tile_barrier();
}
// ---- partials for the phase-2 merge (natural-log max, unnormalized O)
l += xhalf(l);
if (!qlive) return;
const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk;
float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) +
((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2) + HD * g;
#pragma unroll
for (int c = 0; c < OFRAGS; ++c)
#pragma unroll
for (int e = 0; e < 8; ++e)
wp[16 * c + 2 * e + h] = O[c][e] * (1.0f / PSCALE);
if (g == 0 && h == 0) {
wp[D] = m * 0.6931471805599453f;
wp[D + 1] = l;
}
#endif // OCTAVE_PREFILL_STUB
}
// ===========================================================================
// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix
// partials. 8 waves, 16 query rows each, one query head per block.
// ===========================================================================
constexpr int CW = 8; // waves
constexpr int CM = 16 * CW; // query rows per block
[[maybe_unused]] constexpr int FRAGS = D / 16;
[[maybe_unused]] constexpr int X = 8; // fp16 per 16-byte chunk
__device__ __forceinline__ float row16_max(float v) {
v = fmaxf(v, __shfl_xor(v, 1));
v = fmaxf(v, __shfl_xor(v, 2));
v = fmaxf(v, __shfl_xor(v, 4));
return fmaxf(v, __shfl_xor(v, 8));
}
__device__ __forceinline__ float row16_sum(float v) {
v += __shfl_xor(v, 1);
v += __shfl_xor(v, 2);
v += __shfl_xor(v, 4);
return v + __shfl_xor(v, 8);
}
__global__ void __launch_bounds__(D)
chunk_attn(__half* __restrict__ out, const float* __restrict__ ws,
const __half* __restrict__ q, const __half* __restrict__ k,
const __half* __restrict__ v,
const int* __restrict__ cu_seqlens_q,
const int* __restrict__ seq_lens, int num_q_heads,
int num_kv_heads, int max_ctx, float sm_scale, int num_splits,
int total_q_tokens, int64_t sq_tok, int64_t sq_head,
int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head,
int64_t so_tok, int64_t so_head) {
#ifndef OCTAVE_PREFILL_STUB
const int seq = blockIdx.x, head = blockIdx.y;
const int tid = threadIdx.x, wave = tid >> 5, lane = tid & 31;
const int lo = lane & 15, hi = lane >> 4;
const int q0 = cu_seqlens_q[seq];
const int qlen = cu_seqlens_q[seq + 1] - q0;
const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx));
const int tile0 = blockIdx.z * CM;
if (tile0 >= qlen) return;
const int wq0 = tile0 + wave * 16;
const int nq_wave = max(0, min(16, qlen - wq0));
const int kvh = head / (num_q_heads / num_kv_heads);
__shared__ __attribute__((aligned(16))) _Float16 sK[FRAGS * 2 * KT * X];
__shared__ __attribute__((aligned(16))) _Float16 sV[D * KT];
__shared__ __attribute__((aligned(16))) _Float16 sP[CW][16 * KT];
float m_st[8], l_st[8];
f8v acc[FRAGS];
// Merge the prefix splits (one weight per split; exp(0) = 1 with one).
const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2);
const int used = used_splits(ctx, num_splits);
#pragma unroll
for (int i = 0; i < 8; ++i) {
m_st[i] = -INFINITY;
l_st[i] = 0.f;
}
#pragma unroll
for (int f = 0; f < FRAGS; ++f) acc[f] = (f8v){0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int qp = wq0 + 2 * i + hi;
if (qp >= qlen) continue;
const float* w0 = ws + ((int64_t)(q0 + qp) * num_q_heads + head) * (D + 2);
float mx = -INFINITY;
for (int s = 0; s < used; ++s) mx = fmaxf(mx, w0[s * split_stride + D]);
float lsum = 0.f;
for (int s = 0; s < used; ++s) {
const float* wp = w0 + s * split_stride;
if (wp[D] == -INFINITY) continue;
const float wgt = __expf(wp[D] - mx);
lsum += wp[D + 1] * wgt;
#pragma unroll
for (int f = 0; f < FRAGS; ++f) acc[f][i] += wp[16 * f + lo] * wgt;
}
m_st[i] = mx;
l_st[i] = lsum;
}
// Q rows are re-read from global (L1/L2) in the loop instead of being held
// in registers; out-of-range rows are masked anyway.
const __half* qrow = q + (int64_t)(q0 + min(wq0 + lo, qlen - 1)) * sq_tok +
(int64_t)head * sq_head;
const int k_end = min(qlen, tile0 + min(CM, qlen - tile0));
_Float16* sPw = &sP[wave][0];
for (int n0 = 0; n0 < k_end; n0 += KT) {
{ // K tile: [16-dim chunk][token][8] fp16
const int tk = tid / (FRAGS), dc = (tid % FRAGS) * 2;
const bool ok = n0 + tk < qlen;
const __half* row =
k + (int64_t)(q0 + n0 + tk) * sk_tok + (int64_t)kvh * sk_head;
#pragma unroll
for (int d = 0; d < 2; ++d) {
int4 x = ok ? *(const int4*)(row + (dc + d) * X) : int4{0, 0, 0, 0};
*(int4*)&sK[(dc + d) * (KT * X) + tk * X] = x;
}
}
{ // V tile transposed: [dim][token]
#pragma unroll
for (int p = 0; p < 2; ++p) {
const int tk = tid / (D / KT),
dbase = ((tid % (D / KT)) + p * (D / KT)) * 8;
const bool ok = n0 + tk < qlen;
int4 x = ok ? *(const int4*)(v + (int64_t)(q0 + n0 + tk) * sv_tok +
(int64_t)kvh * sv_head + dbase)
: int4{0, 0, 0, 0};
_Float16 t8[8];
__builtin_memcpy(t8, &x, 16);
#pragma unroll
for (int e = 0; e < 8; ++e) sV[(dbase + e) * KT + tk] = t8[e];
}
}
__syncthreads();
if (nq_wave > 0 && n0 <= wq0 + nq_wave - 1) {
f8v s = {0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll 2
for (int f = 0; f < FRAGS; ++f) {
h16v b, a;
const int4 klo = *(const int4*)&sK[(2 * f) * (KT * X) + lo * X];
const int4 khi = *(const int4*)&sK[(2 * f + 1) * (KT * X) + lo * X];
__builtin_memcpy(&b, &klo, 16);
__builtin_memcpy(((char*)&b) + 16, &khi, 16);
__builtin_memcpy(&a, qrow + 16 * f, 32);
s = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, s);
}
const int kabs = n0 + lo;
float p[8];
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int r = 2 * i + hi;
const bool keep = r < nq_wave && kabs < qlen && kabs <= wq0 + r;
const float sc = keep ? s[i] * sm_scale : -INFINITY;
const float mi = row16_max(sc);
const float mn = fmaxf(m_st[i], mi);
const float al = m_st[i] == -INFINITY ? 0.f : __expf(m_st[i] - mn);
p[i] = mn == -INFINITY ? 0.f : __expf(sc - mn);
l_st[i] = l_st[i] * al + row16_sum(p[i]);
m_st[i] = mn;
#pragma unroll
for (int f = 0; f < FRAGS; ++f) acc[f][i] *= al;
sPw[r * KT + lo] = (_Float16)p[i];
}
h16v pa;
const int4 plo = *(const int4*)&sPw[lo * KT];
const int4 phi = *(const int4*)&sPw[lo * KT + 8];
__builtin_memcpy(&pa, &plo, 16);
__builtin_memcpy(((char*)&pa) + 16, &phi, 16);
#pragma unroll
for (int f = 0; f < FRAGS; ++f) {
h16v vb;
const int4 vlo = *(const int4*)&sV[(16 * f + lo) * KT];
const int4 vhi = *(const int4*)&sV[(16 * f + lo) * KT + 8];
__builtin_memcpy(&vb, &vlo, 16);
__builtin_memcpy(((char*)&vb) + 16, &vhi, 16);
acc[f] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(pa, vb, acc[f]);
}
}
__syncthreads();
}
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int qp = wq0 + 2 * i + hi;
if (qp >= qlen) continue;
const float inv = 1.0f / (l_st[i] + 1e-10f);
__half* orow = out + (int64_t)(q0 + qp) * so_tok + (int64_t)head * so_head;
#pragma unroll
for (int f = 0; f < FRAGS; ++f)
orow[16 * f + lo] = __float2half(acc[f][i] * inv);
}
#endif // OCTAVE_PREFILL_STUB
}
// Partials of phase 1, kept alive for the process: a graph that captured a
// pointer must keep finding it valid. Grows geometrically, so it settles.
float* workspace(int64_t need, const at::TensorOptions& opts) {
static std::mutex mtx;
static std::unordered_map<int, std::vector<at::Tensor>> kept;
std::lock_guard<std::mutex> lock(mtx);
std::vector<at::Tensor>& bufs = kept[opts.device().index()];
if (bufs.empty() || bufs.back().numel() < need) {
int64_t n = 8 << 20;
while (n < need) n *= 2;
bufs.push_back(at::empty({n}, opts));
}
return (float*)bufs.back().data_ptr();
}
} // namespace octave_pf
int octave_slot_bytes(int64_t fmt);
void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k,
torch::Tensor v, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor cu_seqlens_q,
torch::Tensor seq_lens, int64_t max_query_len,
double sm_scale, int64_t fmt, float* ws,
int num_splits);
void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k,
torch::Tensor v, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor cu_seqlens_q,
torch::Tensor seq_lens, int64_t max_query_len,
double sm_scale, int64_t fmt) {
using namespace octave_pf;
TORCH_CHECK(q.dtype() == at::kHalf && k.dtype() == at::kHalf &&
v.dtype() == at::kHalf && out.dtype() == at::kHalf,
"octave_prefill: fp16 only");
TORCH_CHECK(q.size(2) == D && q.stride(2) == 1 && k.stride(2) == 1 &&
v.stride(2) == 1 && out.stride(2) == 1);
TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4 &&
cache.stride(3) == 1);
TORCH_CHECK(cache.size(3) == octave_slot_bytes(fmt),
"octave_prefill: slot size mismatch");
const int align = fmt & octave::kCompactFlag ? 4 : 8;
TORCH_CHECK(cache.stride(2) % align == 0 && cache.stride(1) % align == 0 &&
cache.stride(0) % align == 0,
"octave_prefill: misaligned slots");
// 0: gfx1100 (WMMA), 1: gfx942 (MFMA, octave_prefill_cdna.cu).
static const int arch = [] {
const std::string name =
at::cuda::getCurrentDeviceProperties()->gcnArchName;
if (name.rfind("gfx1100", 0) == 0) return 0;
if (name.rfind("gfx942", 0) == 0) return 1;
return -1;
}();
TORCH_CHECK(arch >= 0, "octave_prefill: gfx1100 or gfx942 only");
const int num_seqs = seq_lens.size(0);
if (num_seqs == 0 || q.size(0) == 0) return;
const at::cuda::OptionalCUDAGuard guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream().stream();
const int hq = q.size(1), hkv = cache.size(1);
TORCH_CHECK(hq % hkv == 0);
const int hpk = hq / hkv;
const int block_size = cache.size(2);
const int max_blocks = block_table.size(1);
TORCH_CHECK(block_table.stride(1) == 1 && block_table.stride(0) == max_blocks,
"octave_prefill: block_table rows must be contiguous");
const int total_q = q.size(0);
const int rowtiles = (int)((max_query_len * hpk + 127) / 128);
const int base_blocks = std::max(1, num_seqs * hkv * rowtiles);
// Blocks whose split has no keys exit at once, so ask for plenty.
const int num_splits =
std::min(16, std::max(1, (1024 + base_blocks - 1) / base_blocks));
float* ws = workspace((int64_t)num_splits * total_q * hq * (D + 2),
q.options().dtype(at::kFloat));
if (arch == 1) {
octave_prefill_cdna(out, q, k, v, cache, block_table, cu_seqlens_q,
seq_lens, max_query_len, sm_scale, fmt, ws, num_splits);
return;
}
dim3 grid1(num_seqs, hkv, rowtiles * num_splits);
#define PREFIX(B, KC) \
prefix_attn<B, KC><<<grid1, NTHREADS, 0, stream>>>( \
(const __half*)q.data_ptr(), cache.data_ptr<uint8_t>(), \
block_table.data_ptr<int>(), cu_seqlens_q.data_ptr<int>(), \
seq_lens.data_ptr<int>(), ws, (float)sm_scale, hq, hkv, block_size, \
max_blocks, num_splits, total_q, q.stride(0), q.stride(1), \
cache.stride(0), cache.stride(1), cache.stride(2))
if (fmt == 4)
PREFIX(4, false);
else if (fmt == 3)
PREFIX(3, false);
else
PREFIX(3, true);
#undef PREFIX
dim3 grid2(num_seqs, hq, (max_query_len + CM - 1) / CM);
chunk_attn<<<grid2, D, 0, stream>>>(
(__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(),
(const __half*)k.data_ptr(), (const __half*)v.data_ptr(),
cu_seqlens_q.data_ptr<int>(), seq_lens.data_ptr<int>(), hq, hkv,
max_blocks * block_size, (float)sm_scale, num_splits, total_q,
q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0),
v.stride(1), out.stride(0), out.stride(1));
}
#else
void octave_prefill(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
int64_t, double, int64_t) {
TORCH_CHECK(false, "octave requires ROCm");
}
#endif
csrc/attention/octave_torch_bindings.cpp
+82
View File
@@ -0,0 +1,82 @@
// Op registration for the Octave KV cache kernels (legacy _C extension, full
// torch API). Built only for gfx1100 and gfx942; see CMakeLists.txt.
#include <optional>
#include <torch/all.h>
#include <torch/library.h>
#include "core/registration.h"
// fmt: V bits (3 or 4), plus 16 for compact K.
void octave_cache_store(torch::Tensor key, torch::Tensor value,
torch::Tensor cache, torch::Tensor slot_mapping,
torch::Tensor k_signs, torch::Tensor v_signs,
int64_t fmt);
void octave_decode(torch::Tensor out, torch::Tensor query, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor q_to_req,
torch::Tensor q_to_klen, torch::Tensor mid_o,
torch::Tensor k_signs, torch::Tensor v_signs,
double sm_scale, int64_t num_kv_splits, int64_t fmt,
int64_t query_group, bool use_wmma);
void octave_decode_sparse(torch::Tensor out, torch::Tensor query,
torch::Tensor cache, torch::Tensor block_table,
torch::Tensor q_to_req, torch::Tensor indices,
torch::Tensor mid_o, torch::Tensor k_signs,
torch::Tensor v_signs, double sm_scale,
int64_t num_kv_splits, int64_t fmt,
const std::optional<torch::Tensor>& positions,
const std::optional<torch::Tensor>& wtab,
const std::optional<torch::Tensor>& wtags,
const std::optional<torch::Tensor>& stab,
const std::optional<torch::Tensor>& stags);
void octave_window_store(torch::Tensor key, torch::Tensor value,
torch::Tensor slot_mapping, torch::Tensor positions,
torch::Tensor wtab, torch::Tensor wtags,
torch::Tensor wlocks, torch::Tensor stab,
torch::Tensor stags, torch::Tensor slocks);
void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout,
bool inverse);
void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k,
torch::Tensor v, torch::Tensor cache,
torch::Tensor block_table, torch::Tensor cu_seqlens_q,
torch::Tensor seq_lens, int64_t max_query_len,
double sm_scale, int64_t fmt);
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"octave_cache_store(Tensor key, Tensor value, Tensor! cache, "
"Tensor slot_mapping, Tensor k_signs, Tensor v_signs, int fmt) -> ()");
ops.impl("octave_cache_store", torch::kCUDA, &octave_cache_store);
ops.def(
"octave_decode(Tensor! out, Tensor query, Tensor cache, "
"Tensor block_table, Tensor q_to_req, Tensor q_to_klen, "
"Tensor! mid_o, Tensor k_signs, Tensor v_signs, float sm_scale, "
"int num_kv_splits, int fmt, int query_group, bool use_wmma) -> ()");
ops.impl("octave_decode", torch::kCUDA, &octave_decode);
ops.def(
"octave_decode_sparse(Tensor! out, Tensor query, Tensor cache, "
"Tensor block_table, Tensor q_to_req, Tensor indices, Tensor! mid_o, "
"Tensor k_signs, Tensor v_signs, float sm_scale, int num_kv_splits, "
"int fmt, Tensor? positions=None, Tensor? wtab=None, "
"Tensor? wtags=None, Tensor? stab=None, Tensor? stags=None) -> ()");
ops.impl("octave_decode_sparse", torch::kCUDA, &octave_decode_sparse);
ops.def(
"octave_window_store(Tensor key, Tensor value, Tensor slot_mapping, "
"Tensor positions, Tensor! wtab, Tensor! wtags, Tensor! wlocks, "
"Tensor! stab, Tensor! stags, Tensor! slocks) -> ()");
ops.impl("octave_window_store", torch::kCUDA, &octave_window_store);
ops.def(
"octave_rotate(Tensor! x, Tensor signs, bool k_layout, bool inverse) "
"-> ()");
ops.impl("octave_rotate", torch::kCUDA, &octave_rotate);
ops.def(
"octave_prefill(Tensor! out, Tensor q, Tensor k, Tensor v, Tensor cache, "
"Tensor block_table, Tensor cu_seqlens_q, Tensor seq_lens, "
"int max_query_len, float sm_scale, int fmt) -> ()");
ops.impl("octave_prefill", torch::kCUDA, &octave_prefill);
}
csrc/cache_kernels.cu
+19 -4
View File
@@ -420,7 +420,8 @@ __global__ void reshape_and_cache_flash_kernel(
}
}
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
bool skip_scale>
__global__ void concat_and_cache_mla_kernel(
const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim]
@@ -453,6 +454,9 @@ __global__ void concat_and_cache_mla_kernel(
block_idx * block_stride + block_offset * entry_stride + i + offset;
if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
dst[dst_idx] = src[src_idx];
} else if constexpr (skip_scale) {
dst[dst_idx] =
fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], 1.0f);
} else {
dst[dst_idx] =
fp8::scaled_convert<cache_t, scalar_t, kv_dt>(src[src_idx], *scale);
@@ -910,8 +914,8 @@ void reshape_and_cache_flash(
// KV_T is the data type of key and value tensors.
// CACHE_T is the stored data type of kv-cache.
// KV_DTYPE is the real data type of kv-cache.
#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
vllm::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE> \
#define LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, SKIP_SCALE) \
vllm::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE, SKIP_SCALE> \
<<<grid, block, 0, stream>>>( \
reinterpret_cast<KV_T*>(kv_c.data_ptr()), \
reinterpret_cast<KV_T*>(k_pe.data_ptr()), \
@@ -920,6 +924,13 @@ void reshape_and_cache_flash(
kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \
reinterpret_cast<const float*>(scale.data_ptr()));
#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \
if (skip_scale) { \
LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, true) \
} else { \
LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, false) \
}
// KV_T is the data type of key and value tensors.
// CACHE_T is the stored data type of kv-cache.
#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \
@@ -938,7 +949,8 @@ void concat_and_cache_mla(
torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank
// + pe_dim)]
torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
const std::string& kv_cache_dtype, torch::stable::Tensor& scale) {
const std::string& kv_cache_dtype, torch::stable::Tensor& scale,
bool skip_scale) {
// NOTE(woosuk): In vLLM V1, key.size(0) can be different from
// slot_mapping.size(0) because of padding for CUDA graphs.
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
@@ -1030,6 +1042,9 @@ void concat_and_cache_mla(
}
}
#undef CALL_CONCAT_AND_CACHE_MLA
#undef LAUNCH_CONCAT_AND_CACHE_MLA
void concat_and_cache_mla_grouped(
torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank]
torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim]
csrc/cooperative_topk.cuh
+7 -3
View File
@@ -536,9 +536,13 @@ __device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
return;
}
// Large path: init mbarriers + state, then dispatch fused or twopass
const uint32_t per_block =
(params.stride + CS - 1) / CS; // how many elements per block
// Large path: init mbarriers + state, then dispatch fused or twopass.
// Elements per block, sized by the row's length rather than its stride so a
// short row in a wide logits buffer still takes the single-pass staging.
// large_topk splits the row in 16-byte units, but kFusedMax and
// kSizePerStage are multiples of that unit, so the path and stage count
// come out the same.
const uint32_t per_block = (static_cast<uint32_t>(sl) + CS - 1) / CS;
constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16
: (CS == 8) ? kFusedStagesCS8
: kMaxSinglePassStages) *
csrc/cpu/fp8_quant.cpp
+285
View File
@@ -0,0 +1,285 @@
#include <torch/all.h>
#include <algorithm>
#include <cmath>
#if defined(__AVX512F__)
#include <immintrin.h>
#endif
namespace {
template <typename fp8_t>
constexpr float fp8_max() {
return std::is_same_v<fp8_t, c10::Float8_e5m2> ? 57344.0f : 448.0f;
}
template <typename fp8_t>
inline fp8_t to_fp8(float x) {
constexpr float max = fp8_max<fp8_t>();
return static_cast<fp8_t>(std::min(std::max(x, -max), max));
}
#if defined(__AVX512F__)
inline __m512 load_as_fp32(const float* p) { return _mm512_loadu_ps(p); }
inline __m512 load_as_fp32(const c10::BFloat16* p) {
const __m256i v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(p));
return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(v), 16));
}
inline __m512 load_as_fp32(const c10::Half* p) {
return _mm512_cvtph_ps(
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(p)));
}
// Bit-exact vector form of c10::Float8_e4m3fn(float) (round-to-nearest-even)
// for inputs already clamped to [-448, 448] or NaN.
inline __m128i fp32_to_e4m3fn(__m512 x) {
const __m512i bits = _mm512_castps_si512(x);
const __m512i sign =
_mm512_and_si512(bits, _mm512_set1_epi32(static_cast<int>(0x80000000u)));
const __m512i abs_bits = _mm512_xor_si512(bits, sign);
const __mmask16 nan_mask =
_mm512_cmpge_epu32_mask(abs_bits, _mm512_set1_epi32(1087 << 20));
const __mmask16 denorm_mask =
_mm512_cmplt_epu32_mask(abs_bits, _mm512_set1_epi32(121 << 23));
const __m512i denorm_offset = _mm512_set1_epi32(141 << 23);
const __m512i denorm = _mm512_sub_epi32(
_mm512_castps_si512(_mm512_add_ps(_mm512_castsi512_ps(abs_bits),
_mm512_castsi512_ps(denorm_offset))),
denorm_offset);
const __m512i mant_odd =
_mm512_and_si512(_mm512_srli_epi32(abs_bits, 20), _mm512_set1_epi32(1));
const __m512i normal = _mm512_srli_epi32(
_mm512_add_epi32(
_mm512_add_epi32(abs_bits, _mm512_set1_epi32(static_cast<int>(
(uint32_t(7 - 127) << 23) + 0x7FFFF))),
mant_odd),
20);
__m512i res = _mm512_mask_blend_epi32(denorm_mask, normal, denorm);
res = _mm512_mask_mov_epi32(res, nan_mask, _mm512_set1_epi32(0x7F));
res = _mm512_or_si512(res, _mm512_srli_epi32(sign, 24));
return _mm512_cvtepi32_epi8(res);
}
inline __m512 clamp_fp8_range(__m512 x, float max) {
// Operand order propagates NaN, matching the scalar path.
return _mm512_min_ps(_mm512_set1_ps(max),
_mm512_max_ps(_mm512_set1_ps(-max), x));
}
#endif
// out[i] = fp8(in[i] * scale) when kDivide is false, fp8(in[i] / scale)
// otherwise.
template <bool kDivide, typename scalar_t, typename fp8_t>
inline void quant_row(fp8_t* __restrict__ out, const scalar_t* __restrict__ in,
int64_t n, float scale) {
int64_t i = 0;
#if defined(__AVX512F__)
if constexpr (std::is_same_v<fp8_t, c10::Float8_e4m3fn> &&
!std::is_same_v<scalar_t, double>) {
const __m512 vscale = _mm512_set1_ps(scale);
for (; i + 16 <= n; i += 16) {
const __m512 v = load_as_fp32(in + i);
const __m512 scaled =
kDivide ? _mm512_div_ps(v, vscale) : _mm512_mul_ps(v, vscale);
_mm_storeu_si128(
reinterpret_cast<__m128i*>(out + i),
fp32_to_e4m3fn(clamp_fp8_range(scaled, fp8_max<fp8_t>())));
}
}
#endif
for (; i < n; ++i) {
const float v = static_cast<float>(in[i]);
out[i] = to_fp8<fp8_t>(kDivide ? v / scale : v * scale);
}
}
template <typename scalar_t>
inline float row_absmax(const scalar_t* __restrict__ in, int64_t n) {
int64_t i = 0;
float absmax = 0.0f;
#if defined(__AVX512F__)
if constexpr (!std::is_same_v<scalar_t, double>) {
__m512 vmax = _mm512_setzero_ps();
for (; i + 16 <= n; i += 16) {
// Operand order skips NaN, matching fmaxf in the CUDA kernel.
vmax = _mm512_max_ps(_mm512_abs_ps(load_as_fp32(in + i)), vmax);
}
absmax = _mm512_reduce_max_ps(vmax);
}
#endif
for (; i < n; ++i) {
absmax = std::max(absmax, std::fabs(static_cast<float>(in[i])));
}
return absmax;
}
template <typename scalar_t, typename fp8_t>
void static_quant_impl(fp8_t* __restrict__ out, const scalar_t* __restrict__ in,
const float* __restrict__ scale, int64_t num_tokens,
int64_t hidden_size, int64_t in_row_stride,
int64_t out_row_stride, int64_t group_m, int64_t group_n,
int64_t scale_stride_i, int64_t scale_stride_j) {
#pragma omp parallel for
for (int64_t m = 0; m < num_tokens; ++m) {
const scalar_t* row_in = in + m * in_row_stride;
fp8_t* row_out = out + m * out_row_stride;
const float* row_scale = scale + (m / group_m) * scale_stride_i;
if (scale_stride_j == 0) {
quant_row<false>(row_out, row_in, hidden_size, 1.0f / row_scale[0]);
} else {
for (int64_t n = 0; n < hidden_size; n += group_n) {
quant_row<false>(row_out + n, row_in + n, group_n,
1.0f / row_scale[(n / group_n) * scale_stride_j]);
}
}
}
}
template <typename scalar_t, typename fp8_t>
void dynamic_per_token_quant_impl(fp8_t* __restrict__ out,
float* __restrict__ scale,
const scalar_t* __restrict__ in,
const float* scale_ub, int64_t num_tokens,
int64_t hidden_size, int64_t in_row_stride,
int64_t out_row_stride) {
constexpr float min_scale = 1.0f / (fp8_max<fp8_t>() * 512.0f);
#pragma omp parallel for
for (int64_t m = 0; m < num_tokens; ++m) {
const scalar_t* row_in = in + m * in_row_stride;
fp8_t* row_out = out + m * out_row_stride;
float absmax = row_absmax(row_in, hidden_size);
if (scale_ub != nullptr) {
absmax = std::min(absmax, *scale_ub);
}
const float token_scale = std::max(absmax / fp8_max<fp8_t>(), min_scale);
scale[m] = token_scale;
quant_row<true>(row_out, row_in, hidden_size, token_scale);
}
}
#define DISPATCH_FP8_QUANT(OUT_DTYPE, IN_DTYPE, NAME, ...) \
AT_DISPATCH_FLOATING_TYPES_AND2( \
at::ScalarType::BFloat16, at::ScalarType::Half, IN_DTYPE, NAME, [&] { \
switch (OUT_DTYPE) { \
case at::ScalarType::Float8_e4m3fn: { \
using fp8_t = c10::Float8_e4m3fn; \
__VA_ARGS__(); \
break; \
} \
case at::ScalarType::Float8_e5m2: { \
using fp8_t = c10::Float8_e5m2; \
__VA_ARGS__(); \
break; \
} \
default: \
TORCH_CHECK(false, NAME, ": unsupported fp8 dtype ", OUT_DTYPE); \
} \
})
void check_fp8_quant_args(const torch::Tensor& out,
const torch::Tensor& input) {
TORCH_CHECK(input.dim() >= 2 && input.is_cpu() && out.is_cpu(),
"fp8 quant expects CPU tensors with at least 2 dims");
TORCH_CHECK(input.stride(-1) == 1 && out.stride(-1) == 1,
"last dimension of input and output must be contiguous");
TORCH_CHECK(out.size(-1) == input.size(-1) &&
out.numel() / out.size(-1) >= input.numel() / input.size(-1),
"output must match input hidden size and hold all tokens");
}
} // namespace
void static_scaled_fp8_quant(
torch::Tensor& out, const torch::Tensor& input, const torch::Tensor& scale,
const std::optional<std::vector<int64_t>>& opt_group_shape) {
check_fp8_quant_args(out, input);
TORCH_CHECK(scale.scalar_type() == at::kFloat, "scale must be float32");
const int64_t hidden_size = input.size(-1);
const int64_t num_tokens = input.numel() / hidden_size;
int64_t group_m = num_tokens;
int64_t group_n = hidden_size;
int64_t scale_stride_i = 0;
int64_t scale_stride_j = 0;
auto resolve_group_shape = [&]() {
TORCH_CHECK(opt_group_shape.has_value() && opt_group_shape->size() == 2,
"group_shape with exactly 2 elements is required");
group_m = (*opt_group_shape)[0] == -1 ? num_tokens : (*opt_group_shape)[0];
group_n = (*opt_group_shape)[1] == -1 ? hidden_size : (*opt_group_shape)[1];
};
if (scale.dim() == 0 || (scale.dim() == 1 && scale.numel() == 1)) {
// per-tensor
} else if (scale.dim() == 1) {
resolve_group_shape();
const int64_t scale_m = num_tokens / group_m;
const int64_t scale_n = hidden_size / group_n;
TORCH_CHECK(num_tokens % group_m == 0 && hidden_size % group_n == 0,
"input shape must be divisible by group_shape");
TORCH_CHECK(scale.numel() == scale_m * scale_n,
"1D scale length does not match group_shape");
TORCH_CHECK(scale_m == 1 || scale_n == 1,
"1D scale requires one trivial scale dimension; use a 2D "
"scale for group scaling");
if (scale_m == 1) {
scale_stride_j = scale.stride(0);
} else {
scale_stride_i = scale.stride(0);
}
} else if (scale.dim() == 2) {
TORCH_CHECK(
num_tokens % scale.size(0) == 0 && hidden_size % scale.size(1) == 0,
"input shape must be divisible by 2D scale shape");
group_m = num_tokens / scale.size(0);
group_n = hidden_size / scale.size(1);
scale_stride_i = scale.stride(0);
scale_stride_j = scale.stride(1);
} else {
TORCH_CHECK(false, "scale must be 0D, 1D, or 2D, got ", scale.dim(), "D");
}
DISPATCH_FP8_QUANT(
out.scalar_type(), input.scalar_type(), "static_scaled_fp8_quant", [&] {
static_quant_impl<scalar_t, fp8_t>(
out.data_ptr<fp8_t>(), input.const_data_ptr<scalar_t>(),
scale.const_data_ptr<float>(), num_tokens, hidden_size,
input.stride(-2), out.stride(-2), group_m, group_n, scale_stride_i,
scale_stride_j);
});
}
void dynamic_per_token_scaled_fp8_quant(
torch::Tensor& out, const torch::Tensor& input, torch::Tensor& scale,
const std::optional<torch::Tensor>& scale_ub) {
check_fp8_quant_args(out, input);
TORCH_CHECK(scale.scalar_type() == at::kFloat && scale.is_contiguous(),
"scale must be a contiguous float32 tensor");
const int64_t hidden_size = input.size(-1);
const int64_t num_tokens = input.numel() / hidden_size;
TORCH_CHECK(scale.numel() >= num_tokens,
"scale must hold one element per token");
const float* scale_ub_ptr = nullptr;
if (scale_ub.has_value()) {
TORCH_CHECK(scale_ub->scalar_type() == at::kFloat && scale_ub->numel() == 1,
"scale_ub must be a single float32 element");
scale_ub_ptr = scale_ub->const_data_ptr<float>();
}
DISPATCH_FP8_QUANT(out.scalar_type(), input.scalar_type(),
"dynamic_per_token_scaled_fp8_quant", [&] {
dynamic_per_token_quant_impl<scalar_t, fp8_t>(
out.data_ptr<fp8_t>(), scale.data_ptr<float>(),
input.const_data_ptr<scalar_t>(), scale_ub_ptr,
num_tokens, hidden_size, input.stride(-2),
out.stride(-2));
});
}
csrc/cpu/sgl-kernels/common.h
+10 -2
View File
@@ -2,8 +2,16 @@
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
//
// Synced from
// https://github.com/sgl-project/sglang/tree/7c248dde7fe1f3b5100966f8143f97a9932c22a4/sgl-kernel/csrc/cpu
// Sync date: 2026-07-29
// https://github.com/sgl-project/sglang/tree/c4870d98cfefa80f5cbf83de2ff6a6745032f86c/python/sglang/kernels/aot/csrc/cpu
// Sync date: 2026-10-16
//
// Partial sync: only moe.cpp, moe.h, moe_fp8.cpp, vec.h and gemm.h were merged
// (3-way, on top of the prior sync at
// 7c248dde7fe1f3b5100966f8143f97a9932c22a4). decode.cpp, extend.cpp,
// gemm_fp8.cpp, topk.cpp and vec_pack.h still match the prior sync: upstream
// changed them for fp8 packed KV-cache attention (new signatures that need
// op-schema plumbing) and a topk.cpp rewrite that conflicts with the local
// DeepSeek-V4 biased-topk. Reconcile them together in a full sync.
// clang-format off
csrc/cpu/sgl-kernels/gemm.h
+16 -3
View File
@@ -98,8 +98,18 @@ enum class CPUActMethod : int {
silu_and_mul = 0,
swiglu = 1,
gelu_and_mul = 2,
gelu_tanh_and_mul = 3,
};
inline CPUActMethod act_method_from_int(int64_t act_method) {
TORCH_CHECK(
act_method >= static_cast<int64_t>(CPUActMethod::silu_and_mul) &&
act_method <= static_cast<int64_t>(CPUActMethod::gelu_tanh_and_mul),
"Unsupported act_method: ",
act_method);
return static_cast<CPUActMethod>(act_method);
}
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 };
constexpr bool operator==(CPUQuantMethod a, int64_t b) {
@@ -169,7 +179,8 @@ void fused_experts_int8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
int64_t num_tokens_post_pad,
CPUActMethod act_func);
// moe implementations for fp8 w8a16 and mxfp4
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
@@ -253,7 +264,8 @@ void fused_experts_int4_w4a8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
int64_t num_tokens_post_pad,
CPUActMethod act_func);
template <typename scalar_t>
void shared_expert_fp8_kernel_impl(
@@ -452,4 +464,5 @@ void fused_experts_fp8_a8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
int64_t num_tokens_post_pad,
CPUActMethod act_func);
csrc/cpu/sgl-kernels/moe.cpp
+28 -9
View File
@@ -502,7 +502,8 @@ void fused_experts_kernel_impl(
} else {
const int64_t offset = offsets[mb];
if (act_func == CPUActMethod::swiglu) {
// the fused path below hardcodes SiLU in its store step
if (act_func != CPUActMethod::silu_and_mul) {
tinygemm_kernel(
/* A */ A,
/* B */ B0,
@@ -544,12 +545,16 @@ void fused_experts_kernel_impl(
add_bias_stub(C1 + m * BLOCK_N, B1_bias, n_size);
}
}
// 1.d silu and mul
// 1.d activation and mul
const int64_t offset = offsets[mb];
if (act_func == CPUActMethod::silu_and_mul && use_brgemm) {
for (int64_t m = 0; m < m_size; ++m) {
silu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N);
}
} else if (act_func == CPUActMethod::gelu_and_mul) {
for (int64_t m = 0; m < m_size; ++m) {
gelu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N);
}
} else if (act_func == CPUActMethod::swiglu) {
for (int64_t m = 0; m < m_size; ++m) {
scalar_t* __restrict__ ic1_row = ic1 + (offset + m) * N;
@@ -910,7 +915,20 @@ void fused_experts_cpu(
const std::optional<at::Tensor>& w2_bias,
const std::optional<double>& alpha,
const std::optional<double>& limit,
bool is_vnni) {
bool is_vnni,
int64_t act_method) {
const CPUActMethod act_func = act_method_from_int(act_method);
TORCH_CHECK(
act_func != CPUActMethod::swiglu || (alpha.has_value() && limit.has_value()),
"fused_experts_cpu: swiglu requires alpha and limit");
// the int8 and int4 kernels only support silu or gelu_tanh in their fused store step
const bool is_int_quant =
moe_comp_method == CPUQuantMethod::INT8_W8A8 || moe_comp_method == CPUQuantMethod::INT4_W4A8;
TORCH_CHECK(
!is_int_quant || act_func == CPUActMethod::silu_and_mul || act_func == CPUActMethod::gelu_tanh_and_mul,
"fused_experts_cpu: INT8_W8A8 and INT4_W4A8 support silu or gelu_tanh only, got act_method: ",
act_method);
auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -942,6 +960,7 @@ void fused_experts_cpu(
}
CHECK_DIM(2, topk_weights);
CHECK_DIM(2, topk_ids_);
CHECK_EQ(topk_ids_.scalar_type(), at::kInt);
// TODO: support topk_weights to be bf16 or fp16 in the kernel.
@@ -1097,7 +1116,8 @@ void fused_experts_cpu(
K,
E,
topk,
num_tokens_post_pad);
num_tokens_post_pad,
act_func);
} else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
// here we just ignore C_tmp as it is not used
scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
@@ -1105,7 +1125,6 @@ void fused_experts_cpu(
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
CHECK_MOE_SCALES_FP8(1, 2);
fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
@@ -1145,7 +1164,6 @@ void fused_experts_cpu(
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
// mxfp4 supports only group size of 32 (2^5)
constexpr int64_t group_size = 32;
@@ -1225,7 +1243,8 @@ void fused_experts_cpu(
K,
E,
topk,
num_tokens_post_pad);
num_tokens_post_pad,
act_func);
} else if (moe_comp_method == CPUQuantMethod::FP8_W8A8) {
// Stage 1: FP8×FP8 GEMM (hidden_states is FP8 input)
// Stage 2: BF16×FP8 GEMM (W8A16 path for down-projection)
@@ -1266,12 +1285,12 @@ void fused_experts_cpu(
K,
E,
topk,
num_tokens_post_pad);
num_tokens_post_pad,
act_func);
} else {
scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
fused_experts_kernel_impl<scalar_t>(
out_hidden_states.data_ptr<scalar_t>(),
csrc/cpu/sgl-kernels/moe.h
+40
View File
@@ -164,6 +164,46 @@ inline void silu_and_mul_stub(
}
}
template <typename scalar_t, typename input_t>
inline void gelu_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"gelu_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = fast_gelu(x0) * y0;
x1 = fast_gelu(x1) * y1;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
}
template <typename scalar_t, typename input_t>
inline void gelu_tanh_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"gelu_tanh_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = fast_gelu_tanh(x0) * y0;
x1 = fast_gelu_tanh(x1) * y1;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
}
template <typename scalar_t, typename input_t>
inline void clamp_sigmoid_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) {
csrc/cpu/sgl-kernels/moe_fp8.cpp
+13 -1
View File
@@ -124,13 +124,25 @@ void fused_experts_fp_kernel_impl(
}
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
// stage 1.5: intermediate_cache1 = activation(intermediate_cache0)
if (act_func == CPUActMethod::silu_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else if (act_func == CPUActMethod::gelu_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
gelu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else if (act_func == CPUActMethod::gelu_tanh_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
gelu_tanh_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else if (act_func == CPUActMethod::swiglu) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp
+18 -8
View File
@@ -172,7 +172,8 @@ void fused_experts_fp8_a8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
int64_t num_tokens_post_pad,
CPUActMethod act_func) {
// 1. intermediate_cache1 : [M * topk, N]
// 2. intermediate_cache2 : [M * topk, K]
// 3. A_tmp : [T, BLOCK_M * K] FP8
@@ -316,12 +317,20 @@ void fused_experts_fp8_a8_kernel_impl(
}
});
// stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:]
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
// stage 1.5: intermediate_cache1 = activation(ic0[:, :N]) * ic0[:, N:]
if (act_func == CPUActMethod::gelu_tanh_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
gelu_tanh_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
}
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2: [E, K, N] FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight)
@@ -426,7 +435,8 @@ void fused_experts_fp8_a8_kernel_impl(
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
int64_t num_tokens_post_pad, \
CPUActMethod act_func)
INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/moe_int4.cpp
+18 -8
View File
@@ -52,7 +52,8 @@ void fused_experts_int4_w4a8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
int64_t num_tokens_post_pad,
CPUActMethod act_func) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
int num_threads = at::get_num_threads();
@@ -188,12 +189,20 @@ void fused_experts_int4_w4a8_kernel_impl(
}
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
// stage 1.5: intermediate_cache1 = activation(intermediate_cache0)
if (act_func == CPUActMethod::gelu_tanh_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
gelu_tanh_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
}
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
@@ -317,7 +326,8 @@ void fused_experts_int4_w4a8_kernel_impl(
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
int64_t num_tokens_post_pad, \
CPUActMethod act_func)
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/moe_int8.cpp
+59 -26
View File
@@ -9,7 +9,7 @@
namespace {
template <typename scalar_t, int BLOCK_N>
template <typename scalar_t, int BLOCK_N, CPUActMethod ACT = CPUActMethod::silu_and_mul>
inline void silu_and_mul(
scalar_t* __restrict__ C,
const int32_t* __restrict__ C0, // x: x0, x1
@@ -56,7 +56,11 @@ inline void silu_and_mul(
auto silu_and_mul = [&](auto col) {
__m512 x = vc0[col];
__m512 y = vc1[col];
vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y);
if constexpr (ACT == CPUActMethod::gelu_tanh_and_mul) {
vc0[col] = _mm512_mul_ps(_mm512_rcp14_gelu_tanh_ps(x), y);
} else {
vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y);
}
};
using bVec = at::vec::Vectorized<scalar_t>;
@@ -122,7 +126,7 @@ inline void scale_C(
}
/// gemm for w13
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
template <typename scalar_t, int BLOCK_M, int BLOCK_N, CPUActMethod ACT = CPUActMethod::silu_and_mul>
struct tinygemm_kernel_vnni {
static inline void apply(
const uint8_t* __restrict__ A,
@@ -218,8 +222,13 @@ struct tinygemm_kernel_vnni {
__m512 x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
__m512 y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
__m512 y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
if constexpr (ACT == CPUActMethod::gelu_tanh_and_mul) {
x0 = _mm512_mul_ps(_mm512_rcp14_gelu_tanh_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_gelu_tanh_ps(x1), y1);
} else {
x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
}
bVec out_vec = convert_from_float_ext<scalar_t>(fVec(x0), fVec(x1));
out_vec.store(C + row * ldc + col * 16);
@@ -232,8 +241,8 @@ struct tinygemm_kernel_vnni {
}
};
#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply( \
#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE, ACT>::apply( \
A + mb_start * lda, \
B0 + nb_start * 4, \
B1 + nb_start * 4, \
@@ -248,7 +257,7 @@ struct tinygemm_kernel_vnni {
ldb, \
ldc);
template <typename scalar_t>
template <typename scalar_t, CPUActMethod ACT = CPUActMethod::silu_and_mul>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B0,
@@ -464,7 +473,8 @@ void fused_experts_int8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
int64_t num_tokens_post_pad,
CPUActMethod act_func) {
// handle 2 tiles per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
@@ -562,25 +572,47 @@ void fused_experts_int8_kernel_impl(
// 1.d silu and mul
const int64_t offset = offsets[mb];
silu_and_mul<scalar_t, BLOCK_N>(
ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
if (act_func == CPUActMethod::gelu_tanh_and_mul) {
silu_and_mul<scalar_t, BLOCK_N, CPUActMethod::gelu_tanh_and_mul>(
ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
} else {
silu_and_mul<scalar_t, BLOCK_N, CPUActMethod::silu_and_mul>(
ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
}
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
const int64_t offset = offsets[mb];
tinygemm_kernel(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* As */ As,
/* Bs0 */ Bs0,
/* Bs1 */ Bs1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
if (act_func == CPUActMethod::gelu_tanh_and_mul) {
tinygemm_kernel<scalar_t, CPUActMethod::gelu_tanh_and_mul>(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* As */ As,
/* Bs0 */ Bs0,
/* Bs1 */ Bs1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
} else {
tinygemm_kernel<scalar_t, CPUActMethod::silu_and_mul>(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* As */ As,
/* Bs0 */ Bs0,
/* Bs1 */ Bs1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
}
}
});
@@ -711,7 +743,8 @@ void fused_experts_int8_kernel_impl(
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
int64_t num_tokens_post_pad, \
CPUActMethod act_func)
INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/vec.h
+29
View File
@@ -612,6 +612,15 @@ inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa));
}
// gelu_tanh(x) = x * sigmoid(2z), z = sqrt(2/pi) * (x + 0.044715 * x^3)
inline __attribute__((always_inline)) __m512 _mm512_rcp14_gelu_tanh_ps(__m512 x) {
const __m512 c0 = _mm512_set1_ps(1.5957691216057308f); // 2 * sqrt(2/pi)
const __m512 c1 = _mm512_set1_ps(0.07135481627260025f); // 2 * sqrt(2/pi) * 0.044715
__m512 x2 = _mm512_mul_ps(x, x);
__m512 z2 = _mm512_fmadd_ps(_mm512_mul_ps(c1, x2), x, _mm512_mul_ps(c0, x)); // c0*x + c1*x^3
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(z2));
}
#endif
inline at::vec::Vectorized<float> fast_sigmoid(const at::vec::Vectorized<float>& x) {
@@ -631,6 +640,26 @@ inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x)
#endif
}
// exact (erf) gelu, matching torch.nn.functional.gelu(approximate="none")
inline at::vec::Vectorized<float> fast_gelu(const at::vec::Vectorized<float>& x) {
const auto half = at::vec::Vectorized<float>(0.5f);
const auto one = at::vec::Vectorized<float>(1.f);
const auto inv_sqrt2 = at::vec::Vectorized<float>(0.70710678118654752440f);
return half * x * (one + (x * inv_sqrt2).erf());
}
// tanh-approx gelu: x * sigmoid(2z), z = sqrt(2/pi) * (x + 0.044715 * x^3)
inline at::vec::Vectorized<float> fast_gelu_tanh(const at::vec::Vectorized<float>& x) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_gelu_tanh_ps(x));
#else
const auto c0 = at::vec::Vectorized<float>(1.5957691216057308f); // 2 * sqrt(2/pi)
const auto c1 = at::vec::Vectorized<float>(0.07135481627260025f); // 2 * sqrt(2/pi) * 0.044715
const auto z2 = c0 * x + c1 * x * x * x;
return x * fast_sigmoid(z2);
#endif
}
inline at::vec::Vectorized<float>
fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) {
#if defined(CPU_CAPABILITY_AVX512)
csrc/cpu/torch_bindings.cpp
+19 -4
View File
@@ -78,7 +78,8 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states,
const std::optional<at::Tensor>& w1_bias,
const std::optional<at::Tensor>& w2_bias,
const std::optional<double>& alpha,
const std::optional<double>& limit, bool is_vnni);
const std::optional<double>& limit, bool is_vnni,
int64_t act_method);
at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
at::Tensor& scales2,
@@ -398,6 +399,13 @@ torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits,
torch::Tensor greedy_argmax(const torch::Tensor& logits);
void static_scaled_fp8_quant(
torch::Tensor& out, const torch::Tensor& input, const torch::Tensor& scale,
const std::optional<std::vector<int64_t>>& group_shape);
void dynamic_per_token_scaled_fp8_quant(
torch::Tensor& out, const torch::Tensor& input, torch::Tensor& scale,
const std::optional<torch::Tensor>& scale_ub);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
@@ -659,7 +667,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
"Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, "
"Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
"bool is_vnni) -> "
"bool is_vnni, int act_method) -> "
"()");
ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
ops.def(
@@ -906,9 +914,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
&cpu_attention_with_kv_cache);
// placeholders
ops.def("static_scaled_fp8_quant() -> ()", placeholder_op);
ops.def(
"static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
"int[]? group_shape=None) -> ()");
ops.impl("static_scaled_fp8_quant", torch::kCPU, &static_scaled_fp8_quant);
ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op);
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
ops.def(
"dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, "
"Tensor! scale, Tensor? scale_ub) -> ()");
ops.impl("dynamic_per_token_scaled_fp8_quant", torch::kCPU,
&dynamic_per_token_scaled_fp8_quant);
// WNA16
#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) || \
csrc/fs_io.cpp
+5 -1
View File
@@ -71,7 +71,11 @@ inline int _store_block(const char* tmp_path, const char* dest_path,
if (rename(tmp_path, dest_path) != 0) {
const int err = errno;
unlink(tmp_path);
return err;
// Another writer published this block after the access() above; some
// shared filesystems fail rename(2) with EEXIST instead of replacing.
// File names are content hashes, so the file already in place holds this
// block and the store is done.
return err == EEXIST ? 0 : err;
}
return 0;
csrc/hisparse_kernels.cu
+201 -156
View File
@@ -21,6 +21,7 @@
#include <algorithm>
#include <cstdint>
#include <optional>
#include <vector>
#include <torch/csrc/stable/c/shim.h>
#include <torch/csrc/stable/ops.h>
@@ -185,8 +186,8 @@ constexpr int32_t kEntryMiss = 1 << 29;
constexpr int32_t kEntryInvalid = 1 << 28;
constexpr int32_t kEntrySlotMask = 0xFFFF;
constexpr int kResidencyCounters = 5;
// Per-request table entries per referenced top-k position (load <= 1/4
// until the hash_size cap).
// Per-request table entries per referenced top-k position, within the
// device's shared-memory budget.
constexpr int kHashLoadInverse = 4;
// Translate a request-relative top-k position to its resident GPU row (-1 when
@@ -236,6 +237,12 @@ __device__ __forceinline__ int32_t translate_topk_entry(
return g >= host_rows ? -1 : g;
}
// Phase 1 parks host rows in hot_indices as values below -1, distinct from
// resident rows (>= 0) and invalid positions (-1).
__device__ __forceinline__ int32_t encode_host_row(const int32_t g) {
return -2 - g;
}
// Write the hot index of one reference to a resolved union entry and tally it.
__device__ __forceinline__ void write_entry_index(
const int32_t entry, const int row, const int i, const int32_t top_k,
@@ -285,9 +292,11 @@ __device__ __forceinline__ void write_entry_index(
// s_hash_vals[hash_size] first position, then slot | kEntry* tag
// s_chunk_off[nbc + 1] prefix sums for hit, then miss, compaction
// s_evict_off[nbc + 1] prefix sums for evictable compaction
// s_counters[5] distinct hits, union size, per-row valid count,
// misses, per-entry hits
// s_done[nbc] positions phase 3 resolved, one word per chunk
// s_counters[5] distinct hits, union size, per-entry misses,
// host positions (phase 1) then misses,
// per-entry hits
// s_done[nbc] hot slots holding a union row (phase 2), then
// positions phase 3 resolved, one word per chunk
// s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd]
// Valid global ids must be unique within each row.
__global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
@@ -312,13 +321,14 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
const int32_t request_state_count, const int64_t host_rows,
const int64_t hot_table_stride, const int32_t hot_block_size,
const int32_t top_k, const int32_t hot_size, const int32_t hash_size,
const int64_t region_stride, const int64_t attention_block_stride,
const int64_t source_bt_stride, const int32_t source_num_reqs,
const int32_t source_num_blocks, const int32_t source_block_size,
const int64_t resident_bt_stride, const int32_t resident_num_reqs,
const int32_t resident_num_blocks, const int32_t resident_block_size,
const int32_t resident_null_block, const int64_t input_row_stride,
const int64_t attention_row_stride, const int64_t valid_count_stride) {
const int32_t max_union_rows, const int64_t region_stride,
const int64_t attention_block_stride, const int64_t source_bt_stride,
const int32_t source_num_reqs, const int32_t source_num_blocks,
const int32_t source_block_size, const int64_t resident_bt_stride,
const int32_t resident_num_reqs, const int32_t resident_num_blocks,
const int32_t resident_block_size, const int32_t resident_null_block,
const int64_t input_row_stride, const int64_t attention_row_stride,
const int64_t valid_count_stride) {
const int NUM_WARPS = blockDim.x / kWarpSize;
const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
@@ -377,12 +387,6 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
device_global_indices + static_cast<int64_t>(state_row) * region_stride;
int16_t* row_lru = lru_slots + static_cast<int64_t>(state_row) * hot_size;
// A union within max_union_rows (hash_size - 1) and within the rows' top-k
// leaves an empty slot, so every probe sequence ends. Use a slice sized to
// this request's rows so single-row requests initialize a small table.
const int table_size =
min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k);
extern __shared__ char smem_raw[];
int32_t* s_hash_keys = reinterpret_cast<int32_t*>(smem_raw);
int32_t* s_hash_vals = s_hash_keys + hash_size;
@@ -393,10 +397,6 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
reinterpret_cast<unsigned int*>(s_counters + kResidencyCounters);
int16_t* s_lru_out = reinterpret_cast<int16_t*>(s_done + num_buffer_chunks);
for (int i = tid; i < table_size; i += blockDim.x) {
s_hash_keys[i] = kHashEmpty;
s_hash_vals[i] = INT32_MAX;
}
for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) {
s_chunk_off[i] = 0;
s_evict_off[i] = 0;
@@ -406,129 +406,174 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
}
__syncthreads();
// Phase 1: translate every row's request-relative positions, resolve
// resident rows directly, and collect the union of host rows.
for (int row = first_row; row < end_row; ++row) {
const int32_t* row_topk =
global_indices + static_cast<int64_t>(row) * input_row_stride;
// Phase 1: translate every row's request-relative positions in one pass,
// resolve resident rows directly, and count the host-backed positions.
// s_chunk_off holds per-row valid counts until phase 2.
int32_t* s_row_valid = s_chunk_off;
const int num_positions = (end_row - first_row) * top_k;
for (int position = tid; position < num_positions; position += blockDim.x) {
const int row = first_row + position / top_k;
const int i = position % top_k;
int32_t* row_out = hot_indices + static_cast<int64_t>(row) * top_k;
int32_t* row_attention =
attention_indices != nullptr
? attention_indices +
static_cast<int64_t>(row) * attention_row_stride
: nullptr;
int32_t* row_miss = miss_mask != nullptr
? miss_mask + static_cast<int64_t>(row) * top_k
: nullptr;
for (int i = tid; i < top_k; i += blockDim.x) {
int32_t resident_row;
const int32_t g = translate_topk_entry(
row_topk[i], request_row, state_row, source_block_table,
resident_block_table, host_rows, source_bt_stride, source_num_reqs,
source_num_blocks, source_block_size, resident_bt_stride,
resident_num_reqs, resident_num_blocks, resident_block_size,
resident_null_block, resident_row);
if (resolved_global_indices != nullptr) {
resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
}
if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1);
if (row_miss != nullptr) row_miss[i] = 0;
if (resident_row >= 0) {
store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
attention_block_stride);
} else if (g < 0) {
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
} else {
int h = hash_slot(g, table_size);
bool inserted = false;
for (int probe = 0; probe < table_size; ++probe) {
const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g);
if (old == kHashEmpty) {
atomicAdd(&s_counters[1], 1);
inserted = true;
break;
}
if (old == g) {
inserted = true;
break;
}
h = next_slot(h, table_size);
}
if (inserted) {
atomicMin(&s_hash_vals[h], (row - first_row) * top_k + i);
} else {
// A full table: push the union count past the cap so the check
// below traps.
atomicAdd(&s_counters[1], hash_size);
}
}
int32_t resident_row;
const int32_t g = translate_topk_entry(
global_indices[static_cast<int64_t>(row) * input_row_stride + i],
request_row, state_row, source_block_table, resident_block_table,
host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resolved_global_indices != nullptr) {
resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
}
__syncthreads();
if (tid == 0) {
if (valid_counts != nullptr) {
valid_counts[static_cast<int64_t>(row) * valid_count_stride] =
s_counters[2];
}
s_counters[2] = 0;
if (resident_row >= 0 || g >= 0)
atomicAdd(&s_row_valid[row - first_row], 1);
if (miss_mask != nullptr)
miss_mask[static_cast<int64_t>(row) * top_k + i] = 0;
if (resident_row >= 0) {
store_hot_index(row_out, row_attention, i, resident_row, hot_block_size,
attention_block_stride);
} else if (g < 0) {
store_hot_index(row_out, row_attention, i, -1, hot_block_size,
attention_block_stride);
} else {
// Stash the host row in its output slot; later phases overwrite it.
row_out[i] = encode_host_row(g);
atomicAdd(&s_counters[3], 1);
}
__syncthreads();
}
// The table and eviction slots are sized for max_union_rows (hash_size - 1);
// a larger union would never reach an empty probe slot or run past s_lru_out.
if (s_counters[1] >= hash_size) {
if (tid == 0) {
printf(
"HiSparse residency: request %d resolves %d host rows, above the "
"configured bound of %d\n",
request_row, s_counters[1], hash_size - 1);
__syncthreads();
for (int row = first_row + tid; row < end_row; row += blockDim.x) {
if (valid_counts != nullptr) {
valid_counts[static_cast<int64_t>(row) * valid_count_stride] =
s_row_valid[row - first_row];
}
__trap();
s_row_valid[row - first_row] = 0;
}
// Fully resident rows need only request-relative page translation. Avoid
// scanning or rewriting the hot LRU when no selected row can consult it.
if (s_counters[1] == 0) {
// building the table or touching the hot LRU when no row can consult it.
const int host_positions = s_counters[3];
if (host_positions == 0) {
for (int row = first_row + tid; row < end_row; row += blockDim.x) {
if (swap_counts != nullptr) swap_counts[row] = 0;
}
return;
}
// Phase 2: walk hot slots in LRU order, classify hit / evictable, and
// compact them (hits forward, evictables backward) into s_lru_out. A hit
// records its slot on the union entry.
// Use a slice sized to this request's rows so single-row requests
// initialize a small table. Legal unions always leave an empty slot.
const int table_size =
min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k);
for (int i = tid; i < table_size; i += blockDim.x) {
s_hash_keys[i] = kHashEmpty;
s_hash_vals[i] = INT32_MAX;
}
__syncthreads();
// Phase 1b: collect the union of host rows.
for (int position = tid; position < num_positions; position += blockDim.x) {
const int32_t stashed =
hot_indices[static_cast<int64_t>(first_row) * top_k + position];
if (stashed > -2) continue;
const int32_t g = encode_host_row(stashed);
int h = hash_slot(g, table_size);
bool inserted = false;
for (int probe = 0; probe < table_size; ++probe) {
const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g);
if (old == kHashEmpty) {
atomicAdd(&s_counters[1], 1);
inserted = true;
break;
}
if (old == g) {
inserted = true;
break;
}
h = next_slot(h, table_size);
}
if (inserted) {
atomicMin(&s_hash_vals[h], position);
} else {
// A full table: push the union count past the cap so the check
// below traps.
atomicAdd(&s_counters[1], max_union_rows + 1);
}
}
__syncthreads();
// Hash capacity can exceed the union bound; eviction slots cannot.
if (s_counters[1] > max_union_rows) {
if (tid == 0) {
printf(
"HiSparse residency: request %d resolves %d host rows, above the "
"configured bound of %d\n",
request_row, s_counters[1], max_union_rows);
}
__trap();
}
// Phase 2: scan hot slots in slot order so device_global_indices is read
// coalesced, recording each hit's slot on its union entry and in s_done.
constexpr int kSlotBatch = 4;
for (int base = 0; base < hot_size; base += kSlotBatch * blockDim.x) {
int32_t cached_g[kSlotBatch];
#pragma unroll
for (int b = 0; b < kSlotBatch; ++b) {
const int slot = base + b * blockDim.x + tid;
cached_g[b] = slot < hot_size ? row_dgi[slot] : -1;
}
#pragma unroll
for (int b = 0; b < kSlotBatch; ++b) {
const int slot = base + b * blockDim.x + tid;
// Every lane makes its first probe, then the warp iterates together
// until no lane is still probing, so the warp is converged at the
// ballot. CUDA 12.8-13.2.1 ptxas can fail to reconverge a per-lane
// probe loop before a ballot on sm_100.
const int32_t g = cached_g[b];
int h = hash_slot(g, table_size);
int32_t k = s_hash_keys[h];
bool is_hit = g >= 0 && k == g;
bool probing = g >= 0 && !is_hit && k != kHashEmpty;
while (__any_sync(0xFFFFFFFF, probing)) {
if (probing) h = next_slot(h, table_size);
k = probing ? s_hash_keys[h] : kHashEmpty;
const bool found = probing && k == g;
is_hit |= found;
probing = probing && !found && k != kHashEmpty;
}
if (is_hit) s_hash_vals[h] = kEntryHit | slot;
const unsigned int hit_bits = __ballot_sync(0xFFFFFFFF, is_hit);
if (lane_id == 0 && slot < hot_size) {
s_done[slot / kWarpSize] = hit_bits;
}
}
}
__syncthreads();
// Walk hot slots in LRU order, classify hit / evictable, and compact them
// (hits forward, evictables backward) into s_lru_out.
const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS;
int total_hit_count = 0;
int total_evict_count = 0;
int16_t next_lru = tid < hot_size ? row_lru[tid] : int16_t(-1);
for (int iter = 0; iter < iters_buffer; iter++) {
const int chunk_idx = warp_id + iter * NUM_WARPS;
const bool has_valid_chunk = chunk_idx < num_buffer_chunks;
const int pos = chunk_idx * kWarpSize + lane_id;
const bool has_valid_pos = has_valid_chunk && (pos < hot_size);
const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1);
// Corruption tripwire: lru/dgi are long-lived device state; if some
// external writer (e.g. stray RDMA into reused VRAM) corrupts a slot
// out of [0, region_stride), treat it as no-hit so it degrades to a
// re-miss instead of an unbounded dgi read (phase 3 bounds the write
// side).
const int32_t cached_g =
(slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1;
bool is_hit = false;
if (cached_g >= 0) {
int h = hash_slot(cached_g, table_size);
while (true) {
const int32_t k = s_hash_keys[h];
if (k == cached_g) {
s_hash_vals[h] = kEntryHit | slot;
is_hit = true;
break;
}
if (k == kHashEmpty) break;
h = next_slot(h, table_size);
}
}
const int16_t slot = has_valid_pos ? next_lru : int16_t(-1);
const int next_pos = pos + NUM_WARPS * kWarpSize;
next_lru = next_pos < hot_size ? row_lru[next_pos] : int16_t(-1);
// Corruption tripwire: lru is long-lived device state; an out-of-range
// slot is treated as no-hit so it degrades to a re-miss (phase 3 bounds
// the write side).
const bool is_hit = slot >= 0 && slot < hot_size &&
(s_done[slot / kWarpSize] >> (slot % kWarpSize)) & 1u;
const bool is_evictable = has_valid_pos && !is_hit;
int local_hit_off = 0;
@@ -580,7 +625,6 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
// device_global_indices. Swaps fill the rows' swap segments in order; the
// union never exceeds their combined capacity, and position chunks never
// exceed buffer chunks because hot_size covers every row's top-k.
const int num_positions = (end_row - first_row) * top_k;
const int num_position_chunks = (num_positions + kWarpSize - 1) / kWarpSize;
const int iters_position = (num_position_chunks + NUM_WARPS - 1) / NUM_WARPS;
int miss_running_total = 0;
@@ -596,17 +640,10 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
int32_t entry = 0;
bool is_miss = false;
if (has_valid_chunk && position < num_positions) {
const int row = first_row + position / top_k;
int32_t resident_row;
g = translate_topk_entry(
global_indices[static_cast<int64_t>(row) * input_row_stride +
position % top_k],
request_row, state_row, source_block_table, resident_block_table,
host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resident_row < 0 && g >= 0) {
const int32_t stashed =
hot_indices[static_cast<int64_t>(first_row) * top_k + position];
if (stashed < -1) {
g = encode_host_row(stashed);
h = hash_slot(g, table_size);
while (s_hash_keys[h] != g) {
h = next_slot(h, table_size);
@@ -703,16 +740,10 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
continue;
}
const int row = first_row + position / top_k;
int32_t resident_row;
const int32_t g = translate_topk_entry(
global_indices[static_cast<int64_t>(row) * input_row_stride +
position % top_k],
request_row, state_row, source_block_table, resident_block_table,
host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
source_block_size, resident_bt_stride, resident_num_reqs,
resident_num_blocks, resident_block_size, resident_null_block,
resident_row);
if (resident_row >= 0 || g < 0) continue;
const int32_t stashed =
hot_indices[static_cast<int64_t>(first_row) * top_k + position];
if (stashed > -2) continue;
const int32_t g = encode_host_row(stashed);
int h = hash_slot(g, table_size);
while (s_hash_keys[h] != g) {
h = next_slot(h, table_size);
@@ -1213,13 +1244,12 @@ void hisparse_resolve_residency(
STD_TORCH_CHECK(max_union_rows > 0 && max_union_rows <= hot_size,
"max_union_rows must be in [1, hot_size], got ",
max_union_rows);
// One spare entry above the union bound keeps every probe sequence finite.
const int hash_size = static_cast<int>(max_union_rows) + 1;
const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize;
const size_t smem_bytes =
sizeof(int32_t) * (2 * hash_size + 2 * (num_buffer_chunks + 1) +
kResidencyCounters + num_buffer_chunks) +
const size_t fixed_smem_bytes =
sizeof(int32_t) * (3 * num_buffer_chunks + 2 + kResidencyCounters) +
sizeof(int16_t) * hot_size;
const size_t min_smem_bytes =
fixed_smem_bytes + 2 * sizeof(int32_t) * (max_union_rows + 1);
const torch::stable::accelerator::DeviceGuard device_guard(
hot_cache.get_device_index());
@@ -1228,14 +1258,21 @@ void hisparse_resolve_residency(
&max_smem_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin,
hot_cache.get_device_index()) == cudaSuccess,
"failed to query shared memory per block");
STD_TORCH_CHECK(smem_bytes <= static_cast<size_t>(max_smem_bytes),
STD_TORCH_CHECK(min_smem_bytes <= static_cast<size_t>(max_smem_bytes),
"HiSparse device_buffer_size=", hot_size, " needs ",
smem_bytes,
min_smem_bytes,
" bytes of shared memory to resolve residency, "
"but the device allows ",
max_smem_bytes,
"; lower device_buffer_size or the number of speculative "
"tokens.");
// Keep the table sparse when possible without rejecting configurations
// that fit the minimum table with one spare entry above the union bound.
const int hash_size =
std::min(kHashLoadInverse * static_cast<int>(max_union_rows) + 1,
static_cast<int>((max_smem_bytes - fixed_smem_bytes) /
(2 * sizeof(int32_t))));
const size_t smem_bytes = fixed_smem_bytes + 2 * sizeof(int32_t) * hash_size;
const cudaStream_t stream = get_current_cuda_stream();
const int64_t attention_row_stride =
attention_indices.has_value() ? attention_indices.value().stride(0) : 0;
@@ -1243,11 +1280,19 @@ void hisparse_resolve_residency(
valid_counts.has_value() ? valid_counts.value().stride(0) : 0;
auto kernel = hisparse_resolve_residency_kernel;
if (smem_bytes > 48 * 1024) {
const cudaError_t attribute_error = cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
STD_TORCH_CHECK(attribute_error == cudaSuccess,
"failed to configure HiSparse swap-in shared memory: ",
cudaGetErrorString(attribute_error));
static std::vector<size_t> configured_smem_bytes;
const size_t device_index = hot_cache.get_device_index();
if (configured_smem_bytes.size() <= device_index) {
configured_smem_bytes.resize(device_index + 1, 0);
}
if (configured_smem_bytes[device_index] < smem_bytes) {
const cudaError_t attribute_error = cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
STD_TORCH_CHECK(attribute_error == cudaSuccess,
"failed to configure HiSparse swap-in shared memory: ",
cudaGetErrorString(attribute_error));
configured_smem_bytes[device_index] = smem_bytes;
}
}
kernel<<<launch_rows, kBlockSize, smem_bytes, stream>>>(
hot_block_table.const_data_ptr<int32_t>(),
@@ -1260,10 +1305,10 @@ void hisparse_resolve_residency(
device_global_indices.mutable_data_ptr<int32_t>(),
lru_slots.mutable_data_ptr<int16_t>(), stats_ptr, request_state_ptr,
request_state_count, host_rows, hot_block_table.stride(0), hot_block_size,
top_k, hot_size, hash_size, region_stride, attention_block_stride,
source_bt_stride, source_num_reqs, source_num_blocks,
static_cast<int32_t>(source_block_size), resident_bt_stride,
resident_num_reqs, resident_num_blocks,
top_k, hot_size, hash_size, static_cast<int32_t>(max_union_rows),
region_stride, attention_block_stride, source_bt_stride, source_num_reqs,
source_num_blocks, static_cast<int32_t>(source_block_size),
resident_bt_stride, resident_num_reqs, resident_num_blocks,
static_cast<int32_t>(resident_block_size),
static_cast<int32_t>(resident_null_block), global_indices.stride(0),
attention_row_stride, valid_count_stride);
csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu
+73 -42
View File
@@ -48,6 +48,7 @@
using bf16_t = __bf16;
using bf16x2 = __bf16 __attribute__((ext_vector_type(2)));
using bf16x4 = __bf16 __attribute__((ext_vector_type(4)));
using bf16x8 = __bf16 __attribute__((ext_vector_type(8)));
using f32x4 = float __attribute__((ext_vector_type(4)));
using u32x2 = uint32_t __attribute__((ext_vector_type(2)));
@@ -120,11 +121,12 @@ struct Params {
// separate pointers so that a paged cache (both equal) and a dense
// [N, H, V, K] pair (two tensors, identity rows) share one code path.
// Either pointer may be null.
const float* cache_in;
float* cache_out;
const void* cache_in;
void* cache_out;
bool cache_bf16;
const int* sidx;
const bool* hh0;
// Floats between consecutive rows: the tensor's stride(0), which a paged
// Elements between consecutive rows: the tensor's stride(0), which a paged
// cache may pad beyond H * kV * kK. Only each row itself is dense.
int64_t crow;
int64_t ckrow; // the same, for the checkpoint destination
@@ -151,7 +153,8 @@ struct Params {
// replaying the tail as a second model forward. `ckidx`, when present,
// redirects sequence n's snapshot to row `ckidx[n]`, which lets the store go
// straight into the paged state cache; a negative row disables it.
float* ckpt;
void* ckpt;
bool ckpt_bf16;
const int* ckoff;
const int* ckidx;
int G;
@@ -163,6 +166,17 @@ struct Params {
// with a leading dimension of 0, so nothing inside the chunk loop changes.
[[maybe_unused]] __device__ bf16_t g_uzero[kK];
__device__ __forceinline__ void store_state4(void* dst, int64_t offset,
f32x4 value, bool bf16) {
if (bf16) {
const bf16x4 packed = __builtin_convertvector(value, bf16x4);
__builtin_memcpy(static_cast<bf16_t*>(dst) + offset, &packed,
sizeof(packed));
} else {
*reinterpret_cast<f32x4*>(static_cast<float*>(dst) + offset) = value;
}
}
union Frag {
bf16x8 v;
uint32_t d[4];
@@ -758,13 +772,22 @@ __global__ __launch_bounds__(64 * WAVES * KS,
// state the scan composed for it; in pass one they start at zero, which is
// the B_g the scan expects. Transfer blocks start at the identity.
if (HAS_H0 && g == 0 && warm && !from_sin) {
const float* s0 =
p.cache_in + h0_row +
((static_cast<int64_t>(h) * kV + v0 + lrow) * kK + 4 * lgrp);
const int64_t base =
h0_row + ((static_cast<int64_t>(h) * kV + v0 + lrow) * kK + 4 * lgrp);
for (int vt = 0; vt < kNV; ++vt)
for (int kt = 0; kt < kNKTL; ++kt)
S[kt][vt] = *reinterpret_cast<const f32x4*>(s0 + vt * 16 * kK +
(kt0 + kt) * 16);
for (int kt = 0; kt < kNKTL; ++kt) {
const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16;
if (p.cache_bf16) {
bf16x4 packed;
__builtin_memcpy(&packed,
static_cast<const bf16_t*>(p.cache_in) + offset,
sizeof(packed));
S[kt][vt] = __builtin_convertvector(packed, f32x4);
} else {
S[kt][vt] = *reinterpret_cast<const f32x4*>(
static_cast<const float*>(p.cache_in) + offset);
}
}
} else if (PASS2 && (g > 0 || (from_sin && warm))) {
const float* s0 = p.sin_ +
(static_cast<int64_t>(g) * p.NH() + nh) *
@@ -842,11 +865,12 @@ __global__ __launch_bounds__(64 * WAVES * KS,
// Same store as the final state below, at an interior boundary instead of
// the last one.
auto store_checkpoint = [&]() {
float* dst = p.ckpt + ck_dst + (v0 + lrow) * kK + 4 * lgrp;
const int64_t base = ck_dst + (v0 + lrow) * kK + 4 * lgrp;
for (int vt = 0; vt < kNV; ++vt)
for (int kt = 0; kt < kNKTL; ++kt)
*reinterpret_cast<f32x4*>(dst + vt * 16 * kK + (kt0 + kt) * 16) =
S[kt][vt];
for (int kt = 0; kt < kNKTL; ++kt) {
const int64_t offset = base + vt * 16 * kK + (kt0 + kt) * 16;
store_state4(p.ckpt, offset, S[kt][vt], p.ckpt_bf16);
}
};
const int tail = last_group ? T - (nt_all - 1) * kBT : kBT;
@@ -1049,15 +1073,18 @@ __global__ __launch_bounds__(64 * WAVES * KS,
// With G == 1 this is the final state. With G > 1 it is what the group
// reaches from a zero start, which the scan then composes.
float* s1 = nullptr;
int64_t st_vt = 16 * kK, st_kt = 16; // the cache row is [H][V][K] fp32
if constexpr (PASS2) {
// Only the last group holds the sequence's final state. The row's inner
// layout is [V][K], so this is one f32x4 per lane per tile, strictly
// inside the row this block owns.
if (last_group && p.cache_out != nullptr)
s1 = p.cache_out + h0_row +
((static_cast<int64_t>(h) * kV + v0 + lrow) * kK + 4 * lgrp);
// Only the last group writes the persistent state; group workspaces
// retain the accumulator precision regardless of the cache dtype.
if (last_group && p.cache_out != nullptr) {
const int64_t base =
h0_row + ((static_cast<int64_t>(h) * kV + v0 + lrow) * kK + 4 * lgrp);
for (int vt = 0; vt < kNV; ++vt)
for (int kt = 0; kt < kNKTL; ++kt)
store_state4(p.cache_out, base + vt * 16 * kK + (kt0 + kt) * 16,
S[kt][vt], p.cache_bf16);
}
return;
} else if (xfer) {
// Publish C_g = M_g - diag(A_g). A_g is broadcast through the
// double-buffer half the last chunk did not use, so the diagonal comes off
@@ -1086,18 +1113,15 @@ __global__ __launch_bounds__(64 * WAVES * KS,
} else {
if (slice == 0 && tid < kK)
p.ag[(static_cast<int64_t>(g) * p.NH() + nh) * kK + tid] = aacc;
s1 = p.bg +
(static_cast<int64_t>(g) * p.NH() + nh) *
(static_cast<int64_t>(kV) * kK) +
static_cast<int64_t>(v0 / 16) * kStTile + lane * 4;
st_vt = kStTile;
st_kt = kStKt;
}
if (s1 != nullptr)
float* const s1 = p.bg +
(static_cast<int64_t>(g) * p.NH() + nh) *
(static_cast<int64_t>(kV) * kK) +
static_cast<int64_t>(v0 / 16) * kStTile + lane * 4;
for (int vt = 0; vt < kNV; ++vt)
for (int kt = 0; kt < kNKTL; ++kt)
*reinterpret_cast<f32x4*>(s1 + vt * st_vt + (kt0 + kt) * st_kt) =
*reinterpret_cast<f32x4*>(s1 + vt * kStTile + (kt0 + kt) * kStKt) =
S[kt][vt];
}
#else
kda_unsupported_arch("kda_chunk_fused");
#endif // ifdef ROCM_USE_FUSED_PREFILL_KDA
@@ -2112,7 +2136,7 @@ void fused_kda_chunk(
// Prefill checkpoint export. `checkpoint_state` is either a [N, H, V, K]
// staging buffer, as FlashKDA takes, or the paged state cache itself when
// `checkpoint_state_indices` says which row each sequence owns.
float* ckpt_ptr = nullptr;
void* ckpt_ptr = nullptr;
int const* ckoff_ptr = nullptr;
int const* ckidx_ptr = nullptr;
if (checkpoint_state.has_value() || checkpoint_offsets.has_value()) {
@@ -2121,7 +2145,8 @@ void fused_kda_chunk(
"checkpoint_state and checkpoint_offsets must be given together");
STD_TORCH_CHECK(
checkpoint_state->is_cuda() &&
checkpoint_state->scalar_type() == ScalarType::Float &&
(checkpoint_state->scalar_type() == ScalarType::Float ||
checkpoint_state->scalar_type() == ScalarType::BFloat16) &&
checkpoint_state->dim() == 4 &&
checkpoint_state->size(1) == num_heads &&
checkpoint_state->size(2) == kV &&
@@ -2129,7 +2154,8 @@ void fused_kda_chunk(
checkpoint_state->stride(3) == 1 &&
checkpoint_state->stride(2) == kK &&
checkpoint_state->stride(1) == kV * kK,
"checkpoint_state must be an fp32 [rows, H, 128, 128] tensor whose "
"checkpoint_state must be an fp32 or bf16 [rows, H, 128, 128] tensor "
"whose "
"rows are dense; stride(0) may be padded, as the paged state cache "
"it usually is");
STD_TORCH_CHECK(checkpoint_offsets->is_cuda() &&
@@ -2154,7 +2180,7 @@ void fused_kda_chunk(
checkpoint_state->size(0), " rows for ", num_seqs,
" sequences");
}
ckpt_ptr = static_cast<float*>(checkpoint_state->data_ptr());
ckpt_ptr = checkpoint_state->data_ptr();
ckoff_ptr = static_cast<int const*>(checkpoint_offsets->data_ptr());
}
@@ -2162,8 +2188,8 @@ void fused_kda_chunk(
// reads and writes row `state_indices[n]` of a paged cache. Without it, the
// dense `initial_state` / `final_state` pair is used, which is the same
// layout with one row per sequence in order.
float const* cache_in = h0_ptr;
float* cache_out = ht_ptr;
void const* cache_in = h0_ptr;
void* cache_out = ht_ptr;
int const* sidx_ptr = nullptr;
bool const* hh0_ptr = nullptr;
if (state_cache.has_value()) {
@@ -2175,13 +2201,14 @@ void fused_kda_chunk(
"state_cache needs state_indices and has_initial_state");
STD_TORCH_CHECK(
state_cache->is_cuda() &&
state_cache->scalar_type() == ScalarType::Float &&
(state_cache->scalar_type() == ScalarType::Float ||
state_cache->scalar_type() == ScalarType::BFloat16) &&
state_cache->dim() == 4 && state_cache->size(1) == num_heads &&
state_cache->size(2) == kV && state_cache->size(3) == kK &&
state_cache->stride(3) == 1 && state_cache->stride(2) == kK &&
state_cache->stride(1) == kV * kK,
"state_cache must be an fp32 [slots, H, 128, 128] tensor whose rows "
"are dense; stride(0) may be padded");
"state_cache must be an fp32 or bf16 [slots, H, 128, 128] tensor whose "
"rows are dense; stride(0) may be padded");
STD_TORCH_CHECK(state_indices->is_cuda() &&
state_indices->is_contiguous() &&
state_indices->scalar_type() == ScalarType::Int &&
@@ -2197,8 +2224,8 @@ void fused_kda_chunk(
// A block reads row `state_indices[n]` at entry and writes it at exit, so
// `checkpoint_state_indices` must name no row in `state_indices`; a
// checkpoint aimed at one would be a write racing that read.
cache_in = static_cast<float const*>(state_cache->data_ptr());
cache_out = static_cast<float*>(state_cache->data_ptr());
cache_in = state_cache->data_ptr();
cache_out = state_cache->data_ptr();
sidx_ptr = static_cast<int const*>(state_indices->data_ptr());
hh0_ptr = static_cast<bool const*>(has_initial_state->data_ptr());
}
@@ -2231,6 +2258,8 @@ void fused_kda_chunk(
params.NHv = static_cast<int>(nh_total);
params.cache_in = cache_in;
params.cache_out = cache_out;
params.cache_bf16 = state_cache.has_value() &&
state_cache->scalar_type() == ScalarType::BFloat16;
params.sidx = sidx_ptr;
params.hh0 = hh0_ptr;
params.crow =
@@ -2238,6 +2267,8 @@ void fused_kda_chunk(
params.ckrow = checkpoint_state.has_value() ? checkpoint_state->stride(0)
: num_heads * kV * kK;
params.ckpt = ckpt_ptr;
params.ckpt_bf16 = checkpoint_state.has_value() &&
checkpoint_state->scalar_type() == ScalarType::BFloat16;
params.ckoff = ckoff_ptr;
params.ckidx = ckidx_ptr;
csrc/moe/moe_permute_unpermute_op.cu
+10 -11
View File
@@ -104,18 +104,17 @@ torch::stable::Tensor moe_sort_routing(
torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
CubKeyValueSorter sorter{};
torch::stable::Tensor topk_ids_for_sort = topk_ids;
if (expert_map.has_value() || build_inverse) {
const int* expert_map_ptr =
expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr;
topk_ids_for_sort = maybe_allocate_tensor(
maybe_topk_ids_for_sort, topk_ids.sizes(),
torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
torch::stable::copy_(topk_ids_for_sort, topk_ids);
preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
expert_map_ptr, n_expert, stream);
}
// Map invalid routes (e.g. -1 for padding tokens) to the n_expert sentinel so
// the offset scan sees keys in ascending order.
const int* expert_map_ptr =
expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr;
auto topk_ids_for_sort = maybe_allocate_tensor(
maybe_topk_ids_for_sort, topk_ids.sizes(),
torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
torch::stable::copy_(topk_ids_for_sort, topk_ids);
preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
expert_map_ptr, n_expert, stream);
sortAndScanExpert(
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+3
View File
@@ -171,6 +171,9 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
void preprocessTopkIdLauncher(int* topk_id_ptr, int size,
const int* expert_map_ptr, int num_experts,
cudaStream_t stream) {
if (size == 0) {
return;
}
int block = std::min(size, 1024);
int grid = (size + block - 1) / block;
int smem_size = (num_experts) * sizeof(int);
csrc/ops.h
+3 -2
View File
@@ -85,7 +85,8 @@ void get_cutlass_moe_mm_data(
torch::stable::Tensor& output_permutation, const int64_t num_experts,
const int64_t n, const int64_t k,
const std::optional<torch::stable::Tensor>& blockscale_offsets,
const bool is_gated);
const bool is_gated,
const std::optional<torch::stable::Tensor>& sorted_expert_ids);
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
const torch::stable::Tensor& expert_first_token_offset,
@@ -712,7 +713,7 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c,
torch::stable::Tensor& kv_cache,
torch::stable::Tensor& slot_mapping,
const std::string& kv_cache_dtype,
torch::stable::Tensor& scale);
torch::stable::Tensor& scale, bool skip_scale);
void concat_and_cache_mla_grouped(
torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe,
csrc/persistent_topk.cuh
+266 -472
View File
File diff suppressed because it is too large. Load diff
csrc/persistent_topk_overflow.cuh
+262
View File
@@ -0,0 +1,262 @@
#ifndef PERSISTENT_TOPK_OVERFLOW_CUH_
#define PERSISTENT_TOPK_OVERFLOW_CUH_
#include <cstddef>
#include <cstdint>
#include <cuda_runtime.h>
#include <cub/cub.cuh>
#include "topk_histogram_4096.cuh"
#include "topk_histogram_4096_overflow.cuh"
namespace vllm::persistent::overflow {
// Where exact recovery resumes after the medium path's coarse histogram:
// the selected 11-bit FP32 prefix and how many keys are still needed in it.
struct MediumResume {
bool needed = false;
uint32_t prefix = 0;
uint32_t remaining = 0;
};
template <int TopK, int BlockSize, int VecSize>
__device__ __noinline__ void recover_medium(const float* scores,
int32_t* output, uint32_t length,
void* smem, uint32_t prefix,
uint32_t remaining) {
topk_histogram_4096::exact_topk_rescan<TopK, BlockSize, VecSize, 1, 4096>(
scores, output, length, smem, prefix, remaining);
}
} // namespace vllm::persistent::overflow
namespace vllm::filtered_topk {
template <typename T, size_t N>
struct vec_t;
template <typename DType, bool UseWideCoarse>
struct FilteredTopKTraits;
namespace overflow {
// Sampled recovery pays off only for long rows in large batches. topk.cu
// currently routes such batches to sampled_topk_kernel, so this matters for
// direct FilteredTopKRaggedTransform callers and future dispatch changes.
constexpr uint32_t kSamplingMinRows = 96;
constexpr uint32_t kSamplingMinLength = 192 * 1024;
// 8-bit digits over the ordered key: the sampled candidate set is small, so a
// 256-bin histogram selected by one thread is cheaper than a block scan.
struct SampledRadixDigits {
static constexpr uint32_t kBins = 256;
static constexpr uint32_t kRounds = 4;
__device__ static constexpr uint32_t shift(uint32_t round) {
return 24 - round * 8;
}
__device__ static constexpr uint32_t digit_mask(uint32_t) { return 0xFFu; }
__device__ static constexpr uint32_t prefix_mask(uint32_t round) {
return round == 0 ? 0u : (~0u << (32 - round * 8));
}
};
// Sample a pivot, bound the candidate set, then select it exactly by radix.
template <typename DType, typename IdType, int VecSize, uint32_t BlockSize,
uint32_t CandidateCapacity>
__device__ bool try_sampled_exact_topk(const DType* __restrict__ scores,
IdType* __restrict__ output,
uint32_t length, uint32_t top_k,
void* dynamic_smem, int* scratch) {
using Traits = FilteredTopKTraits<DType, false>;
using Radix = SampledRadixDigits;
constexpr uint32_t kItemsPerThread = 2;
constexpr uint32_t kSampleSize = BlockSize * kItemsPerThread;
using SampleSort = cub::BlockRadixSort<uint32_t, BlockSize, kItemsPerThread>;
static_assert(sizeof(typename SampleSort::TempStorage) <=
2 * CandidateCapacity * sizeof(int));
// scratch: the radix histogram followed by four counters.
int* histogram = scratch;
int* candidate_count = scratch + Radix::kBins;
int* output_counter = candidate_count + 1;
int* radix_prefix = candidate_count + 2;
int* radix_remaining = candidate_count + 3;
const uint32_t tx = threadIdx.x;
uint32_t sample_keys[kItemsPerThread];
#pragma unroll
for (uint32_t i = 0; i < kItemsPerThread; ++i) {
const uint32_t sample_idx = tx * kItemsPerThread + i;
const uint32_t score_idx = static_cast<uint32_t>(
(static_cast<uint64_t>(2 * sample_idx + 1) * length) /
(2 * kSampleSize));
sample_keys[i] = Traits::ToOrdered(scores[score_idx]);
}
auto* sort_smem =
static_cast<typename SampleSort::TempStorage*>(dynamic_smem);
SampleSort(*sort_smem).SortDescending(sample_keys);
__syncthreads();
const uint32_t target_candidates =
min(CandidateCapacity / 2, (top_k < 2048) ? 4 * top_k : 2 * top_k);
uint32_t sample_rank = static_cast<uint32_t>(
(static_cast<uint64_t>(target_candidates) * kSampleSize + length - 1) /
length);
sample_rank = max(1u, min(sample_rank, kSampleSize));
const uint32_t rank_index = sample_rank - 1;
if (tx == rank_index / kItemsPerThread) {
*radix_prefix = static_cast<int>(sample_keys[rank_index % kItemsPerThread]);
}
if (tx == 0) *candidate_count = 0;
__syncthreads();
auto* candidates = static_cast<int*>(dynamic_smem);
const uint32_t proposed_pivot = static_cast<uint32_t>(*radix_prefix);
// Every key tied with the pivot is a candidate. When the sample already
// predicts more candidates than fit, skip the doomed collection scan.
uint32_t samples_at_or_above = 0;
#pragma unroll
for (uint32_t i = 0; i < kItemsPerThread; ++i) {
samples_at_or_above +=
__syncthreads_count(sample_keys[i] >= proposed_pivot);
}
if (static_cast<uint64_t>(samples_at_or_above) * length >
static_cast<uint64_t>(CandidateCapacity) * kSampleSize) {
return false;
}
const auto collect_candidate = [&](uint32_t idx, DType score) {
if (Traits::ToOrdered(score) >= proposed_pivot) {
const int pos = atomicAdd(candidate_count, 1);
if (pos < static_cast<int>(CandidateCapacity)) candidates[pos] = idx;
}
};
vec_t<DType, VecSize> score_vec;
const uint32_t aligned_length = length / VecSize * VecSize;
for (uint32_t base = tx * VecSize; base < aligned_length;
base += BlockSize * VecSize) {
score_vec.cast_load(scores + base);
#pragma unroll
for (uint32_t i = 0; i < VecSize; ++i) {
collect_candidate(base + i, score_vec[i]);
}
}
for (uint32_t idx = aligned_length + tx; idx < length; idx += BlockSize) {
collect_candidate(idx, scores[idx]);
}
__syncthreads();
const int num_candidates = *candidate_count;
if (num_candidates < static_cast<int>(top_k) ||
num_candidates > static_cast<int>(CandidateCapacity)) {
return false;
}
auto* candidate_keys = candidates + CandidateCapacity;
for (int i = tx; i < num_candidates; i += BlockSize) {
candidate_keys[i] =
static_cast<int>(Traits::ToOrdered(scores[candidates[i]]));
}
if (tx == 0) {
*radix_prefix = 0;
*radix_remaining = static_cast<int>(top_k);
}
__syncthreads();
#pragma unroll
for (uint32_t round = 0; round < Radix::kRounds; ++round) {
if (tx < Radix::kBins) histogram[tx] = 0;
__syncthreads();
const uint32_t prefix = static_cast<uint32_t>(*radix_prefix);
const uint32_t prefix_mask = Radix::prefix_mask(round);
const uint32_t shift = Radix::shift(round);
for (int i = tx; i < num_candidates; i += BlockSize) {
const uint32_t ordered = static_cast<uint32_t>(candidate_keys[i]);
if ((ordered & prefix_mask) == prefix) {
atomicAdd(&histogram[(ordered >> shift) & Radix::digit_mask(round)], 1);
}
}
__syncthreads();
if (tx == 0) {
int count_above = 0;
const int remaining = *radix_remaining;
for (int bin = static_cast<int>(Radix::kBins) - 1; bin >= 0; --bin) {
const int count = histogram[bin];
if (count_above + count >= remaining) {
*radix_remaining = remaining - count_above;
*radix_prefix =
static_cast<int>(prefix | (static_cast<uint32_t>(bin) << shift));
break;
}
count_above += count;
}
}
__syncthreads();
}
if (tx == 0) {
*output_counter = 0;
histogram[0] = 0;
}
__syncthreads();
const uint32_t pivot = static_cast<uint32_t>(*radix_prefix);
const int equal_count = *radix_remaining;
const int equal_base = static_cast<int>(top_k) - equal_count;
for (int i = tx; i < num_candidates; i += BlockSize) {
const int idx = candidates[i];
const uint32_t ordered = static_cast<uint32_t>(candidate_keys[i]);
if (ordered > pivot) {
output[atomicAdd(output_counter, 1)] = static_cast<IdType>(idx);
} else if (ordered == pivot) {
const int pos = atomicAdd(&histogram[0], 1);
if (pos < equal_count)
output[equal_base + pos] = static_cast<IdType>(idx);
}
}
__syncthreads();
return true;
}
// Exact selection after the filtered path's coarse bin overflowed the
// candidate stash. The coarse histogram holds suffix counts over the top
// CoarseBits of the ordered key; recovery resumes from its top 11 bits.
// scratch must hold Radix::kBins + 4 ints for the sampled path.
template <typename DType, typename IdType, int VecSize, uint32_t BlockSize,
uint32_t MaxK, uint32_t CandidateCapacity, int CoarseBits>
__device__ __noinline__ void recover_coarse(const DType* scores, IdType* output,
uint32_t length, uint32_t top_k,
uint32_t num_rows,
const int* coarse_suffix,
uint32_t threshold_bin,
void* dynamic_smem, int* scratch) {
static_assert(CoarseBits >= 11);
constexpr uint32_t kShift = CoarseBits - 11;
const uint32_t prefix_bin = threshold_bin >> kShift;
const uint32_t remaining =
top_k - static_cast<uint32_t>(coarse_suffix[(prefix_bin + 1) << kShift]);
// The coarse histogram shares dynamic smem with the recovery buffers.
__syncthreads();
// Sampling amortizes its fixed block-sort cost only for long, busy batches.
if (num_rows >= kSamplingMinRows && length >= kSamplingMinLength &&
try_sampled_exact_topk<DType, IdType, VecSize, BlockSize,
CandidateCapacity>(scores, output, length, top_k,
dynamic_smem, scratch)) {
return;
}
topk_histogram_4096::exact_topk_rescan<MaxK, BlockSize, VecSize, 1,
CandidateCapacity>(
scores, output, length, dynamic_smem, prefix_bin << 21, remaining);
}
} // namespace overflow
} // namespace vllm::filtered_topk
#endif // PERSISTENT_TOPK_OVERFLOW_CUH_
csrc/quantization/fp4/mxfp4_experts_quant.cu
+60 -30
View File
@@ -22,6 +22,7 @@
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <optional>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
@@ -66,7 +67,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts, bool low_latency) {
int n_experts, int32_t const* sorted_expert_ids,
bool low_latency) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
@@ -89,7 +91,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
int rowIdx_in_expert = 0;
int expert_idx = 0;
if constexpr (SMALL_NUM_EXPERTS) {
if (sorted_expert_ids != nullptr) {
expert_idx = __ldca(&sorted_expert_ids[rowIdx]);
rowIdx_in_expert = rowIdx - __ldca(&input_offset_by_experts[expert_idx]);
} else if constexpr (SMALL_NUM_EXPERTS) {
for (int i = 0; i < n_experts; i++) {
uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
@@ -168,7 +173,7 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts) {
int n_experts, int32_t const* sorted_expert_ids) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
@@ -208,21 +213,26 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
int rowIdx_in_expert = 0;
int expert_idx = 0;
// Binary search through experts using shared memory
int left = 0, right = n_experts - 1;
while (left <= right) {
int mid = (left + right) / 2;
uint32_t mid_offset = shared_input_offsets[mid];
uint32_t next_offset = shared_input_offsets[mid + 1];
if (sorted_expert_ids != nullptr) {
expert_idx = __ldca(&sorted_expert_ids[rowIdx]);
rowIdx_in_expert = rowIdx - shared_input_offsets[expert_idx];
} else {
// Binary search through experts using shared memory
int left = 0, right = n_experts - 1;
while (left <= right) {
int mid = (left + right) / 2;
uint32_t mid_offset = shared_input_offsets[mid];
uint32_t next_offset = shared_input_offsets[mid + 1];
if (rowIdx >= mid_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - mid_offset;
expert_idx = mid;
break;
} else if (rowIdx < mid_offset) {
right = mid - 1;
} else {
left = mid + 1;
if (rowIdx >= mid_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - mid_offset;
expert_idx = mid;
break;
} else if (rowIdx < mid_offset) {
right = mid - 1;
} else {
left = mid + 1;
}
}
}
@@ -259,7 +269,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
template <typename T, bool FUSE_SILU_MUL = false>
void mxfp4_quant_impl(void* output, void* output_scale, void* input,
void* input_offset_by_experts,
void* output_scale_offset_by_experts, int m_topk, int k,
void* output_scale_offset_by_experts,
int32_t const* sorted_expert_ids, int m_topk, int k,
int n_experts, cudaStream_t stream) {
int multiProcessorCount =
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -288,7 +299,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
n_experts, sorted_expert_ids);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
<<<grid, block, shared_mem_size, stream>>>(
@@ -297,7 +308,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
n_experts, sorted_expert_ids);
}
} else {
if (n_experts >= 16) {
@@ -308,7 +319,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
n_experts, sorted_expert_ids,
/* bool low_latency */ true);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
@@ -316,7 +328,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
n_experts, sorted_expert_ids,
/* bool low_latency */ true);
}
}
}
@@ -345,12 +358,21 @@ static void validate_mxfp4_experts_quant_inputs(
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
std::optional<torch::stable::Tensor> const& sorted_expert_ids,
int64_t n_experts, int64_t m_topk, int64_t k) {
CHECK_INPUT(output, "output");
CHECK_INPUT(output_scale, "output_scale");
CHECK_INPUT(input, "input");
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
if (sorted_expert_ids.has_value()) {
CHECK_INPUT(sorted_expert_ids.value(), "sorted_expert_ids");
STD_TORCH_CHECK(sorted_expert_ids->dim() == 1);
STD_TORCH_CHECK(sorted_expert_ids->scalar_type() == INT);
STD_TORCH_CHECK(sorted_expert_ids->size(0) == m_topk);
STD_TORCH_CHECK(sorted_expert_ids->get_device_index() ==
input.get_device_index());
}
STD_TORCH_CHECK(output.dim() == 2);
STD_TORCH_CHECK(output_scale.dim() == 2);
@@ -393,7 +415,8 @@ void mxfp4_experts_quant(
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
int64_t n_experts,
std::optional<torch::stable::Tensor> const& sorted_expert_ids) {
#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
int32_t sm = get_sm_version_num();
STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
@@ -405,7 +428,7 @@ void mxfp4_experts_quant(
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -417,8 +440,11 @@ void mxfp4_experts_quant(
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
output_scale_offset_by_experts.data_ptr(),
sorted_expert_ids.has_value()
? static_cast<int32_t const*>(sorted_expert_ids->data_ptr())
: nullptr,
m_topk, k, n_experts, stream);
});
#else
STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
@@ -431,7 +457,8 @@ void silu_and_mul_mxfp4_experts_quant(
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
int64_t n_experts,
std::optional<torch::stable::Tensor> const& sorted_expert_ids) {
#if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
int32_t sm = get_sm_version_num();
STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
@@ -445,7 +472,7 @@ void silu_and_mul_mxfp4_experts_quant(
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
@@ -457,8 +484,11 @@ void silu_and_mul_mxfp4_experts_quant(
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
output_scale_offset_by_experts.data_ptr(),
sorted_expert_ids.has_value()
? static_cast<int32_t const*>(sorted_expert_ids->data_ptr())
: nullptr,
m_topk, k, n_experts, stream);
});
#else
STD_TORCH_CHECK_NOT_IMPLEMENTED(
csrc/quantization/w8a8/cutlass/moe/moe_data.cu
+9 -1
View File
@@ -88,6 +88,7 @@ __global__ void compute_expert_blockscale_offsets(
__global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
int32_t* input_permutation,
int32_t* output_permutation,
int32_t* sorted_expert_ids,
int32_t* atomic_buffer, const int topk_length,
const int topk) {
int const blk_expert_id = blockIdx.x;
@@ -108,6 +109,9 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
int start = atomicAdd(&atomic_buffer[expert_id], 1);
input_permutation[start] = i / topk;
output_permutation[i] = start;
if (sorted_expert_ids != nullptr) {
sorted_expert_ids[start] = expert_id;
}
}
}
}
@@ -242,7 +246,8 @@ void get_cutlass_moe_mm_data_caller(
torch::stable::Tensor& output_permutation, const int64_t num_experts,
const int64_t n, const int64_t k,
const std::optional<torch::stable::Tensor>& blockscale_offsets,
const bool is_gated) {
const bool is_gated,
const std::optional<torch::stable::Tensor>& sorted_expert_ids) {
auto device = topk_ids.device();
const torch::stable::accelerator::DeviceGuard device_guard(device.index());
auto stream = get_current_cuda_stream(device.index());
@@ -278,6 +283,9 @@ void get_cutlass_moe_mm_data_caller(
static_cast<const int32_t*>(topk_ids.data_ptr()),
static_cast<int32_t*>(input_permutation.data_ptr()),
static_cast<int32_t*>(output_permutation.data_ptr()),
sorted_expert_ids.has_value()
? static_cast<int32_t*>(sorted_expert_ids->data_ptr())
: nullptr,
static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
topk_ids.size(1));
}
csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu
+8 -6
View File
@@ -97,7 +97,8 @@ void get_cutlass_moe_mm_data_caller(
torch::stable::Tensor& output_permutation, const int64_t num_experts,
const int64_t n, const int64_t k,
const std::optional<torch::stable::Tensor>& blockscale_offsets,
const bool is_gated);
const bool is_gated,
const std::optional<torch::stable::Tensor>& sorted_expert_ids);
void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
const torch::stable::Tensor& expert_first_token_offset,
@@ -337,17 +338,18 @@ void get_cutlass_moe_mm_data(
torch::stable::Tensor& output_permutation, const int64_t num_experts,
const int64_t n, const int64_t k,
const std::optional<torch::stable::Tensor>& blockscale_offsets,
const bool is_gated) {
const bool is_gated,
const std::optional<torch::stable::Tensor>& sorted_expert_ids) {
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
problem_sizes2, input_permutation,
output_permutation, num_experts, n, k,
blockscale_offsets, is_gated);
get_cutlass_moe_mm_data_caller(
topk_ids, expert_offsets, problem_sizes1, problem_sizes2,
input_permutation, output_permutation, num_experts, n, k,
blockscale_offsets, is_gated, sorted_expert_ids);
return;
#endif
STD_TORCH_CHECK_NOT_IMPLEMENTED(
csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu
+254
View File
@@ -0,0 +1,254 @@
// MXFP4 x int8 (W4A8) GEMV for RDNA3/RDNA3.5 (gfx11) small-batch decode.
//
// Activations are quantized to int8 per 32-element group (symmetric, one fp32
// scale per group), E2M1 weight codes map exactly to int8 (2 * e2m1), and each
// 32-element group is reduced with eight v_dot4_i32_iu8 instructions. One fp32
// multiply per group applies the E8M0 weight scale and the activation scale.
// The scheme follows ggml's MMVQ / q8_1 path (llama.cpp, MIT License).
//
// Layouts (Quark OCP MX checkpoints, no repack):
// a [M, K] fp16 / bf16, 1 <= M <= 8
// b_q [N, K/2] uint8, two E2M1 codes per byte, low nibble = even k
// b_scale [N, K/32] uint8 E8M0
// out [M, N] same dtype as a
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <string>
#include "core/cuda_compat.h"
#if defined(__GFX11__)
#define __HIP__MXFP4_W4A8_GFX11__
#endif
namespace {
constexpr int kMaxM = 8;
constexpr int kGroup = 32;
constexpr int kWarpsPerBlock = 8;
template <typename T>
__device__ __forceinline__ float to_float(T v);
template <>
__device__ __forceinline__ float to_float(half v) {
return __half2float(v);
}
template <>
__device__ __forceinline__ float to_float(__hip_bfloat16 v) {
return __bfloat162float(v);
}
template <typename T>
__device__ __forceinline__ T from_float(float v);
template <>
__device__ __forceinline__ half from_float(float v) {
return __float2half(v);
}
template <>
__device__ __forceinline__ __hip_bfloat16 from_float(float v) {
return __float2bfloat16(v);
}
#if defined(__HIP__MXFP4_W4A8_GFX11__) || !defined(__HIP_DEVICE_COMPILE__)
// 2 * E2M1 value for each 4-bit code (sign in bit 3). Exact in int8.
__device__ __constant__ int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12,
0, -1, -2, -3, -4, -6, -8, -12};
template <typename T>
__global__ void mxfp4_w4a8_quant_act_kernel(const T* __restrict__ a,
int8_t* __restrict__ a_q,
float* __restrict__ a_s, int M,
int K) {
const int groups = K / kGroup;
const int m = blockIdx.x;
const int g = blockIdx.y * blockDim.x + threadIdx.x;
if (m >= M || g >= groups) return;
const T* row = a + (int64_t)m * K + g * kGroup;
float v[kGroup];
float amax = 0.f;
#pragma unroll
for (int i = 0; i < kGroup; ++i) {
v[i] = to_float(row[i]);
amax = fmaxf(amax, fabsf(v[i]));
}
const float scale = fmaxf(amax / 127.0f, 1e-12f);
a_s[(int64_t)m * groups + g] = scale;
int8_t* out = a_q + (int64_t)m * K + g * kGroup;
#pragma unroll
for (int i = 0; i < kGroup; ++i) {
out[i] = (int8_t)fmaxf(-127.f, fminf(127.f, rintf(v[i] / scale)));
}
}
// One wave (32 lanes) per output column n. Lanes stride over the K groups;
// each lane decodes its weight group once and reuses it for all M rows.
template <int M, typename T>
__global__ void mxfp4_w4a8_gemv_rdna3_kernel(
const int* __restrict__ a_q, const float* __restrict__ a_s,
const uint8_t* __restrict__ b_q, const uint8_t* __restrict__ b_scale,
T* __restrict__ out, int N, int K) {
const int groups = K / kGroup;
const int lane = threadIdx.x;
const int n = blockIdx.x * blockDim.y + threadIdx.y;
if (n >= N) return;
const uint8_t* w_row = b_q + (int64_t)n * (K / 2);
const uint8_t* s_row = b_scale + (int64_t)n * groups;
float acc[M];
#pragma unroll
for (int m = 0; m < M; ++m) acc[m] = 0.f;
for (int g = lane; g < groups; g += 32) {
const uint8_t* gp = w_row + g * (kGroup / 2);
int w4[8];
#pragma unroll
for (int j = 0; j < 8; ++j) {
const uint32_t two = gp[2 * j] | (gp[2 * j + 1] << 8);
const int w0 = kE2M1x2[two & 0xF];
const int w1 = kE2M1x2[(two >> 4) & 0xF];
const int w2 = kE2M1x2[(two >> 8) & 0xF];
const int w3 = kE2M1x2[(two >> 12) & 0xF];
w4[j] = (w0 & 0xFF) | ((w1 & 0xFF) << 8) | ((w2 & 0xFF) << 16) |
((w3 & 0xFF) << 24);
}
// 2^(e - 127) / 2: undoes the factor 2 folded into kE2M1x2.
const float w_scale = __int_as_float((int)s_row[g] << 23) * 0.5f;
#pragma unroll
for (int m = 0; m < M; ++m) {
const int* ap = a_q + (int64_t)m * (K / 4) + g * (kGroup / 4);
int isum = 0;
#pragma unroll
for (int j = 0; j < 8; ++j) {
isum = __builtin_amdgcn_sudot4(true, w4[j], true, ap[j], isum, false);
}
acc[m] += (float)isum * (w_scale * a_s[(int64_t)m * groups + g]);
}
}
#pragma unroll
for (int m = 0; m < M; ++m) {
float v = acc[m];
#pragma unroll
for (int o = 16; o > 0; o >>= 1) v += VLLM_SHFL_XOR_SYNC_WIDTH(v, o, 32);
if (lane == 0) out[(int64_t)m * N + n] = from_float<T>(v);
}
}
#else // non-gfx11 device pass: empty kernels for symbol parity.
template <typename T>
__global__ void mxfp4_w4a8_quant_act_kernel(const T*, int8_t*, float*, int,
int) {}
template <int M, typename T>
__global__ void mxfp4_w4a8_gemv_rdna3_kernel(const int*, const float*,
const uint8_t*, const uint8_t*, T*,
int, int) {}
#endif // __HIP__MXFP4_W4A8_GFX11__ || !__HIP_DEVICE_COMPILE__
bool mxfp4_w4a8_on_gfx11() {
static const bool result = [] {
const auto* dprops = at::cuda::getCurrentDeviceProperties();
const std::string arch = dprops->gcnArchName;
return arch.rfind("gfx11", 0) == 0;
}();
return result;
}
template <typename T>
void launch(const at::Tensor& a, const at::Tensor& b_q,
const at::Tensor& b_scale, at::Tensor& out, int M, int N, int K) {
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const int groups = K / kGroup;
auto a_q = torch::empty({M, K}, a.options().dtype(torch::kInt8));
auto a_s = torch::empty({M, groups}, a.options().dtype(torch::kFloat32));
const dim3 q_block(256);
const dim3 q_grid(M, (groups + 255) / 256);
mxfp4_w4a8_quant_act_kernel<T><<<q_grid, q_block, 0, stream>>>(
reinterpret_cast<const T*>(a.data_ptr()), a_q.data_ptr<int8_t>(),
a_s.data_ptr<float>(), M, K);
const dim3 block(32, kWarpsPerBlock);
const dim3 grid((N + kWarpsPerBlock - 1) / kWarpsPerBlock);
const int* a_q32 = reinterpret_cast<const int*>(a_q.data_ptr<int8_t>());
const float* a_s_ptr = a_s.data_ptr<float>();
const uint8_t* b_q_ptr = b_q.data_ptr<uint8_t>();
const uint8_t* b_s_ptr = b_scale.data_ptr<uint8_t>();
T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
#define MXFP4_W4A8_LAUNCH(MM) \
case MM: \
mxfp4_w4a8_gemv_rdna3_kernel<MM, T><<<grid, block, 0, stream>>>( \
a_q32, a_s_ptr, b_q_ptr, b_s_ptr, out_ptr, N, K); \
break;
switch (M) {
MXFP4_W4A8_LAUNCH(1)
MXFP4_W4A8_LAUNCH(2)
MXFP4_W4A8_LAUNCH(3)
MXFP4_W4A8_LAUNCH(4)
MXFP4_W4A8_LAUNCH(5)
MXFP4_W4A8_LAUNCH(6)
MXFP4_W4A8_LAUNCH(7)
MXFP4_W4A8_LAUNCH(8)
}
#undef MXFP4_W4A8_LAUNCH
}
} // namespace
torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q,
const at::Tensor& b_scale) {
TORCH_CHECK(a.is_cuda() && b_q.is_cuda() && b_scale.is_cuda(),
"mxfp4_w4a8_gemv_rdna3: all inputs must be on the GPU");
TORCH_CHECK(mxfp4_w4a8_on_gfx11(),
"mxfp4_w4a8_gemv_rdna3 requires an RDNA3/RDNA3.5 (gfx11) GPU");
TORCH_CHECK(a.scalar_type() == at::kHalf || a.scalar_type() == at::kBFloat16,
"mxfp4_w4a8_gemv_rdna3: a must be float16 or bfloat16");
TORCH_CHECK(
b_q.scalar_type() == at::kByte && b_scale.scalar_type() == at::kByte,
"mxfp4_w4a8_gemv_rdna3: b_q and b_scale must be uint8");
TORCH_CHECK(a.dim() == 2 && b_q.dim() == 2 && b_scale.dim() == 2,
"mxfp4_w4a8_gemv_rdna3: inputs must be 2D");
TORCH_CHECK(
a.is_contiguous() && b_q.is_contiguous() && b_scale.is_contiguous(),
"mxfp4_w4a8_gemv_rdna3: inputs must be contiguous");
const int M = a.size(0);
const int K = a.size(1);
const int N = b_q.size(0);
TORCH_CHECK(M >= 1 && M <= kMaxM,
"mxfp4_w4a8_gemv_rdna3 supports 1 <= M <= ", kMaxM,
", got M=", M);
TORCH_CHECK(K % kGroup == 0,
"mxfp4_w4a8_gemv_rdna3: K must be a multiple of 32");
TORCH_CHECK(b_q.size(1) * 2 == K,
"mxfp4_w4a8_gemv_rdna3: b_q must be [N, K/2]");
TORCH_CHECK(b_scale.size(0) == N && b_scale.size(1) == K / kGroup,
"mxfp4_w4a8_gemv_rdna3: b_scale must be [N, K/32]");
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
auto out = torch::empty({M, N}, a.options());
if (a.scalar_type() == at::kHalf) {
launch<half>(a, b_q, b_scale, out, M, N, K);
} else {
launch<__hip_bfloat16>(a, b_q, b_scale, out, M, N, K);
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return out;
}
csrc/rocm/ops.h
+3
View File
@@ -24,6 +24,9 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
const at::Tensor& scale_a, const at::Tensor& scale_b,
const int64_t CuCount);
torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q,
const at::Tensor& b_scale);
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros, torch::Tensor b_scales,
bool use_v2_format);
csrc/rocm/torch_bindings.cpp
+6
View File
@@ -54,6 +54,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
#endif // VLLM_SKIP_SKINNY_GEMMS
#ifdef VLLM_ROCM_GFX1100
// MXFP4 weights x int8 activations GEMV for RDNA3/RDNA3.5 (gfx11) decode,
// 1 <= M <= 8. Weights [N, K/2] uint8 E2M1, scales [N, K/32] uint8 E8M0.
rocm_ops.def(
"mxfp4_w4a8_gemv_rdna3(Tensor a, Tensor b_q, Tensor b_scale) -> Tensor");
rocm_ops.impl("mxfp4_w4a8_gemv_rdna3", torch::kCUDA, &mxfp4_w4a8_gemv_rdna3);
// W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
rocm_ops.def(
"gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
csrc/sampled_topk.cuh
+20 -17
View File
@@ -17,6 +17,8 @@
namespace vllm::sampled_topk {
namespace hist4096 = topk_histogram_4096;
constexpr int kThreads = 1024;
// Conservative crossover bounds from the B300 batch/length sweep.
template <int K>
@@ -35,11 +37,6 @@ struct Storage {
uint32_t prefix;
};
__device__ __forceinline__ uint32_t ordered(float value) {
const uint32_t bits = __float_as_uint(value);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
__device__ __forceinline__ uint32_t coarse(uint32_t key) {
const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key;
return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits));
@@ -83,7 +80,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count,
for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0;
__syncthreads();
for (int i = tid; i < count; i += kThreads) {
const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
const uint32_t key =
Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]);
atomicAdd(&s.histogram[coarse(key)], 1);
}
__syncthreads();
@@ -95,7 +93,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count,
__syncthreads();
const uint32_t prefix = s.prefix;
for (int i = tid; i < count; i += kThreads) {
const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]);
const uint32_t key =
Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]);
if (coarse(key) == coarse_bin && (key & mask) == prefix) {
atomicAdd(&s.histogram[(key >> shift) & 255], 1);
}
@@ -127,9 +126,11 @@ __device__ void emit(Storage& s, const float* row, int32_t* dst, int count,
__syncthreads();
for (int base = 0; base < count; base += kThreads) {
const int i = base + tid;
const uint2 pair = i < count ? (Buffered ? s.candidates[i]
: make_uint2(i, ordered(row[i])))
: make_uint2(0, 0);
const uint2 pair =
i < count
? (Buffered ? s.candidates[i]
: make_uint2(i, hist4096::convert_to_uint32_v2(row[i])))
: make_uint2(0, 0);
bool keep = i < count && pair.y > cutoff;
if (i < count && pair.y == cutoff) {
keep = atomicSub(&s.remaining, 1) > 0;
@@ -176,14 +177,16 @@ __global__ void __launch_bounds__(kThreads)
// Sample contiguous warps spread across the row to keep loads coalesced.
for (int i = tid; i < kSample; i += kThreads) {
const int index = (i / 32) * (length / (kSample / 32)) + i % 32;
atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1);
atomicAdd(&s.histogram[coarse(hist4096::convert_to_uint32_v2(row[index]))],
1);
}
__syncthreads();
select_bin<2048>(s);
const uint16_t half_key = s.bin << 5;
const uint16_t half_bits =
(half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key;
const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits)));
const uint32_t cutoff =
hist4096::convert_to_uint32_v2(__half2float(__ushort_as_half(half_bits)));
if (tid == 0) s.count = 0;
__syncthreads();
constexpr int kItems = 16;
@@ -198,16 +201,16 @@ __global__ void __launch_bounds__(kThreads)
for (int j = 0; j < kItems / 4; ++j) {
const float4 values =
reinterpret_cast<const float4*>(row + base)[tid + j * kThreads];
keys[4 * j] = ordered(values.x);
keys[4 * j + 1] = ordered(values.y);
keys[4 * j + 2] = ordered(values.z);
keys[4 * j + 3] = ordered(values.w);
keys[4 * j] = hist4096::convert_to_uint32_v2(values.x);
keys[4 * j + 1] = hist4096::convert_to_uint32_v2(values.y);
keys[4 * j + 2] = hist4096::convert_to_uint32_v2(values.z);
keys[4 * j + 3] = hist4096::convert_to_uint32_v2(values.w);
}
} else {
#pragma unroll
for (int j = 0; j < kItems; ++j) {
const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4;
keys[j] = i < length ? ordered(row[i]) : 0;
keys[j] = i < length ? hist4096::convert_to_uint32_v2(row[i]) : 0;
}
}
#pragma unroll
csrc/topk.cu
+5 -2
View File
@@ -82,10 +82,13 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
uint32_t max_chunk_elements =
static_cast<uint32_t>(available_for_ordered / sizeof(uint32_t));
// Views like x[:, 1:] keep the stride but shift the base.
const auto base =
reinterpret_cast<uintptr_t>(logits.const_data_ptr<float>());
uint32_t vec_size = 1;
if (stride % 4 == 0)
if (stride % 4 == 0 && base % 16 == 0)
vec_size = 4;
else if (stride % 2 == 0)
else if (stride % 2 == 0 && base % 8 == 0)
vec_size = 2;
max_chunk_elements = (max_chunk_elements / vec_size) * vec_size;
csrc/topk_histogram_4096.cuh
+36 -7
View File
@@ -69,7 +69,11 @@ __device__ __forceinline__ void load_float4_predicated(const float* ptr,
// converts the float32 score to a 32-bit ordered unsigned integer — the full
// precision key for radix sorting
__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
uint32_t bits = __float_as_uint(x);
// Read the bits through an opaque move: otherwise nvcc lowers
// `bits | 0x80000000u` to the FP op -|x|, which canonicalizes NaN to
// 0x7FFFFFFF (the key of -0.0) while the FP16 coarse bin puts NaN on top.
uint32_t bits;
asm("mov.b32 %0, %1;" : "=r"(bits) : "f"(x));
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
@@ -334,12 +338,26 @@ __device__ __noinline__ void exact_topk_rescan_cold(
const float* __restrict__ scores, int32_t* __restrict__ output,
uint32_t length, void* smem);
enum class OverflowRecovery { kDisabled, kFp32Rescan };
// kRegisterValues refines from the values already held in registers and
// falls back to the FP32 rescan when the coarse bin cannot bound them.
enum class OverflowRecovery { kFp32Rescan, kRegisterValues };
// Defined in topk_histogram_4096_overflow.cuh, like exact_topk_rescan_cold.
template <uint32_t TopK, uint32_t HistBits, uint32_t VecsPerThread>
__device__ __forceinline__ bool exact_topk_refine_registers(
const float4 (&vecs)[VecsPerThread], int32_t* __restrict__ output,
uint32_t length, uint32_t coarse_bin, uint32_t num_above, void* _smem);
template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize,
uint32_t StartRound, uint32_t FinalCandidateCapacity>
__device__ void exact_topk_rescan(const float* __restrict__ scores,
int32_t* __restrict__ output, uint32_t length,
void* _smem, uint32_t initial_prefix,
uint32_t initial_remaining);
template <uint32_t TopK, uint32_t HIST_BITS,
uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
bool UsePredicatedLoads = false,
OverflowRecovery Recovery = OverflowRecovery::kDisabled>
bool UsePredicatedLoads = false, OverflowRecovery Recovery>
__device__ void histogram_4096_topk(const float* __restrict__ scores,
int32_t* __restrict__ output,
uint32_t length, void* _smem) {
@@ -464,6 +482,16 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores,
return;
}
}
if constexpr (Recovery == OverflowRecovery::kRegisterValues) {
if (need_tie && num_equal > TopK) {
if (!exact_topk_refine_registers<TopK, HIST_BITS, VECS_PER_THREAD>(
vecs, output, length, thr_bin, num_above, _smem)) {
exact_topk_rescan<TopK, kBlockSize, UsePredicatedLoads ? 1 : 4, 0, 0>(
scores, output, length, _smem, 0, TopK);
}
return;
}
}
done = false;
#pragma unroll
@@ -567,12 +595,13 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores,
}
template <uint32_t TopK, uint32_t HIST_BITS,
uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
OverflowRecovery Recovery>
__device__ __noinline__ void histogram_4096_topk_predicated(
const float* __restrict__ scores, int32_t* __restrict__ output,
uint32_t length, void* _smem) {
histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
length, _smem);
histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true, Recovery>(
scores, output, length, _smem);
}
} // namespace topk_histogram_4096
csrc/topk_histogram_4096_overflow.cuh
+292 -32
View File
@@ -62,36 +62,51 @@ struct ExactRadixTraits {
}
};
// Final-round candidates are staged past the rescan's own shared state.
constexpr uint32_t kExactCandidateOffset = 16 * 1024;
// Exact bounded-memory radix selection. Each round rescans the row, so memory
// use and correctness are independent of coarse-bin occupancy.
template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize = 1>
// use and correctness are independent of coarse-bin occupancy. Recovery can
// resume at StartRound from a known prefix, and with FinalCandidateCapacity
// the last round stages its candidates so the output pass skips a row scan.
template <uint32_t TopK, uint32_t BlockSize, uint32_t VecSize = 1,
uint32_t StartRound = 0, uint32_t FinalCandidateCapacity = 0>
__device__ void exact_topk_rescan(const float* __restrict__ scores,
int32_t* __restrict__ output, uint32_t length,
void* _smem) {
void* _smem, uint32_t initial_prefix = 0,
uint32_t initial_remaining = TopK) {
static_assert(VecSize == 1 || VecSize == 2 || VecSize == 4);
using Radix = ExactRadixTraits;
constexpr uint32_t kHistogramBins = Radix::kBins;
static_assert(kHistogramBins % BlockSize == 0);
using BlockScan = cub::BlockScan<uint32_t, BlockSize>;
static_assert(StartRound < Radix::kRounds);
static_assert(BlockSize == kWarpSize * kWarpSize, "one warp sum per lane");
struct ExactSmem {
uint32_t histogram[kHistogramBins];
typename BlockScan::TempStorage scan;
uint32_t warp_sum[BlockSize / kWarpSize];
uint32_t prefix;
uint32_t remaining;
uint32_t output_counter;
uint32_t candidate_count;
uint32_t pivot_count;
};
static_assert(FinalCandidateCapacity == 0 ||
sizeof(ExactSmem) <= kExactCandidateOffset);
static_assert(FinalCandidateCapacity == 0 || TopK <= FinalCandidateCapacity);
auto* smem = static_cast<ExactSmem*>(_smem);
auto* final_candidates = reinterpret_cast<int32_t*>(
static_cast<uint8_t*>(_smem) + kExactCandidateOffset);
const uint32_t tx = threadIdx.x;
if (tx == 0) {
smem->prefix = 0;
smem->remaining = TopK;
smem->prefix = initial_prefix;
smem->remaining = initial_remaining;
}
__syncthreads();
#pragma unroll
for (uint32_t round = 0; round < Radix::kRounds; ++round) {
for (uint32_t round = StartRound; round < Radix::kRounds; ++round) {
for (uint32_t bin = tx; bin < kHistogramBins; bin += BlockSize) {
smem->histogram[bin] = 0;
}
@@ -101,13 +116,39 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
const uint32_t digit_mask = Radix::digit_mask(round);
const uint32_t prefix_mask = Radix::prefix_mask(round);
const uint32_t prefix = smem->prefix;
const bool final_round = round == Radix::kRounds - 1;
scan_scores<BlockSize, VecSize>(scores, length, [&](uint32_t, float score) {
const uint32_t ordered = convert_to_uint32_v2(score);
if ((ordered & prefix_mask) == prefix) {
atomicAdd(&smem->histogram[(ordered >> shift) & digit_mask], 1);
if constexpr (FinalCandidateCapacity > 0) {
if (final_round) {
if (tx == 0) {
smem->output_counter = 0;
smem->candidate_count = 0;
}
__syncthreads();
}
});
}
scan_scores<BlockSize, VecSize>(
scores, length, [&](uint32_t idx, float score) {
const uint32_t ordered = convert_to_uint32_v2(score);
const uint32_t ordered_prefix = ordered & prefix_mask;
if (ordered_prefix == prefix) {
atomicAdd(&smem->histogram[(ordered >> shift) & digit_mask], 1);
if constexpr (FinalCandidateCapacity > 0) {
if (final_round) {
const uint32_t pos = atomicAdd(&smem->candidate_count, 1);
if (pos < FinalCandidateCapacity) {
final_candidates[pos] = static_cast<int32_t>(idx);
}
}
}
} else if constexpr (FinalCandidateCapacity > 0) {
if (final_round && ordered_prefix > prefix) {
const uint32_t pos = atomicAdd(&smem->output_counter, 1);
output[pos] = static_cast<int32_t>(idx);
}
}
});
__syncthreads();
constexpr uint32_t kBinsPerThread = kHistogramBins / BlockSize;
@@ -119,12 +160,19 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
local_sum += counts[i];
}
uint32_t lower_prefix;
uint32_t total;
// Snapshot before the block scan; its barrier precedes any update.
// Shuffle scan: lighter on registers than a CUB block scan. Snapshot
// remaining first; the barrier below precedes any update.
const uint32_t remaining = smem->remaining;
BlockScan(smem->scan).ExclusiveSum(local_sum, lower_prefix, total);
uint32_t count_above = total - lower_prefix - local_sum;
const uint32_t lane = tx % kWarpSize;
const uint32_t warp = tx / kWarpSize;
const uint32_t warp_inclusive = warp_inclusive_sum(lane, local_sum);
if (lane == kWarpSize - 1) smem->warp_sum[warp] = warp_inclusive;
__syncthreads();
const uint32_t warp_total = smem->warp_sum[lane];
uint32_t count_above =
warp_reduce_sum_full(lane > warp ? warp_total : 0) +
__shfl_sync(0xFFFFFFFFu, warp_inclusive, kWarpSize - 1) -
warp_inclusive;
#pragma unroll
for (int i = static_cast<int>(kBinsPerThread) - 1; i >= 0; --i) {
const uint32_t count = counts[i];
@@ -132,6 +180,7 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
const uint32_t bin = tx * kBinsPerThread + i;
smem->remaining = remaining - count_above;
smem->prefix = prefix | (bin << shift);
if constexpr (FinalCandidateCapacity > 0) smem->pivot_count = count;
}
count_above += count;
}
@@ -140,26 +189,54 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores,
const uint32_t pivot = smem->prefix;
if (tx == 0) {
smem->output_counter = 0;
if constexpr (FinalCandidateCapacity == 0) smem->output_counter = 0;
smem->histogram[0] = 0;
}
__syncthreads();
const uint32_t equal_count = smem->remaining;
const uint32_t equal_base = TopK - equal_count;
scan_scores<BlockSize, VecSize>(
scores, length, [&](uint32_t idx, float score) {
const uint32_t ordered = convert_to_uint32_v2(score);
if (ordered > pivot) {
const uint32_t pos = atomicAdd(&smem->output_counter, 1);
output[pos] = static_cast<int32_t>(idx);
} else if (ordered == pivot) {
const uint32_t pos = atomicAdd(&smem->histogram[0], 1);
if (pos < equal_count) {
output[equal_base + pos] = static_cast<int32_t>(idx);
}
}
});
const auto collect = [&](uint32_t idx, float score) {
const uint32_t ordered = convert_to_uint32_v2(score);
if (ordered > pivot) {
const uint32_t pos = atomicAdd(&smem->output_counter, 1);
output[pos] = static_cast<int32_t>(idx);
} else if (ordered == pivot) {
const uint32_t pos = atomicAdd(&smem->histogram[0], 1);
if (pos < equal_count) {
output[equal_base + pos] = static_cast<int32_t>(idx);
}
}
};
if constexpr (FinalCandidateCapacity > 0) {
// Keys above the final round's prefix are already in the output.
const uint32_t candidate_count = smem->candidate_count;
if (candidate_count <= FinalCandidateCapacity) {
for (uint32_t i = tx; i < candidate_count; i += BlockSize) {
const uint32_t idx = final_candidates[i];
collect(idx, scores[idx]);
}
} else if (smem->pivot_count == candidate_count) {
// Every key under the final prefix equals the pivot (heavy ties), so
// any equal_count staged candidates complete the output.
for (uint32_t i = tx; i < equal_count; i += BlockSize) {
output[equal_base + i] = final_candidates[i];
}
} else {
constexpr uint32_t kFinalPrefixMask =
Radix::prefix_mask(Radix::kRounds - 1);
scan_scores<BlockSize, VecSize>(
scores, length, [&](uint32_t idx, float score) {
if ((convert_to_uint32_v2(score) & kFinalPrefixMask) ==
(pivot & kFinalPrefixMask)) {
collect(idx, score);
}
});
}
} else {
scan_scores<BlockSize, VecSize>(scores, length, collect);
}
__syncthreads();
}
@@ -191,6 +268,189 @@ coarse_refine_range(uint32_t coarse_bin) {
.finite = (lower_half_bits & 0x7C00u) != 0x7C00u};
}
constexpr uint32_t kCoarseRefineBins = 2048;
constexpr uint32_t kCoarseRefineDigitBits = 11;
using CoarseRefineScan = cub::BlockScan<uint32_t, kBlockSize>;
struct CoarseRefineSmem {
uint32_t histogram[kCoarseRefineBins];
CoarseRefineScan::TempStorage scan;
uint32_t prefix;
uint32_t remaining;
uint32_t output_counter;
uint32_t equal_counter;
uint32_t valid;
uint32_t key_min;
uint32_t key_max;
};
__device__ __forceinline__ uint32_t warp_reduce_min_full(uint32_t v) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
return __reduce_min_sync(0xFFFFFFFFu, v);
#else
#pragma unroll
for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) {
v = min(v, __shfl_xor_sync(0xFFFFFFFFu, v, mask));
}
return v;
#endif
}
__device__ __forceinline__ uint32_t warp_reduce_max_full(uint32_t v) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
return __reduce_max_sync(0xFFFFFFFFu, v);
#else
#pragma unroll
for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) {
v = max(v, __shfl_xor_sync(0xFFFFFFFFu, v, mask));
}
return v;
#endif
}
__device__ __forceinline__ void select_coarse_refine_digit(
CoarseRefineSmem* smem, uint32_t prefix, uint32_t shift) {
const uint32_t bin0 = 2 * threadIdx.x;
const uint32_t count0 = smem->histogram[bin0];
const uint32_t count1 = smem->histogram[bin0 + 1];
const uint32_t local_sum = count0 + count1;
const uint32_t remaining = smem->remaining;
uint32_t lower_prefix;
uint32_t total;
CoarseRefineScan(smem->scan).ExclusiveSum(local_sum, lower_prefix, total);
uint32_t count_above = total - lower_prefix - local_sum;
if (count_above < remaining && count_above + count1 >= remaining) {
smem->remaining = remaining - count_above;
smem->prefix = prefix | ((bin0 + 1) << shift);
}
count_above += count1;
if (count_above < remaining && count_above + count0 >= remaining) {
smem->remaining = remaining - count_above;
smem->prefix = prefix | (bin0 << shift);
}
}
// Refine an overflowing FP16 histogram bin from values already resident in
// registers. A 12-bit FP16 bin spans at most 22 ordered-FP32 bits for finite
// values, so two radix-2048 passes select the exact pivot. A bin holding a
// single key (e.g. all-equal scores) needs no refinement: that key is the
// pivot, so the first pass also tracks the bin's min and max key.
template <uint32_t TopK, uint32_t HistBits, uint32_t VecsPerThread>
__device__ __forceinline__ bool exact_topk_refine_registers(
const float4 (&vecs)[VecsPerThread], int32_t* __restrict__ output,
uint32_t length, uint32_t coarse_bin, uint32_t num_above, void* _smem) {
constexpr uint32_t kBins = kCoarseRefineBins;
constexpr uint32_t kDigitBits = kCoarseRefineDigitBits;
constexpr uint32_t kDigitMask = kBins - 1;
constexpr uint32_t kDeltaLimit = 1u << (2 * kDigitBits);
static_assert(HistBits == 12);
static_assert(kBins == 2 * kBlockSize);
auto* smem = static_cast<CoarseRefineSmem*>(_smem);
const uint32_t tx = threadIdx.x;
const auto range = coarse_refine_range<HistBits>(coarse_bin);
const uint32_t base_key = range.base_key;
if (tx == 0) {
smem->prefix = 0;
smem->remaining = TopK - num_above;
smem->valid = range.finite;
smem->key_min = 0xFFFFFFFFu;
smem->key_max = 0;
}
__syncthreads();
#pragma unroll
for (uint32_t round = 0; round < 2; ++round) {
for (uint32_t bin = tx; bin < kBins; bin += kBlockSize) {
smem->histogram[bin] = 0;
}
__syncthreads();
const uint32_t prefix = smem->prefix;
const uint32_t shift = round == 0 ? kDigitBits : 0;
uint32_t key_min = 0xFFFFFFFFu;
uint32_t key_max = 0;
bool done = false;
#pragma unroll
for (uint32_t v = 0; v < VecsPerThread && !done; ++v) {
const float* values = reinterpret_cast<const float*>(&vecs[v]);
#pragma unroll
for (uint32_t e = 0; e < 4 && !done; ++e) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if (idx >= length) {
done = true;
} else if (extract_coarse_bin_N<HistBits>(values[e]) == coarse_bin) {
const uint32_t key = convert_to_uint32_v2(values[e]);
const uint32_t delta = key - base_key;
if (round == 0) {
key_min = min(key_min, key);
key_max = max(key_max, key);
}
if (key < base_key || delta >= kDeltaLimit) {
atomicExch(&smem->valid, 0);
} else if (round == 0 || (delta & ~kDigitMask) == prefix) {
atomicAdd(&smem->histogram[(delta >> shift) & kDigitMask], 1);
}
}
}
}
if (round == 0) {
key_min = warp_reduce_min_full(key_min);
key_max = warp_reduce_max_full(key_max);
if (tx % kWarpSize == 0) {
atomicMin(&smem->key_min, key_min);
atomicMax(&smem->key_max, key_max);
}
}
__syncthreads();
if (smem->valid == 0) return false;
if (round == 0 && smem->key_min == smem->key_max) {
if (tx == 0) smem->prefix = smem->key_min - base_key;
__syncthreads();
break;
}
select_coarse_refine_digit(smem, prefix, shift);
__syncthreads();
}
if (tx == 0) {
smem->output_counter = 0;
smem->equal_counter = 0;
}
__syncthreads();
const uint32_t pivot = base_key + smem->prefix;
const uint32_t equal_base = TopK - smem->remaining;
bool done = false;
#pragma unroll
for (uint32_t v = 0; v < VecsPerThread && !done; ++v) {
const float* values = reinterpret_cast<const float*>(&vecs[v]);
#pragma unroll
for (uint32_t e = 0; e < 4 && !done; ++e) {
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
if (idx >= length) {
done = true;
} else {
const uint32_t key = convert_to_uint32_v2(values[e]);
if (key > pivot) {
const uint32_t pos = atomicAdd(&smem->output_counter, 1);
output[pos] = static_cast<int32_t>(idx);
} else if (key == pivot) {
const uint32_t pos = atomicAdd(&smem->equal_counter, 1);
if (equal_base + pos < TopK) {
output[equal_base + pos] = static_cast<int32_t>(idx);
}
}
}
}
}
__syncthreads();
return true;
}
} // namespace topk_histogram_4096
} // namespace vllm
csrc/torch_bindings.cpp
+8 -4
View File
@@ -151,7 +151,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
" Tensor! input_permutation, "
" Tensor! output_permutation, int num_experts, "
" int n, int k, Tensor? blockscale_offsets, "
" bool is_gated) -> ()");
" bool is_gated, "
" Tensor!? sorted_expert_ids=None) -> ()");
// compute per-expert problem sizes from expert_first_token_offset
// produced by vLLM's moe_permute kernel
@@ -235,14 +236,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def(
"mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
"Tensor output_scale_offset_by_experts, int n_experts, "
"Tensor? sorted_expert_ids=None) -> ()");
// Fused SiLU+Mul+MXFP4 experts quantization.
ops.def(
"silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
"output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
"Tensor output_scale_offset_by_experts, int n_experts, "
"Tensor? sorted_expert_ids=None) -> ()");
// Fused SiLU+Mul+NVFP4 quantization.
ops.def(
@@ -1021,7 +1024,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
" Tensor! kv_cache,"
" Tensor slot_mapping,"
" str kv_cache_dtype,"
" Tensor scale) -> ()");
" Tensor scale,"
" bool skip_scale=False) -> ()");
// Grouped concat_and_cache_mla across all layers. Each layer's cache base
// pointer and optional plain-FP8 scale are read from device tensors.
docker/Dockerfile
+2 -1
View File
@@ -31,7 +31,8 @@ ARG TRITON_INSTALL_FROM_SOURCE_REVISION=
# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend).
# This version is used for CUDA 12+ builds to keep the DeepEP build and
# runtime NCCL versions compatible.
ARG NCCL_VERSION=2.30.7
# DeepEPv2 GDA over EFA requires NCCL 2.31.x or later.
ARG NCCL_VERSION=2.31.2
# By parameterizing the base images, we allow third-party to use their own
# base images. One use case is hermetic builds with base images stored in
docker/Dockerfile.efa
+156
View File
@@ -0,0 +1,156 @@
# AWS EFA variant of a vllm-openai image. Adds the EFA userspace stack
# (libfabric, aws-ofi-nccl, hwloc), the EFA Mooncake wheel (x86_64 only),
# and DeepEP rebuilt from the amazon-contributing fork for NCCL GIN over EFA.
#
# The base must be built with INSTALL_KV_CONNECTORS=true (NIXL and Mooncake):
# docker build -f docker/Dockerfile.efa \
# --build-arg BASE_IMAGE=vllm/vllm-openai:latest -t vllm:openai-efa docker/
#
# The host still provides the EFA kernel module and memlock limits.
ARG BASE_IMAGE=vllm/vllm-openai:latest
FROM ${BASE_IMAGE}
ARG TARGETARCH
ARG EFA_INSTALLER_VERSION=1.51.0
ARG EFA_INSTALLER_SHA256=83993044fa4df3dc2f8b5c2f7051d61b9689441be023a2c92ea51d3f53f2465c
# Without the KV connectors there is no NIXL plugin or Mooncake wheel to use EFA.
RUN set -eu; \
if ! python3 -c "import importlib.metadata as m; m.version('nixl')" \
2>/dev/null; then \
echo "Dockerfile.efa needs a base image built with \
--build-arg INSTALL_KV_CONNECTORS=true" >&2; \
exit 1; \
fi
# Install the AWS EFA userspace stack.
RUN set -eux; \
EFA_INSTALLER_ARCHIVE="/tmp/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz"; \
curl -fsSL --retry 3 --retry-all-errors \
"https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz" \
-o "${EFA_INSTALLER_ARCHIVE}"; \
echo "${EFA_INSTALLER_SHA256} ${EFA_INSTALLER_ARCHIVE}" | sha256sum -c -; \
tar -xzf "${EFA_INSTALLER_ARCHIVE}" -C /tmp; \
# The installer apt-installs dependencies; the base image purged the lists.
apt-get update -y; \
# Minimal install: libfabric, aws-ofi-nccl and hwloc only. --build-ngc
# implies --skip-kmod --skip-limit-conf --no-verify.
(cd /tmp/aws-efa-installer && \
./efa_installer.sh --skip-mpi --skip-rdma-core --build-ngc); \
test -x /opt/amazon/efa/bin/fi_info; \
/opt/amazon/efa/bin/fi_info --version; \
# `-l` lists providers without EFA hardware, so this works in a build.
/opt/amazon/efa/bin/fi_info -l | grep -q '^efa:'; \
EFA_NIXL_PLUGIN=$(find /usr/local/lib -name libplugin_LIBFABRIC.so \
-print -quit); \
test -n "${EFA_NIXL_PLUGIN}"; \
ldd "${EFA_NIXL_PLUGIN}" | grep -q 'libfabric\.so\.1 => /opt/amazon/efa/lib/'; \
# Plugins that NCCL_NET_PLUGIN=ofi and NCCL_TUNER_PLUGIN=ofi resolve to.
ldconfig -p | grep -q '/opt/amazon/ofi-nccl/lib/libnccl-net-ofi\.so$'; \
ldconfig -p | grep -q '/opt/amazon/ofi-nccl/lib/libnccl-tuner-ofi\.so$'; \
apt-get clean; \
rm -rf "${EFA_INSTALLER_ARCHIVE}" /tmp/aws-efa-installer \
/var/lib/apt/lists/*
# Route NCCL over EFA. Override with `-e NCCL_NET_PLUGIN= -e NCCL_TUNER_PLUGIN=`
# to run this image on a non-EFA host.
ENV NCCL_NET_PLUGIN=ofi \
NCCL_TUNER_PLUGIN=ofi
# Swap in the EFA Mooncake wheel at the base image's Mooncake version. Both
# wheels own the `mooncake` module; the EFA wheels are published for x86_64 only.
RUN --mount=type=cache,target=/opt/uv/cache \
set -eu; \
pkg_version() { python3 -c "import importlib.metadata as m, sys; \
print(m.version(sys.argv[1]))" "$1" 2>/dev/null; }; \
if MOONCAKE_VERSION=$(pkg_version mooncake-transfer-engine-cuda13); then \
MOONCAKE_PKG=mooncake-transfer-engine-cuda13; \
MOONCAKE_EFA_PKG=mooncake-transfer-engine-efa-cuda13; \
elif MOONCAKE_VERSION=$(pkg_version mooncake-transfer-engine); then \
MOONCAKE_PKG=mooncake-transfer-engine; \
MOONCAKE_EFA_PKG=mooncake-transfer-engine-efa; \
else \
echo "No Mooncake wheel in the base image; build it with \
--build-arg INSTALL_KV_CONNECTORS=true" >&2; \
exit 1; \
fi; \
if [ "${TARGETARCH}" != "amd64" ]; then \
echo "TARGETARCH=${TARGETARCH}: ${MOONCAKE_EFA_PKG} is x86_64-only, \
keeping ${MOONCAKE_PKG} (Mooncake EFA protocol unavailable)"; \
else \
echo "Swapping ${MOONCAKE_PKG}==${MOONCAKE_VERSION} for \
${MOONCAKE_EFA_PKG}==${MOONCAKE_VERSION}"; \
uv pip uninstall --system "${MOONCAKE_PKG}"; \
uv pip install --system --no-deps \
"${MOONCAKE_EFA_PKG}==${MOONCAKE_VERSION}"; \
test "$(pkg_version "${MOONCAKE_EFA_PKG}")" = "${MOONCAKE_VERSION}"; \
# `! cmd` is exempt from `set -e`, so check explicitly.
if pkg_version "${MOONCAKE_PKG}" >/dev/null; then \
echo "${MOONCAKE_PKG} is still installed" >&2; \
exit 1; \
fi; \
fi
# Replace the base image's upstream DeepEP with the amazon-contributing fork,
# whose deepep_v2 backend carries EP traffic over EFA through NCCL GIN.
ARG DEEPEP_REPO="https://github.com/amazon-contributing/DeepEP.git"
ARG DEEPEP_REF="874779c9ccd2294b56304bd6cc5f138f1f71d097"
ARG DEEPEP_CUDA_ARCH_LIST="9.0a;10.0a;10.3a"
# Headers and git for DeepEP's nvcc build. `reinstall` also installs packages
# that are absent.
RUN set -eux; \
CUDA_PKG_VERSION=$(nvcc --version \
| grep -oP 'release \K[0-9]+\.[0-9]+' | tr '.' '-'); \
DEBIAN_FRONTEND=noninteractive apt-get update -y; \
DEBIAN_FRONTEND=noninteractive apt-get reinstall -y --no-install-recommends \
git \
libcusparse-dev-${CUDA_PKG_VERSION} \
libcusolver-dev-${CUDA_PKG_VERSION} \
cuda-nvrtc-dev-${CUDA_PKG_VERSION}; \
rm -rf /var/lib/apt/lists/*
RUN --mount=type=cache,target=/opt/uv/cache \
set -eux; \
git clone "${DEEPEP_REPO}" /tmp/deepep; \
cd /tmp/deepep; \
git checkout "${DEEPEP_REF}"; \
git submodule update --init --recursive; \
uv pip uninstall --system deep_ep 2>/dev/null || true; \
NCCL_ROOT="$(python3 -c "import importlib.util, os; \
s = importlib.util.find_spec('nvidia.nccl'); \
print(os.path.dirname(s.origin or s.submodule_search_locations[0]))")/nccl"; \
export EP_NCCL_ROOT_DIR="${NCCL_ROOT}"; \
export LD_LIBRARY_PATH="${NCCL_ROOT}/lib:${LD_LIBRARY_PATH:-}"; \
# The CUDA driver stubs let DeepEP link libcuda in a builder with no driver.
export LIBRARY_PATH="${NCCL_ROOT}/lib:/usr/local/cuda/lib64/stubs:\
/usr/local/cuda/targets/x86_64-linux/lib/stubs:${LIBRARY_PATH:-}"; \
export DISABLE_AGGRESSIVE_PTX_INSTRS=1; \
export TORCH_CUDA_ARCH_LIST="${DEEPEP_CUDA_ARCH_LIST}"; \
uv pip install --system --no-deps -vv --no-build-isolation .; \
cd /; \
rm -rf /tmp/deepep; \
# vLLM only logs when the deepep_v2 backend is unavailable, so fail here.
python3 -c "import deep_ep; assert hasattr(deep_ep, 'ElasticBuffer'), \
'DeepEP built without the ElasticBuffer (v2) API'"
# Provenance. Build args from docker-build-metadata-args.sh override the values
# inherited from the base image; without them the base's values are kept.
ARG BASE_IMAGE
ARG VLLM_BUILD_COMMIT
ARG VLLM_BUILD_PIPELINE
ARG VLLM_BUILD_URL
ARG VLLM_IMAGE_TAG
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT} \
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE} \
VLLM_BUILD_URL=${VLLM_BUILD_URL} \
VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG}
LABEL org.opencontainers.image.revision="${VLLM_BUILD_COMMIT}" \
org.opencontainers.image.version="${VLLM_IMAGE_TAG}" \
org.opencontainers.image.url="${VLLM_BUILD_URL}" \
ai.vllm.build.commit="${VLLM_BUILD_COMMIT}" \
ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}" \
ai.vllm.build.url="${VLLM_BUILD_URL}" \
ai.vllm.image.tag="${VLLM_IMAGE_TAG}" \
ai.vllm.efa.base-image="${BASE_IMAGE}" \
ai.vllm.efa.installer-version="${EFA_INSTALLER_VERSION}"
docker/Dockerfile.rocm
+22 -3
View File
@@ -338,6 +338,23 @@ RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproj
&& echo "=== ccache stats after ROCm native build ===" \
&& (ccache --show-stats || true)
FROM csrc-build AS kernel-symbol-map-build
ARG VLLM_KERNEL_SYMBOL_MAP=0
ARG buildkite_commit=""
COPY tools/ci/kernel_symbol_map.py tools/ci/rocm_kernel_symbol_map.py tools/ci/
RUN if [ "${VLLM_KERNEL_SYMBOL_MAP}" = "1" ]; then \
uv venv --system-site-packages /tmp/kernel-map/.venv \
&& /tmp/kernel-map/.venv/bin/python tools/ci/kernel_symbol_map.py \
--backend rocm --build-root build --source-root . \
--commit "${buildkite_commit}" \
--out dist/kernel_symbol_map.rocm.json.gz \
&& rm -rf /tmp/kernel-map; \
fi
FROM scratch AS kernel-symbol-map-rocm
ARG COMMON_WORKDIR
COPY --from=kernel-symbol-map-build ${COMMON_WORKDIR}/vllm/dist/kernel_symbol_map.rocm.json.gz /
# Build the full vLLM ROCm wheel by reusing the native extension wheel from
# csrc-build. This stage still rebuilds for Python/package changes, but skips
# the expensive dependency installation and HIP/C++ compilation when their
@@ -370,6 +387,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
# tests/tools/test_kernel_symbol_map.py loads the symbol-map producers from here.
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/ci /tools/ci
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
@@ -387,7 +406,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
FROM base AS build_nixl
ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV UCX_HOME=/usr/local/ucx
ENV NIXL_HOME=/usr/local/nixl
@@ -842,8 +861,8 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -978,8 +997,8 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
ENV TOKENIZERS_PARALLELISM=false
docker/Dockerfile.rocm_72
+3 -3
View File
@@ -384,7 +384,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
FROM base AS build_nixl
ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
@@ -862,8 +862,8 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -978,8 +978,8 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf
ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf
ENV TOKENIZERS_PARALLELISM=false
docker/ci-rocm.hcl
+13
View File
@@ -25,6 +25,10 @@ variable "BUILDKITE_BUILD_ID" {
default = ""
}
variable "VLLM_KERNEL_SYMBOL_MAP" {
default = "0"
}
variable "PARENT_COMMIT" {
default = ""
}
@@ -267,6 +271,8 @@ target "_ci-rocm" {
CI_BASE_IMAGE = CI_BASE_IMAGE
ROCM_SMOKE_ID = BUILDKITE_BUILD_ID
max_jobs = CI_MAX_JOBS
VLLM_KERNEL_SYMBOL_MAP = VLLM_KERNEL_SYMBOL_MAP
buildkite_commit = BUILDKITE_COMMIT
}
}
@@ -302,6 +308,13 @@ target "csrc-rocm-ci" {
output = ["type=cacheonly"]
}
target "kernel-symbol-map-rocm" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "kernel-symbol-map-rocm"
cache-from = get_cache_from_rocm_csrc()
output = ["type=local,dest=./kernel-symbol-map-rocm"]
}
# Cache-only target for the Rust frontend build stage. Final-image cache
# exports use mode=min and do not reliably persist intermediate cargo layers,
# so Rust gets its own source-scoped cache target.
docker/versions.json
+1 -1
View File
@@ -14,7 +14,7 @@
"default": "false"
},
"NCCL_VERSION": {
"default": "2.30.7"
"default": "2.31.2"
},
"BUILD_BASE_IMAGE": {
"default": "pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359"
docs/benchmarking/cli.md
+19
View File
@@ -9,6 +9,25 @@ It's a living document, updated as new features and datasets become available.
For benchmarking production vLLM servers, we recommend [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns.
## Model Forward Benchmark
For exact-batch model timing, including steady-state pipeline-parallel forwards,
see [vllm bench forward](../cli/README.md#forward).
## Attention Kernel Benchmark
Run the attention kernel smoke benchmark from the repository root using the
exact backend names:
```bash
python3 benchmarks/attention_benchmarks/benchmark.py \
--backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k"
```
If any configuration fails, the benchmark exits with status `1` after reporting
results and writing any requested CSV or JSON output. Successful runs exit with
status `0`. Profiling runs also propagate failures, but do not save timing metrics.
## Dataset Overview
<style>
docs/cli/README.md
+42 -1
View File
@@ -136,7 +136,7 @@ vllm bench {forward, latency, serve, throughput}
### forward
Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; sampling and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP. CUDA only for now. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case.
Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; target sampling, rejection sampling, and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP/PCP/PP. CUDA only for now. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case.
```bash
vllm bench forward \
@@ -145,6 +145,47 @@ vllm bench forward \
--case mix=2048,32*1@4096
```
PCP uses `--prefill-context-parallel-size` with the model runner's supported
MLA models. Cases describe the global batch per DP rank, before PCP splits it;
`num_local_tokens`, `num_tokens_after_padding`, and `batch_order` in each rank's
JSON result describe the local partition (a request can appear twice).
Latency is the maximum across all ranks. For example, with four GPUs:
```bash
torchrun --nproc-per-node=4 -m vllm.entrypoints.cli.main bench forward \
--model MODEL --tensor-parallel-size 2 --prefill-context-parallel-size 2 \
--case prefill=8192 --case mix=2048,8*1@4096
```
With `--pipeline-parallel-size`, the benchmark streams identical microbatches
through captured stages. `median_us` measures each microbatch from its first
stage's forward start to its last stage's forward completion, including PP
transfers while other microbatches are in flight. Pipeline fill and drain are
excluded. This is latency under steady-state load, not the completion interval.
Per-rank JSON includes stage timelines and the estimated clock synchronization
uncertainty in microseconds.
```bash
torchrun --nproc-per-node=4 -m vllm.entrypoints.cli.main bench forward \
--model MODEL --pipeline-parallel-size 2 --tensor-parallel-size 2 \
--case prefill8x1k=8*1024 --case prefill1x8k=8192
```
With `--speculative-config`, timing includes the target forward and the MTP,
EAGLE, DSpark, or DFlash proposal, including the drafter's token sampling.
Use `--exclude-drafter` to measure only the target forward. Verification assumes
all draft tokens are accepted and uses a fixed bonus token; target sampling and
rejection sampling are excluded. Both paths use CUDA graphs and restore the
batch's state before each replay.
Balanced MoE routing fixtures apply to the target model; the drafter uses its
model routing, so use real weights when measuring a MoE drafter.
```bash
vllm bench forward --model Qwen/Qwen3-1.7B \
--speculative-config '{"model":"AngelSlim/Qwen3-1.7B_eagle3","method":"eagle3","num_speculative_tokens":3}' \
--case verify=16*4@4096:3
```
See [vllm bench forward](./bench/forward.md) for the full reference of all available arguments.
### latency
docs/contributing/dockerfile/dockerfile.md
+19
View File
@@ -34,6 +34,25 @@ EP extension for the installed PyTorch, with `tools/build_nixl_with_torch215_for
`vllm-openai-base` then replaces the NIXL packages installed from the KV-connector
requirements with these wheels.
The AWS EFA images are built separately by `docker/Dockerfile.efa` on top of a
`vllm-openai` image built with `INSTALL_KV_CONNECTORS=true`. It adds a minimal
AWS EFA userspace install (libfabric, aws-ofi-nccl and hwloc), the EFA Mooncake
wheel (x86_64 only), and DeepEP built for NCCL GIN over EFA. Kernel modules and
host limits remain host-managed.
```bash
docker build -f docker/Dockerfile.efa \
--build-arg BASE_IMAGE=vllm/vllm-openai:latest -t vllm:openai-efa docker/
```
The image does not change vLLM's Mooncake protocol default (`rdma`). To use
EFA, set it per connector:
- `MooncakeConnector`: `"mooncake_protocol": "efa"` in `kv_connector_extra_config`.
- `MooncakeStoreConnector`: `"protocol": "efa"` in the JSON file at `MOONCAKE_CONFIG_PATH`.
The EFA Mooncake wheel is x86_64-only, so use `tcp` on arm64.
> <figure markdown="span">
> ![](../../assets/contributing/dockerfile-stages-dependency.png){ align="center" alt="query" width="100%" }
> </figure>
docs/contributing/profiling.md
+5 -1
View File
@@ -109,7 +109,11 @@ curl -X POST http://localhost:8000/start_profile \
-d '{
"profile_prefix": "sharegpt_run_1",
"delay_iterations": 5000,
"max_iterations": 20
"max_iterations": 20,
"profiler_kwargs": {
"record_shapes": true,
"with_stack": true
}
}'
```
docs/design/rocm_octave.md
+197
View File
@@ -0,0 +1,197 @@
# Octave KV: a 3-bit KV cache for ROCm
`ROCM_OCTAVE` is an attention backend for AMD GPUs that keeps the KV cache of
full-attention layers at 3-4 bits per value, with native kernels (HIP today). It is
selected with `--kv-cache-dtype octave_k3v3_compact`, `octave_k3v3` or
`octave_k3v4` (the backend is picked from the cache dtype;
`--attention-backend ROCM_OCTAVE` also works). The name comes from the 3-bit
codes: each one indexes an eight-level codebook.
## Why it exists
Long-context serving on AMD GPUs is bounded by KV-cache memory long before it
is bounded by compute. The existing options on ROCm were:
- fp16/bf16 KV (`TRITON_ATTN`, `ROCM_ATTN`): 16 bits per value.
- fp8 / int8 per-token-head: 8 bits per value.
- `TURBOQUANT`: 3-4 bits per value, but its kernels are Triton (store and
decode) and its prefill goes through FlashAttention; it does not support the
fused multi-step MTP draft.
Octave KV targets the same bit budget as TurboQuant with HIP kernels written for
the AMD matrix instructions (WMMA on RDNA3, MFMA on CDNA as a contract slot),
reading the packed cache directly in both decode and prefill, and with
multi-token verification (MTP / speculative decoding) as a first-class case.
## Supported models
The format exploits *partial rotary embeddings*: models whose attention heads
apply RoPE to a fraction of the head dims. Today the kernels are built for
head size 256 with 64 rotary dims (`partial_rotary_factor = 0.25`), the
attention shape of the Qwen3-Next / Qwen3.5 / Qwen3.6 / Qwen3.8 family (dense
and MoE, e.g. Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.5-122B-A10B). Like the MLA
backends, which serve only models with multi-head latent attention, Octave KV is
tied to an attention shape, not to a model.
Models with full RoPE (Llama, Qwen3, Mistral, ...) or other head sizes are
rejected at startup with an explicit error. Sliding window, ALiBi, attention
sinks, soft-capping and encoder attention are not supported.
The decode kernels cover GQA groups of 1, 2, 4, 6 or 8 query heads per KV
head on each GPU. Under tensor parallelism the group is counted per GPU, so
models with 32 query heads and 2 KV heads (Qwen3.5-122B-A10B) need TP ≥ 4.
## Format
One slot per (token, KV head). V is always rotated as one 256-dim block
(random signs, then a Walsh-Hadamard transform) and gets 3 or 4 bits per dim.
K has two layouts:
- **Block K** (`octave_k3v3`, `octave_k3v4`): four 64-dim blocks, each rotated
on its own. The first block holds exactly the 64 rotary dims and gets 4 bits
per dim, the three NoPE blocks 3 bits; each block has its own scale.
- **Compact K** (`octave_k3v3_compact`): rotated over all 256 dims like V,
3 bits per dim, one scale. The rotation spreads the rotary dims over the
whole head, so they need no extra bits; slots pad to 4 bytes instead of 8.
| Bytes (block K) | Content |
| --- | --- |
| `[0, 32)` | K RoPE block, 4-bit codes |
| `[32, 104)` | K NoPE blocks, 3-bit codes |
| next `256·v/8` | V, `v`-bit codes |
| next 10 | fp16 scales (K blocks 0-3, V) |
| Bytes (compact K) | Content |
| --- | --- |
| `[0, 96)` | K, 3-bit codes |
| `[96, 192)` | V, 3-bit codes |
| `[192, 196)` | fp16 scales (K, V) |
| Cache dtype | Slot bytes | Bits per value |
| --- | --- | --- |
| `octave_k3v3_compact` | 196 | 3.06 |
| `octave_k3v3` | 216 | 3.38 |
| `octave_k3v4` | 248 | 3.88 |
For reference, `turboquant_3bit_nc` uses 198 bytes and `turboquant_k3v4_nc`
230.
After the rotation the coordinates are close to Gaussian, so each code
indexes the Lloyd-Max codebook of N(0, 1) for its width, stored as int8
(`LUT` in the reference): a code `c` of a block with scale `s` stands for
`LUT[c] · s / 64`. The scale is chosen so that `x · x̂ = |x|²`. A
least-squares fit minimizes the reconstruction error but shrinks every
reconstructed vector (by about 3% at 3 bits), which flattens the softmax and
scales down the attention output; at 380k tokens it cost 2-3× more NLL than
this scale, at the same size. The quantizer is dynamic: no calibration, no per-model
dictionaries. The kernels map four codes to their int8 values with one
`v_perm_b32` (two for 4-bit codes), so the codebook costs nothing in
bandwidth and the dot products stay integer. The byte layout's source of
truth is the PyTorch reference in `vllm/v1/attention/ops/rocm_octave.py`.
Why these layouts: measured on real Q/K/V of Qwen3.8-27B (32k tokens, every
full-attention layer, attention-output error), both K layouts sit on the same
bytes-versus-error curve, about 25-30% below TurboQuant at equal size; block
K reaches the 4-bit rotary block, compact K the smallest slot. Rotating the
rotary block alone at 3 bits does not work (one layer reaches 32% error);
rotating it with the rest of the head does. Moving bits between K and V or
weighting K channels by query energy did not help at a fixed size. The int8
codebook loses nothing against the floating-point one.
Attention runs in the rotated space: the query is rotated like K and the
output is rotated back once per query. The rotations are orthogonal, so the
scores are those of the unrotated vectors.
## Kernels and data flow
| Step | Kernel | File |
| --- | --- | --- |
| Cache write | `octave_cache_store`: one wave per slot, rotate + quantize + pack | `csrc/attention/octave_attn.cu` |
| Decode / MTP verify (≤128 query tokens per request) | `octave_decode`: split-KV; WMMA tier on gfx11, portable dot4 tier elsewhere; consecutive query tokens of a request share each K/V read | `octave_decode_wmma_rdna3.cu`, `octave_attn.cu` |
| Split reduce + inverse V rotation | `reduce_kernel` | `octave_attn.cu` |
| Prefill | `octave_prefill`: phase 1 reads the cached prefix from the packed cache (int8 WMMA for QK with per-segment scales, fp16 WMMA for PV), phase 2 attends to the chunk's own unquantized K/V and merges | `octave_prefill_rdna3.cu` |
| Q/K/V chunk rotation for prefill | `octave_rotate` | `octave_attn.cu` |
| Codebooks (`v_perm_b32` tables) and byte layout | — | `octave_format.cuh` |
The metadata builder builds, on the device and without host syncs, the
per-query request and causal-length maps the kernels index by. They are what
makes MTP verification batches (several query tokens per request, each with
its own causal length) work without special cases, and they give cudagraph
padding queries a K length of 0.
## Architecture contract
Kernels are native: HIP or FlyDSL, never Triton. There is no fallback: an
architecture without a required kernel is rejected at startup with an error
naming what is missing. The kernels in this tree are HIP; a FlyDSL kernel
(for example a CDNA prefill) must keep the same signature and semantics, pass
the same tests, and report a clear error when `flydsl` is not installed.
Octave is enabled only on RDNA3 and CDNA3, validated on gfx1100 and gfx942
(MI300). The backend rejects any other family at selection
(`supports_compute_capability`, via `get_rdna_version()` and
`get_cdna_version()`), and the kernel tests skip elsewhere. CMake builds the
Octave sources and registers the ops only when the target list includes
gfx1100 or gfx942; other targets in the same build compile the
architecture-specific kernels as empty stubs, and the kernel dispatch reports
any other RDNA3 chip as not validated.
| Piece | gfx1100 (RDNA3) | gfx942 (MI300) |
| --- | --- | --- |
| int8 dot4 (`octave_arch.cuh`) | `v_dot4_i32_iu8` | `v_dot4_i32_i8` |
| Codebook lookup (`octave_format.cuh`) | `v_perm_b32` | `v_perm_b32` |
| Cache store, reduce, rotate | portable | portable (decode reduce: gfx942 kernel) |
| Decode | WMMA | MFMA (`v_mfma_i32_16x16x32_i8`) |
| Prefill | WMMA | MFMA |
The portable dot4 decode stays as the reference tier the tests compare the
native tiers against.
To add an architecture:
1. Implement `octave_prefill` for it, in HIP or FlyDSL (same signature and semantics as
`octave_prefill_rdna3.cu`: phase-1 partials over the packed prefix, phase-2
chunk attention and merge, rotated space in and out).
2. Add a decode tier, or validate the portable dot4 kernel there.
3. Add the architecture to the CMake guard, to `_kernel_tier` and
`supports_compute_capability` in
`vllm/v1/attention/backends/rocm_octave_attn.py`, and to the test gate.
4. Run `tests/kernels/attention/test_rocm_octave.py`: every kernel is checked
against the PyTorch reference, and `test_decode_matches_reference` runs
both the fast and the portable decode tier.
## Results
All measurements on 4× RX 7900 XTX (gfx1100), Qwen3.8-27B W4A16, tensor
parallel 4, MTP k=3, 8 GB of KV cache per GPU.
### Capacity
| KV cache | Tokens that fit |
| --- | --- |
| fp16 | TBD |
| `turboquant_k3v4_nc` | 1,790,625 |
| `octave_k4v4` | 1,402k |
| `octave_k3v3` | 1,679k |
### Speed
TBD: `vllm bench serve` (random prompts, no prefix reuse) TTFT and TPOT at
8k / 32k / 100k / 380k input, and per-layer kernel times from
`benchmarks/attention_benchmarks`.
### Quality
TBD: GSM8K (`tests/evals/gsm8k/gsm8k_eval.py`, 1319 questions, 5-shot,
temperature 0.6) and next-token NLL over a 380k-token document.
## Reproducing
```bash
vllm serve Qwen/Qwen3.8-27B --kv-cache-dtype octave_k3v4 ...
pytest tests/kernels/attention/test_rocm_octave.py
python benchmarks/attention_benchmarks/benchmark.py \
--backends ROCM_OCTAVE --kv-cache-dtype octave_k3v4 \
--model Qwen/Qwen3.8-27B --head-dim 256 --num-q-heads 6 --num-kv-heads 1 \
--batch-specs q1s100k q4s100k q2ks100k
```
Loaded 100 of 849 files, more files were not shown because too many files have changed in this diff. Show more