diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index f630f698a14..6da792aba4d 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -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 diff --git a/.buildkite/hardware_tests/ascend_npu.yaml b/.buildkite/hardware_tests/ascend_npu.yaml index 1ec6578bd83..f5bf6c53a89 100644 --- a/.buildkite/hardware_tests/ascend_npu.yaml +++ b/.buildkite/hardware_tests/ascend_npu.yaml @@ -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 diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 64d5a654f53..5321f285738 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -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 \ diff --git a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml index ed98d440851..96ac9304969 100644 --- a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml +++ b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml @@ -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: diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 0016023dc74..7ad8de9a22c 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -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"' diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index 92311638fcf..a9fa467caee 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -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 diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 0cdb7d78b66..34e9067a3ab 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -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 diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index c6a59ca99ed..48f84b34844 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -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"' diff --git a/.buildkite/intel_jobs/models_language_intel.yaml b/.buildkite/intel_jobs/models_language_intel.yaml new file mode 100644 index 00000000000..c956e0e1323 --- /dev/null +++ b/.buildkite/intel_jobs/models_language_intel.yaml @@ -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' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 74d6eee4900..c2c0428c000 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -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: diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 3bd86a92e98..2de3a4bfa74 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -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" diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 5935938616b..db3d0a4acfe 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -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 <&2 + fi } if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index 865e9957468..bdd24be4ee4 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -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 ---- diff --git a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh index 528247964ba..bed4a3e09ae 100755 --- a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh +++ b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh @@ -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}" diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index aa16302da8c..8f04c5c13c4 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -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}" diff --git a/.buildkite/scripts/upload-release-wheels-pypi.sh b/.buildkite/scripts/upload-release-wheels-pypi.sh index a45d2aa3c9f..2b5a4abbf86 100644 --- a/.buildkite/scripts/upload-release-wheels-pypi.sh +++ b/.buildkite/scripts/upload-release-wheels-pypi.sh @@ -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 diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 3628e4462bd..8ba09f55ae7 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -134,23 +134,24 @@ steps: commands: - pytest -v -s tools/test_docker_build_metadata_args.py -- label: ":amd: (MI300) Torch Stable ABI Audit" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - no_gpu: true - optional: 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. +# - label: ":amd: (MI300) Torch Stable ABI Audit" # TBD +# timeout_in_minutes: 35 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# no_gpu: true +# optional: 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 ######################################################################################################################################### # # @@ -207,22 +208,23 @@ steps: #---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------# -- label: ":amd: (MI250) Helion Kernels Shard %N" - timeout_in_minutes: 45 - parallelism: 2 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/utils/import_utils.py - - tools/benchmark_helion_kernels.py - - tests/kernels/helion/ - - vllm/platforms/rocm.py - commands: - - pip install helion==1.4.0 - - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# Disabled until helion transitive dependencies are locked. +# - label: ":amd: (MI250) Helion Kernels Shard %N" +# timeout_in_minutes: 45 +# parallelism: 2 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] +# dind: false +# agent_pool: mi250_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/utils/import_utils.py +# - tools/benchmark_helion_kernels.py +# - tests/kernels/helion/ +# - vllm/platforms/rocm.py +# commands: +# - pip install helion==1.4.0 +# - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT #----------------------------------------------------- mi250 · models / language -----------------------------------------------------# @@ -272,28 +274,29 @@ steps: #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# -- label: ":amd: (MI250) Batch Invariance" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - 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 - commands: - - export VLLM_GPU_MEMORY_UTILIZATION=0.8 - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - 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 +# Disabled until the pytest plugins installation is locked. +# - label: ":amd: (MI250) Batch Invariance" # TBD +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] +# dind: false +# agent_pool: mi250_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 +# commands: +# - export VLLM_GPU_MEMORY_UTILIZATION=0.8 +# - export VLLM_TEST_CLEAN_GPU_MEMORY=1 +# - 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 - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 @@ -857,51 +860,25 @@ steps: - pytest -v -s -x lora/test_mixtral.py - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' -- label: ":amd: (MI300) Distributed Torchrun + Examples" # TBD - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - tests/distributed/test_torchrun_example.py - - tests/distributed/test_torchrun_example_moe.py - - examples/rl/ - - examples/features/data_parallel/data_parallel_offline.py - - vllm/platforms/rocm.py - commands: - - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py - - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py - - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py - - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py - - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py - - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py - # 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 - # rlhf examples - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_nccl.py - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_ipc.py - -- label: ":amd: (MI300) Elastic EP Scaling" # TBD - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/v1/executor/ - - vllm/compilation/ - - tests/distributed/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - pytest -v -s distributed/test_elastic_ep.py +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI300) Elastic EP Scaling" # TBD +# timeout_in_minutes: 55 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/v1/executor/ +# - vllm/compilation/ +# - tests/distributed/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - pytest -v -s distributed/test_elastic_ep.py - label: ":amd: (MI300) RayExecutorV2" # TBD timeout_in_minutes: 75 @@ -925,25 +902,6 @@ steps: - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray" -- label: ":amd: (MI300) Distributed DP + EP" # TBD - timeout_in_minutes: 50 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_8 - num_gpus: 8 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - examples/features/torchrun/torchrun_dp_example_offline.py - - vllm/config/parallel.py - - vllm/distributed/ - - vllm/v1/engine/llm_engine.py - - vllm/v1/executor/uniproc_executor.py - - vllm/v1/worker/gpu_worker.py - - vllm/platforms/rocm.py - commands: - - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep - - label: ":amd: (MI300) Distributed Torchrun + Shutdown" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1027,21 +985,6 @@ steps: commands: - pytest -v -s entrypoints/openai/responses -- label: ":amd: (MI300) Entrypoints Integration (Speech to Text)" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - fast_check: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/entrypoints/speech_to_text - optional: true - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/speech_to_text - - label: ":amd: (MI300) Entrypoints Integration (Multimodal)" timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1135,24 +1078,6 @@ steps: commands: - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt -- label: ":amd: (MI300) LM Eval KV-Offload Medium" # TBD - timeout_in_minutes: 70 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py - - vllm/v1/kv_offload/ - - vllm/v1/simple_kv_offload/ - - tests/evals/gsm8k/test_gsm8k_offloading.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" - - label: ":amd: (MI300) LM Eval Small Models Harness" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1195,162 +1120,9 @@ steps: commands: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt -- label: ":amd: (MI300) LM Eval Small Models Distributed" - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt - optional: true - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - -- label: ":amd: (MI300) DeepSeek V2-Lite Sync EPLB Accuracy" # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/backends/mla/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 - -- label: ":amd: (MI300) LM Eval Large Models Harness" # TBD - timeout_in_minutes: 85 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 - -- label: ":amd: (MI300) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 - -- label: ":amd: (MI300) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 - # TODO(akaratza): Retire this standalone ROCm coverage after # tests/distributed/test_eplb_spec_decode.py is enabled and validated on ROCm. # The upstream parity job currently skips that entire module on ROCm. -- label: ":amd: (MI300) Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/spec_decode/ - - vllm/distributed/eplb - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - .buildkite/scripts/scheduled_integration_test/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 - -- label: ":amd: (MI300) LM Eval Large Models FP8" # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_USE_DEEP_GEMM=0 - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm-tp4.txt --tp-size=4 - #---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------# - label: ":amd: (MI300) Core Operation Kernels Shard %N" # TBD @@ -1601,9 +1373,12 @@ steps: - vllm/models/minimax_m3/amd/ - vllm/models/deepseek_v41/amd/mono/common/ - tests/models/minimax_m3/amd_specialized_decode/ + - tests/models/minimax_m3/test_index_cache_overrides.py + - vllm/model_executor/models/config.py - vllm/platforms/rocm.py commands: - pytest -v -s models/minimax_m3/amd_specialized_decode + - pytest -v -s models/minimax_m3/test_index_cache_overrides.py # the long TP4 hand-off stress - pytest -v -s --optional -m optional models/minimax_m3/amd_specialized_decode @@ -1949,20 +1724,21 @@ steps: - pytest -v -s plugins_tests/test_oot_registration_offline.py - pytest -v -s plugins_tests/lora_resolvers -- label: ":amd: (MI300) GGUF Plugin" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/gguf - - vllm/platforms/rocm.py - commands: - - pip install "vllm-gguf-plugin >= 0.0.2" - - pytest -v -s plugins_tests/gguf +# Disabled until the GGUF plugin installation is locked. +# - label: ":amd: (MI300) GGUF Plugin" # TBD +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/gguf +# - vllm/platforms/rocm.py +# commands: +# - pip install "vllm-gguf-plugin >= 0.0.2" +# - pytest -v -s plugins_tests/gguf #------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# @@ -2068,67 +1844,39 @@ steps: - 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" -- label: ":amd: (MI300) Rust Frontend Distributed" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - rust/ - - vllm/distributed/ - - vllm/engine/ - - vllm/v1/executor/ - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/utils.py - - tests/v1/distributed/test_dense_dp_world_size.py - - tests/v1/distributed/test_external_lb_dp.py - - tests/v1/distributed/test_hybrid_lb_dp.py - - tests/v1/distributed/test_internal_lb_dp.py - - vllm/platforms/rocm.py - commands: - - export VLLM_USE_RUST_FRONTEND=1 - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py - - VLLM_ENGINE_READY_TIMEOUT_S=1800 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" - #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# -- label: ":amd: (MI300) Quantization Shard %N" - timeout_in_minutes: 115 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/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/quantization - - tests/rocm/test_moe_weight_replay.py - - tests/rocm/test_moe_padding.py - optional: true - commands: - # temporary install here since we need nightly, will move to requirements/test.in - # after torchao 0.12 release, and pin a working version of torchao nightly here - # since torchao nightly is only compatible with torch nightly currently - # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now - # we can only upgrade after this is resolved - # TODO(jerryzh168): resolve the above comment - - uv pip install --system torchao==0.18.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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi +# Disabled until torchao transitive dependencies are locked. +# - label: ":amd: (MI300) Quantization Shard %N" +# timeout_in_minutes: 115 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - csrc/ +# - vllm/model_executor/layers/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/quantization +# - tests/rocm/test_moe_weight_replay.py +# - tests/rocm/test_moe_padding.py +# optional: true +# commands: +# # temporary install here since we need nightly, will move to requirements/test.in +# # after torchao 0.12 release, and pin a working version of torchao nightly here +# # since torchao nightly is only compatible with torch nightly currently +# # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now +# # we can only upgrade after this is resolved +# # TODO(jerryzh168): resolve the above comment +# - uv pip install --system torchao==0.18.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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi - label: ":amd: (MI300) Quantized Fusions" # TBD timeout_in_minutes: 45 @@ -2170,30 +1918,31 @@ steps: #------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# -- label: ":amd: (MI300) Regression" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/test_regression - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - optional: true - commands: - - pip install 'modelscope<1.38' - - pytest -v -s test_regression.py +# Disabled until the modelscope installation is locked. +# - label: ":amd: (MI300) Regression" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/ +# - tests/test_regression +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# optional: true +# commands: +# - pip install 'modelscope<1.38' +# - pytest -v -s test_regression.py #--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# @@ -2214,32 +1963,33 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# -- label: ":amd: (MI300) Batch Invariance" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) Batch Invariance" # TBD +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) Batch Invariance (Distributed)" # TBD timeout_in_minutes: 45 @@ -2264,33 +2014,6 @@ steps: - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py -- label: ":amd: (MI300) Spec Decode DeepSeek MTP Parallel Load" - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/llm_base_proposer.py - - vllm/v1/spec_decode/eagle.py - - vllm/v1/worker/gpu/spec_decode/eagle/ - - vllm/model_executor/models/deepseek_mtp.py - - vllm/model_executor/models/deepseek_v2.py - - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py - - vllm/distributed/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/ - - vllm/v1/attention/ - - vllm/v1/engine/ - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/v1/worker/gpu_model_runner.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py - - label: ":amd: (MI300) V1 Spec Decode Shard %N" # TBD timeout_in_minutes: 50 parallelism: 2 @@ -2316,21 +2039,6 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not slow_test' v1/spec_decode --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI300) Acceptance Length (Large Models)" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - tests/v1/spec_decode/test_acceptance_length.py - - vllm/platforms/rocm.py - commands: - - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test - - label: ":amd: (MI300) E2E Scheduling" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2432,38 +2140,39 @@ steps: - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":amd: (MI300) Metrics, Tracing" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - optional: true - num_gpus: 2 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/v1/tracing - - tests/tracing/ - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/tracing/ - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - 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 +# Disabled until the tracing packages installation is locked. +# - label: ":amd: (MI300) Metrics, Tracing" # TBD +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_2 +# optional: true +# num_gpus: 2 +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/ +# - tests/v1/tracing +# - tests/tracing/ +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/tracing/ +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# 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 - label: ":amd: (MI300) V1 E2E" # TBD timeout_in_minutes: 30 @@ -2552,38 +2261,40 @@ steps: - export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json' - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh -- label: ":amd: (MI300) NixlConnector PD edge case" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) NixlConnector PD edge case" # TBD +# timeout_in_minutes: 65 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_2 +# num_gpus: 2 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) Distributed DP Extended" # TBD timeout_in_minutes: 75 @@ -2609,126 +2320,133 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py -- label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD +# timeout_in_minutes: 55 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD - timeout_in_minutes: 80 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 80 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 -- label: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_2 +# num_gpus: 2 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - 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_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_2 +# num_gpus: 2 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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_accuracy_test.sh -- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_2 +# num_gpus: 2 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 ######################################################################################################################################### # # @@ -3142,126 +2860,132 @@ steps: - 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]' -- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_2 +# num_gpus: 2 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 - - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 +# - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py -- label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 - - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 +# - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py -- label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD +# timeout_in_minutes: 40 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" - timeout_in_minutes: 105 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export ENABLE_HMA_FLAG='1' - - export DP_EP='1' - - export GPU_MEMORY_UTILIZATION='0.9' - - export VLLM_SSM_CONV_STATE_LAYOUT='DS' - - export PREFILLER_TP_SIZE='2' - - export DECODER_TP_SIZE='2' - - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' - - export PREFILL_BLOCK_SIZE='2048' - - export DECODE_BLOCK_SIZE='2048' - - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - - 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: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" +# timeout_in_minutes: 105 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 +# commands: +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - export ENABLE_HMA_FLAG='1' +# - export DP_EP='1' +# - export GPU_MEMORY_UTILIZATION='0.9' +# - export VLLM_SSM_CONV_STATE_LAYOUT='DS' +# - export PREFILLER_TP_SIZE='2' +# - export DECODER_TP_SIZE='2' +# - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' +# - export PREFILL_BLOCK_SIZE='2048' +# - export DECODE_BLOCK_SIZE='2048' +# - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - label: ":amd: (MI355) Distributed Features" # TBD timeout_in_minutes: 90 @@ -3375,8 +3099,75 @@ steps: not +quant_fp8 and not +rms_norm and DeepSeek-Coder and ROCM_AITER_MLA' -- label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD - timeout_in_minutes: 85 +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD +# timeout_in_minutes: 85 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_8 +# num_gpus: 8 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355) Distributed Torchrun + Examples" # TBD + timeout_in_minutes: 55 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - tests/distributed/test_torchrun_example.py + - tests/distributed/test_torchrun_example_moe.py + - examples/rl/ + - examples/features/data_parallel/data_parallel_offline.py + - vllm/platforms/rocm.py + commands: + - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py + - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py + - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py + - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py + - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py + - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py + # 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 + # rlhf examples + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_nccl.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_http_ipc.py + +- label: ":amd: (MI355) Distributed DP + EP" # TBD + timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_8 @@ -3384,35 +3175,44 @@ steps: optional: true working_dir: "/vllm-workspace/tests" 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 + - examples/features/torchrun/torchrun_dp_example_offline.py + - vllm/config/parallel.py + - vllm/distributed/ + - vllm/v1/engine/llm_engine.py + - vllm/v1/executor/uniproc_executor.py + - vllm/v1/worker/gpu_worker.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 + - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep + +- label: ":amd: (MI355) Rust Frontend Distributed" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/distributed/ + - vllm/engine/ + - vllm/v1/executor/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/utils.py + - tests/v1/distributed/test_dense_dp_world_size.py + - tests/v1/distributed/test_external_lb_dp.py + - tests/v1/distributed/test_hybrid_lb_dp.py + - tests/v1/distributed/test_internal_lb_dp.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py + - VLLM_ENGINE_READY_TIMEOUT_S=1800 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" #---------------------------------------------------------- mi355 · engine -----------------------------------------------------------# @@ -3783,25 +3583,26 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep-mi355.txt -- label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" - timeout_in_minutes: 105 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - 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 - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" +# timeout_in_minutes: 105 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 +# commands: +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - 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: ":amd: (MI355 DPX) LM Eval TurboQuant k8v4" timeout_in_minutes: 30 @@ -3908,6 +3709,25 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt +- label: ":amd: (MI300) LM Eval Octave KV Cache" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + 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: ":amd: (MI355) LM Eval Spec Decode (Fixed-length FP8)" # TBD timeout_in_minutes: 145 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3986,6 +3806,28 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt +- label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy TP2" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 + - label: ":amd: (MI355) Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy" # TBD timeout_in_minutes: 25 @@ -4009,6 +3851,28 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 +- label: ":amd: (MI355) Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050 + - label: ":amd: (MI355) LM Eval Large Models FP8" # TBD timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4280,6 +4144,91 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400 +- label: ":amd: (MI355) LM Eval Large Models Harness" # TBD + timeout_in_minutes: 85 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 + +- label: ":amd: (MI355) DeepSeek V2-Lite Sync EPLB Accuracy" # TBD + timeout_in_minutes: 40 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/backends/mla/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 + +- label: ":amd: (MI355) Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/spec_decode/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 + +- label: ":amd: (MI355) LM Eval KV-Offload Medium" # TBD + timeout_in_minutes: 70 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" + #--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# - label: ":amd: (MI355 DPX) Examples" # TBD @@ -4570,8 +4519,10 @@ steps: --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py --ignore=kernels/moe/test_deepep_moe.py --ignore=kernels/moe/test_moe_layer.py + --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355) FusedMoE Layer Kernels" # TBD timeout_in_minutes: 35 @@ -5132,28 +5083,29 @@ steps: #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# -- label: ":amd: (MI355 DPX) Quantization Shard %N" - timeout_in_minutes: 115 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - working_dir: "/vllm-workspace/tests" - 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 - optional: true - commands: - - uv pip install --system torchao==0.18.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 --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi +# Disabled until torchao transitive dependencies are locked. +# - label: ":amd: (MI355 DPX) Quantization Shard %N" +# timeout_in_minutes: 115 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# working_dir: "/vllm-workspace/tests" +# 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 +# optional: true +# commands: +# - uv pip install --system torchao==0.18.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 --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi #--------------------------------------------------------- mi355 · samplers ----------------------------------------------------------# @@ -5178,36 +5130,37 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -- label: ":amd: (MI355) Batch Invariance" # TBD - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_1 - optional: true - working_dir: "/vllm-workspace/tests" - 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: ":amd: (MI355) Batch Invariance" # TBD +# timeout_in_minutes: 60 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" timeout_in_minutes: 60 @@ -5458,69 +5411,71 @@ steps: - pytest -v -s v1/simple_kv_offload optional: true -- label: ":amd: (MI355 DPX) V1 KV Connectors Shard %N" - timeout_in_minutes: 60 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - working_dir: /vllm-workspace/tests - source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/simple_kv_offload - - tests/v1/worker - - tests/v1/streaming_input - - tests/v1/kv_connector/unit - - tests/v1/ec_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/entrypoints/pooling/ - - vllm/inputs/ - - vllm/lora/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/outputs.py - - vllm/platforms/ - - vllm/pooling_params.py - - vllm/profiler/ - - vllm/sampling_params.py - - vllm/tokenizers/ - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - 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 --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - optional: true +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355 DPX) V1 KV Connectors Shard %N" +# timeout_in_minutes: 60 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# working_dir: /vllm-workspace/tests +# source_file_dependencies: +# - vllm/ +# - tests/v1/core +# - tests/v1/executor +# - tests/v1/kv_offload +# - tests/v1/simple_kv_offload +# - tests/v1/worker +# - tests/v1/streaming_input +# - tests/v1/kv_connector/unit +# - tests/v1/ec_connector/unit +# - tests/v1/metrics +# - tests/entrypoints/openai/correctness/test_lmeval.py +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/entrypoints/pooling/ +# - vllm/inputs/ +# - vllm/lora/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/outputs.py +# - vllm/platforms/ +# - vllm/pooling_params.py +# - vllm/profiler/ +# - vllm/sampling_params.py +# - vllm/tokenizers/ +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# 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 --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +# - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +# optional: true -- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - num_gpus: 1 - optional: true - working_dir: "/vllm-workspace/tests" - 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" +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# num_gpus: 1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 - label: ":amd: (MI355 DPX) UMBP Connector Unit Tests" timeout_in_minutes: 30 @@ -5542,27 +5497,52 @@ steps: - python3 -c "import mori" - pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py -- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - parallelism: 6 - optional: true - working_dir: "/vllm-workspace/tests" - 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 --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_2 +# num_gpus: 2 +# parallelism: 6 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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 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" +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# 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: ":amd: (MI355 DPX) V1 Metrics + LM Eval" timeout_in_minutes: 65 @@ -5917,32 +5897,75 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ -#----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# +- label: ":amd: (MI355) Spec Decode DeepSeek MTP Parallel Load" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/llm_base_proposer.py + - vllm/v1/spec_decode/eagle.py + - vllm/v1/worker/gpu/spec_decode/eagle/ + - vllm/model_executor/models/deepseek_mtp.py + - vllm/model_executor/models/deepseek_v2.py + - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py + - vllm/distributed/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/v1/attention/ + - vllm/v1/engine/ + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/worker/gpu_model_runner.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py -- label: ":amd: (MI355 DPX) Regression" # TBD - timeout_in_minutes: 25 +- label: ":amd: (MI355 DPX) Acceptance Length (Large Models)" # TBD + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_dpx optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/ - - tests/test_regression - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ + - vllm/v1/spec_decode/ + - tests/v1/spec_decode/test_acceptance_length.py + - vllm/platforms/rocm.py commands: - - pip install 'modelscope<1.38' - - pytest -v -s test_regression.py + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test + +#----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# + +# Disabled until the modelscope installation is locked. +# - label: ":amd: (MI355 DPX) Regression" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/ +# - tests/test_regression +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# commands: +# - pip install 'modelscope<1.38' +# - pytest -v -s test_regression.py - label: ':amd: (MI355) DeepSeek V4 MoE Integration' timeout_in_minutes: 20 diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 7c95577e037..bd381080094 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -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" diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index c833263e704..30560237fed 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -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 diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index 87c004ecfb5..d5ced93f34b 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.buildkite/test_areas/disaggregated_mooncake.yaml @@ -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 diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 303798f25fb..ff2b7d67ecd 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -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). diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 4c80f26ca52..632c3aa1322 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -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 diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 8fc86e65b98..1ca3eacd48c 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -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 diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml index 514ccf90e2b..f9439880d45 100644 --- a/.buildkite/test_areas/fault_tolerance.yaml +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -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 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 44811a4085c..9af2dc4ee2a 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -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" diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 4d610da611a..3855cbcfdc1 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -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 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index f5ff70d9503..dea972df6de 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -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 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index d6c239e4f64..cfa925b2d4c 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -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 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 65df467acf3..b778280bf5e 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -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 diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 22189078fa4..c84a1776d45 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -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)' diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 50cfc2ee86f..1d066c21f0a 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -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 diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 024ee465b0b..dbc6af87f10 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -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). diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 51c89264074..d451ec07ea9 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -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 diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml index 576f9b8c9e1..fb2ff09cdbf 100644 --- a/.buildkite/test_areas/torch_abi.yaml +++ b/.buildkite/test_areas/torch_abi.yaml @@ -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 diff --git a/.github/mergify.yml b/.github/mergify.yml index 39c4731d251..0f545a4a078 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -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: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 1852bd73b68..43a458df767 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -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" }], diff --git a/.github/workflows/recipes-preview.yml b/.github/workflows/recipes-preview.yml new file mode 100644 index 00000000000..42127e99b79 --- /dev/null +++ b/.github/workflows/recipes-preview.yml @@ -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-`, 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}" diff --git a/CMakeLists.txt b/CMakeLists.txt index e1c6169d92b..7694e6f8bc4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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( diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 943133252be..c42ebe895b3 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -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() diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index b6fd26c8034..0af316d59a5 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -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 diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index a1b7d045050..abb8858abe6 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -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 diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index cc96ab291b0..3d6f52d4f18 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -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 ], diff --git a/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py index 115459d4150..2ec5788b881 100644 --- a/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py +++ b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py @@ -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)) diff --git a/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py b/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py new file mode 100644 index 00000000000..12f7571e93c --- /dev/null +++ b/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py @@ -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() diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 7ce49e2e113..580d061a991 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -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, ) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index c4d111e77bf..d92ad58a9cf 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -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" diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index cde756127af..2fbfa157de8 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -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) diff --git a/cmake/hipify.py b/cmake/hipify.py index 2cf58a7b7d3..2135a625dfa 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -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)) diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 83eed340eb7..7d6987122f7 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -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() diff --git a/csrc/attention/octave_arch.cuh b/csrc/attention/octave_arch.cuh new file mode 100644 index 00000000000..536f36b2594 --- /dev/null +++ b/csrc/attention/octave_arch.cuh @@ -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 + +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 diff --git a/csrc/attention/octave_attn.cu b/csrc/attention/octave_attn.cu new file mode 100644 index 00000000000..b4056a1662a --- /dev/null +++ b/csrc/attention/octave_attn.cu @@ -0,0 +1,1180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave KV cache kernels, head_size 256 with 64 RoPE dims. Portable across +// the AMD targets vLLM builds (octave_arch.cuh); the WMMA decode tier for +// gfx11 lives in octave_decode_wmma_rdna3.cu. +// Format: vllm/v1/attention/ops/rocm_octave.py (the PyTorch reference there +// is the source of truth for the byte layout). +// +// octave_cache_store quantize + pack K/V into slots (one wave per slot). +// octave_decode split-KV attention per query group; K dot products +// with int8 dot4 on the codebook values of the packed +// codes, the query quantized to int8 once per block. +// Consecutive query tokens of one request (MTP verify) +// share the KV read. +// octave_rotate the store-time rotations, for the prefill path. +// +// Everything is in the rotated space: K and V are sign-flipped and +// Hadamard-rotated at store time; the query is rotated like K and the +// attention output is rotated back in the reduce kernel. + +#include +#include +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #include "octave_arch.cuh" + #include "octave_format.cuh" + +namespace octave { + +constexpr int DPL = D / 32; // dims per lane when a wave owns one slot + +template +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(half x) { + return __half2float(x); +} +template +__device__ __forceinline__ T from_f(float x) { + return (T)x; +} +template <> +__device__ __forceinline__ half from_f(float x) { + return __float2half(x); +} + +__device__ __forceinline__ float half_bits_to_f(uint32_t bits16) { + return __half2float(__ushort_as_half((unsigned short)bits16)); +} +__device__ __forceinline__ uint16_t f_to_half_bits(float x) { + return __half_as_ushort(__float2half(x)); +} + +template +__device__ __forceinline__ float xor_sum(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v += __shfl_xor(v, m); + return v; +} +template +__device__ __forceinline__ float xor_max(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v = fmaxf(v, __shfl_xor(v, m)); + return v; +} +template +__device__ __forceinline__ uint32_t xor_or(uint32_t v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v |= (uint32_t)__shfl_xor((int)v, m); + return v; +} + +// Walsh-Hadamard butterflies over DPL consecutive dims per lane, then across +// lanes with xor distance 1..LANE_MASK. Unnormalized. +template +__device__ __forceinline__ void fwht_wave(float (&x)[DPL], int lane) { + #pragma unroll + for (int h = 1; h < DPL; h <<= 1) { + #pragma unroll + for (int i = 0; i < DPL; ++i) { + if ((i & h) == 0) { + float a = x[i], b = x[i + h]; + x[i] = a + b; + x[i + h] = a - b; + } + } + } + #pragma unroll + for (int m = 1; m <= LANE_MASK; m <<= 1) { + const bool hi = lane & m; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float o = __shfl_xor(x[i], m); + x[i] = hi ? o - x[i] : x[i] + o; + } + } +} + +__device__ __forceinline__ float sign_of(const int* bits, int i) { + return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f; +} + +// Rotate a lane's 8 dims like K (64-dim blocks) or like V (one 256-dim +// block). +__device__ __forceinline__ void rotate_k(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<4>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.125f; +} + +__device__ __forceinline__ void rotate_full(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<16>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.0625f; +} + +// Midpoints of the int8 codebooks (rocm_octave.LUT), in codebook units. +__device__ constexpr float kThr3[7] = {-103.f, -62.f, -29.5f, 0.f, + 29.5f, 62.f, 103.f}; +__device__ constexpr float kThr4[15] = {-110.5f, -84.5f, -66.f, -50.5f, -36.5f, + -24.f, -12.f, 0.f, 12.f, 24.f, + 36.5f, 50.5f, 66.f, 84.5f, 110.5f}; + +// Nearest codebook entry on the group RMS (4- or 3-bit per lane), codes in +// [0, 2^bits). Returns the stored scale, times LUT_ONE: the one that makes +// x . x_hat = |x|^2, so quantization does not shrink scores or outputs. +template +__device__ __forceinline__ float quantize_group(const float (&x)[DPL], + int (&codes)[DPL], + float group_dims, bool four) { + float ss = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) ss += x[i] * x[i]; + ss = xor_sum(ss); + const float to_lut = + (four ? 46.0f : 59.0f) / fmaxf(sqrtf(ss / group_dims), 1e-12f); + float num = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const float y = x[i] * to_lut; + int c = 0; + #pragma unroll + for (int k = 0; k < 15; ++k) + c += (four ? y > kThr4[k] : (k < 7 && y > kThr3[k])) ? 1 : 0; + codes[i] = c; + num += x[i] * (float)(int8_t)(four ? lut4((uint32_t)c) : lut3((uint32_t)c)); + } + num = xor_sum(num); + return num > 0.f ? ss / num * LUT_ONE : 0.f; +} + +// 4-bit: one word per 8 dims (lane-local). 3-bit: this lane's contribution to +// the 2-bit plane word (16 dims, 2 lanes) and the 1-bit plane word (32 dims, +// 4 lanes); callers OR across lanes. +__device__ __forceinline__ uint32_t pack4(const int (&c)[DPL]) { + uint32_t w = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) + w |= (uint32_t)c[i] << (8 * (i & 3) + 4 * (i >> 2)); + return w; +} +__device__ __forceinline__ void pack3(const int (&c)[DPL], int idx, + uint32_t& lo, uint32_t& hi) { + lo = 0; + hi = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const int e16 = (idx & 1) * DPL + i; // dim within the 16-dim lo word + const int e32 = (idx & 3) * DPL + i; // dim within the 32-dim hi word + lo |= (uint32_t)(c[i] & 3) << (8 * (e16 & 3) + 2 * (e16 >> 2)); + hi |= (uint32_t)(c[i] >> 2) << (8 * (e32 & 3) + (e32 >> 2)); + } +} + +// --------------------------------------------------------------------------- +// Store: grid (num_tokens, num_kv_heads), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + store_kernel(const KT* __restrict__ key, const KT* __restrict__ value, + uint8_t* __restrict__ cache, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ k_signs, + const int* __restrict__ v_signs, int block_size, int64_t sk0, + int64_t sk1, int64_t sv0, int64_t sv1, int64_t scb, + int64_t sch, int64_t sct) { + using F = Format; + const int t = blockIdx.x; + const int h = blockIdx.y; + const int lane = threadIdx.x; + const int64_t slot = slot_mapping[t]; + if (slot < 0) return; + uint8_t* dst = + cache + (slot / block_size) * scb + h * sch + (slot % block_size) * sct; + + float k[DPL], v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + k[i] = to_f(key[t * sk0 + h * sk1 + lane * DPL + i]); + v[i] = to_f(value[t * sv0 + h * sv1 + lane * DPL + i]); + } + + // K, block layout: rotate, then quantize per 64-dim block (8 lanes): the + // RoPE block at 4 bits, the NoPE blocks at 3. Compact layout: rotate over + // all dims like V, 3 bits, one scale. + const bool rope_lane = !KC && lane < R / DPL; + int kc[DPL]; + float k_scale; + if constexpr (KC) { + rotate_full(k, lane, k_signs); + k_scale = quantize_group<16>(k, kc, (float)D, false); + } else { + rotate_k(k, lane, k_signs); + k_scale = quantize_group<4>(k, kc, 64.f, rope_lane); + } + const int nl = lane - F::KR / DPL; // 3-bit lane index (valid when >= 0) + { + uint32_t lo, hi; + pack3(kc, nl & 3, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if (rope_lane) { + *reinterpret_cast(dst + 4 * lane) = pack4(kc); + } else { + if ((nl & 1) == 0) + *reinterpret_cast(dst + F::OFF_KN + 4 * (nl >> 1)) = lo; + if ((nl & 3) == 0) + *reinterpret_cast(dst + F::OFF_KN + F::KN / 4 + + 4 * (nl >> 2)) = hi; + } + } + + // V: rotate all 256 dims, one scale. + rotate_full(v, lane, v_signs); + int vc[DPL]; + const float v_scale = quantize_group<16>(v, vc, (float)D, VB == 4); + if constexpr (VB == 4) { + *reinterpret_cast(dst + F::OFF_V + 4 * lane) = pack4(vc); + } else { + uint32_t lo, hi; + pack3(vc, lane, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if ((lane & 1) == 0) + *reinterpret_cast(dst + F::OFF_V + 4 * (lane >> 1)) = lo; + if ((lane & 3) == 0) + *reinterpret_cast(dst + F::OFF_V + D / 4 + 4 * (lane >> 2)) = + hi; + } + + uint16_t* sc = reinterpret_cast(dst + F::OFF_SC); + if (KC ? lane == 0 : (lane & 7) == 0) sc[lane >> 3] = f_to_half_bits(k_scale); + if (lane == 0) sc[F::NUM_KSC] = f_to_half_bits(v_scale); +} + +// --------------------------------------------------------------------------- +// Decode stage 1: grid (num_q, num_kv_heads, num_splits), block 128. +// A block serves up to QG consecutive query tokens of one request (only the +// first token of each group of a run launches work) and all G query heads of +// its KV head: ROWS = QG * G attention rows share every K/V load. +// --------------------------------------------------------------------------- +constexpr int WAVES = 4; +constexpr int CHUNK = 32 * WAVES; + +template +__global__ void __launch_bounds__(128) + decode_kernel(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_log2, + int num_q, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int num_splits, int64_t sq0, int64_t sq1, + int64_t scb, int64_t sch, int64_t sct, int64_t smo, + int64_t smh, int64_t sms) { + using F = Format; + constexpr int ROWS = QG * G; + constexpr int QW = D / 4; // int8 query words per row + + __shared__ int q8[ROWS][QW]; + __shared__ float q_bs[ROWS][NKB]; // query scale per K block + __shared__ float S[ROWS][CHUNK]; + __shared__ float m_s[ROWS], l_s[ROWS], alpha_s[ROWS]; + __shared__ const uint8_t* tok_ptr[CHUNK]; + __shared__ float tok_vs[CHUNK]; + + const int qi = blockIdx.x; + const int kvh = blockIdx.y; + const int si = blockIdx.z; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int wave = tid >> 5; + + int req = q_to_req[qi]; + const bool req_ok = req >= 0 && req < num_reqs; + int nq = 1; + if constexpr (QG > 1) { + if (req_ok) { + int s = qi; + while (s > 0 && qi - s < 256 && q_to_req[s - 1] == req) --s; + if ((qi - s) % QG != 0) return; // served by the group leader + while (nq < QG && qi + nq < num_q && q_to_req[qi + nq] == req) ++nq; + } + } + const int max_kv = max_blocks * block_size; + int klen[QG]; + int kv_max = 0; + #pragma unroll + for (int j = 0; j < QG; ++j) { + int kl = (req_ok && j < nq) ? q_to_klen[qi + j] : 0; + kl = min(max(kl, 0), max_kv); + klen[j] = kl; + kv_max = max(kv_max, kl); + } + if (!req_ok) req = 0; + + const int tps = (kv_max + num_splits - 1) / num_splits; + const int start = si * tps; + const int end = min(start + tps, kv_max); + + auto write_empty = [&]() { + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + for (int d = lane; d < D; d += 32) op[d] = 0.f; + if (lane == 0) { + op[D] = -INFINITY; + op[D + 1] = 0.f; + } + } + }; + if (start >= end) { + write_empty(); + return; + } + + // Query: rotate like K, quantize to int8 with one scale per 64-dim block. + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + float x[DPL]; + if (j < nq) { + const QT* qp = Q + (qi + j) * sq0 + (kvh * G + g) * sq1 + lane * DPL; + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = to_f(qp[i]); + } else { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = 0.f; + } + if constexpr (KC) + rotate_full(x, lane, k_signs); + else + rotate_k(x, lane, k_signs); + float amax = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) amax = fmaxf(amax, fabsf(x[i])); + const float s = fmaxf(xor_max<4>(amax), 1e-20f) / 127.0f; + uint32_t w[2] = {0, 0}; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + int q = __float2int_rn(x[i] / s); + q = min(max(q, -127), 127); + w[i >> 2] |= (uint32_t)(q & 0xFF) << (8 * (i & 3)); + } + q8[r][2 * lane] = (int)w[0]; + q8[r][2 * lane + 1] = (int)w[1]; + if ((lane & 7) == 0) q_bs[r][lane >> 3] = s; + } + __syncthreads(); + + if (tid < ROWS) { + m_s[tid] = -INFINITY; + l_s[tid] = 0.f; + } + float o0[ROWS], o1[ROWS]; // this lane's two V dims (rotated space) + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + o0[r] = 0.f; + o1[r] = 0.f; + } + const int my_d = wave * 64 + 2 * lane; // first of the two V dims owned + + for (int base = start; base < end; base += CHUNK) { + // ---- scores: lane = token ---- + const int tok = base + wave * 32 + lane; + const bool valid = tok < end; + if (valid) { + const int lb = tok / block_size; + int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) pb = 0; + const uint8_t* p = + cache + pb * scb + kvh * sch + (int64_t)(tok - lb * block_size) * sct; + tok_ptr[wave * 32 + lane] = p; + + constexpr int KW = F::OFF_V / 4; // K code words + uint32_t kw[KW]; + const uint32_t* p4 = reinterpret_cast(p); + #pragma unroll + for (int w = 0; w < KW; ++w) kw[w] = p4[w]; + const uint32_t* kn = kw + F::OFF_KN / 4; + const uint32_t* ps = reinterpret_cast(p + F::OFF_SC); + float ks[NKB]; + if constexpr (KC) { + const uint32_t u = ps[0]; + #pragma unroll + for (int b = 0; b < NKB; ++b) ks[b] = half_bits_to_f(u & 0xFFFF); + tok_vs[wave * 32 + lane] = half_bits_to_f(u >> 16) * (1.0f / LUT_ONE); + } else { + const uint32_t u0 = ps[0], u1 = ps[1]; + ks[0] = half_bits_to_f(u0 & 0xFFFF); + ks[1] = half_bits_to_f(u0 >> 16); + ks[2] = half_bits_to_f(u1 & 0xFFFF); + ks[3] = half_bits_to_f(u1 >> 16); + tok_vs[wave * 32 + lane] = + half_bits_to_f(ps[2] & 0xFFFF) * (1.0f / LUT_ONE); + } + + // Rows two at a time: unrolling all of them kept too much live and + // spilled to scratch on gfx11. + #pragma unroll 2 + for (int r = 0; r < ROWS; ++r) { + const int* qr = q8[r]; + float sk = 0.f; + if constexpr (F::KR > 0) { + int dr = 0; + #pragma unroll + for (int u = 0; u < F::KR / 8; ++u) { + const uint32_t w = kw[u]; + dr = dot4_i8(qr[2 * u], (int)lut4(w & 0x0F0F0F0Fu), dr); + dr = dot4_i8(qr[2 * u + 1], (int)lut4((w >> 4) & 0x0F0F0F0Fu), dr); + } + sk = q_bs[r][0] * ks[0] * (float)dr; + } + #pragma unroll + for (int b = 0; b < F::KN / 64; ++b) { + constexpr int B0 = F::KR / 64; // first 3-bit block + int dn = 0; + #pragma unroll + for (int v = 0; v < 4; ++v) { // 2-bit plane word: 16 dims + const uint32_t lo = kn[4 * b + v]; + const uint32_t hi = kn[F::KN / 16 + 2 * b + (v >> 1)]; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t c = + ((lo >> (2 * i)) & 0x03030303u) | + (((hi >> (4 * (v & 1) + i)) & 0x01010101u) << 2); + dn = + dot4_i8(qr[F::KR / 4 + 16 * b + 4 * v + i], (int)lut3(c), dn); + } + } + sk += q_bs[r][B0 + b] * ks[B0 + b] * (float)dn; + } + const float sc = sk * (sm_scale_log2 / LUT_ONE); + int kl = klen[0]; // klen by constant index only (no scratch) + #pragma unroll + for (int j = 1; j < QG; ++j) kl = r / G == j ? klen[j] : kl; + S[r][wave * 32 + lane] = tok < kl ? sc : -INFINITY; + } + } else { + #pragma unroll + for (int r = 0; r < ROWS; ++r) S[r][wave * 32 + lane] = -INFINITY; + } + __syncthreads(); + + // ---- online softmax: each wave owns rows r = wave, wave + WAVES, ... ---- + for (int r = wave; r < ROWS; r += WAVES) { + float x[CHUNK / 32]; + float mx = -INFINITY; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + x[c] = S[r][32 * c + lane]; + mx = fmaxf(mx, x[c]); + } + mx = xor_max<16>(mx); + const float m_old = m_s[r]; + const float m_new = fmaxf(m_old, mx); + float ps = 0.f; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + const float p = x[c] == -INFINITY ? 0.f : exp2f(x[c] - m_new); + S[r][32 * c + lane] = p; + ps += p; + } + ps = xor_sum<16>(ps); + if (lane == 0) { + const float a = m_new == -INFINITY ? 1.f : exp2f(m_old - m_new); + alpha_s[r] = a; + l_s[r] = l_s[r] * a + ps; + m_s[r] = m_new; + } + } + __syncthreads(); + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float a = alpha_s[r]; + o0[r] *= a; + o1[r] *= a; + } + + // ---- P @ V: wave owns 64 dims, lane owns 2 ---- + const int n_valid = min(CHUNK, end - base); + for (int t = 0; t < n_valid; ++t) { + const uint8_t* p = tok_ptr[t]; + uint32_t c; // codes of dims my_d, my_d + 1 in bytes 0, 1 + if constexpr (VB == 4) { + const uint32_t w = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 3)); + const int e = my_d & 7; // even + c = ((w >> (8 * (e & 3) + 4 * (e >> 2))) & 15) | + (((w >> (8 * ((e + 1) & 3) + 4 * ((e + 1) >> 2))) & 15) << 8); + } else { + const uint32_t lo = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 4)); + const uint32_t hi = *reinterpret_cast( + p + F::OFF_V + D / 4 + 4 * (my_d >> 5)); + const int e16 = my_d & 15, e32 = my_d & 31; + const int f16 = e16 + 1, f32 = e32 + 1; + c = (((lo >> (8 * (e16 & 3) + 2 * (e16 >> 2))) & 3) | + (((hi >> (8 * (e32 & 3) + (e32 >> 2))) & 1) << 2)) | + ((((lo >> (8 * (f16 & 3) + 2 * (f16 >> 2))) & 3) | + (((hi >> (8 * (f32 & 3) + (f32 >> 2))) & 1) << 2)) + << 8); + } + const uint32_t val = lut(c); + const float vs = tok_vs[t]; + const float c0 = (float)(int8_t)(val & 0xFF) * vs; + const float c1 = (float)(int8_t)((val >> 8) & 0xFF) * vs; + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float pr = S[r][t]; + o0[r] += pr * c0; + o1[r] += pr * c1; + } + } + __syncthreads(); + } + + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + op[my_d] = o0[r]; + op[my_d + 1] = o1[r]; + if (tid == 0) { + op[D] = m_s[r]; + op[D + 1] = l_s[r]; + } + } +} + +// --------------------------------------------------------------------------- +// Decode stage 2: combine splits, rotate V back. grid (num_q, num_q_heads), +// block 256. Up to 1024 splits: their weights at once, then 8 waves take the +// splits round-robin (lane = 8 dims) and skip the ones that weigh zero. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(256) + reduce_kernel(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) { + __shared__ float sw[1024], swm[8], swl[8]; + __shared__ float so[8][D]; + const int qi = blockIdx.x, hi = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const float* base = mid_o + qi * smo + hi * smh; + + float ms[4], ls[4]; + float m = -INFINITY; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const int s = tid + 256 * k; + ms[k] = s < num_splits ? base[s * sms + D] : -INFINITY; + ls[k] = s < num_splits ? base[s * sms + D + 1] : 0.0f; + m = fmaxf(m, ms[k]); + } + m = xor_max<16>(m); + if (lane == 0) swm[w] = m; + __syncthreads(); + float mg = swm[0]; + #pragma unroll + for (int i = 1; i < 8; ++i) mg = fmaxf(mg, swm[i]); + float l = 0.0f; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float a = ms[k] == -INFINITY ? 0.0f : exp2f(ms[k] - mg); + if (tid + 256 * k < 1024) sw[tid + 256 * k] = a; + l += ls[k] * a; + } + l = xor_sum<16>(l); + if (lane == 0) swl[w] = l; + __syncthreads(); + + float o[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int s = w; s < num_splits; s += 8) { + const float as = sw[s]; + if (as == 0.0f) continue; + const float4* v = (const float4*)(base + s * sms + lane * 8); + const float4 x = v[0], y = v[1]; + o[0] += x.x * as; + o[1] += x.y * as; + o[2] += x.z * as; + o[3] += x.w * as; + o[4] += y.x * as; + o[5] += y.y * as; + o[6] += y.z * as; + o[7] += y.w * as; + } + #pragma unroll + for (int d = 0; d < 8; ++d) so[w][lane * 8 + d] = o[d]; + __syncthreads(); + float lg = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) lg += swl[i]; + float acc = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) acc += so[i][tid]; + __syncthreads(); + // Inverse rotation of V over the 256 dims (so[0] reused as scratch). + float* x = so[0]; + x[tid] = acc; + __syncthreads(); + for (int hh = 1; hh < D; hh <<= 1) { + float a = 0.f, b = 0.f; + const bool lo = (tid & hh) == 0; + if (lo) { + a = x[tid]; + b = x[tid + hh]; + } + __syncthreads(); + if (lo) { + x[tid] = a + b; + x[tid + hh] = a - b; + } + __syncthreads(); + } + const float inv_l = lg > 0.f ? 1.0f / lg : 0.f; + out[qi * soo + hi * soh + tid] = + from_f(x[tid] * 0.0625f * sign_of(v_signs, tid) * inv_l); +} + +// --------------------------------------------------------------------------- +// In-place rotation of [T, H, 256] vectors for the prefill path, like K +// (signs, then 64-dim Hadamard blocks) or like V (one 256-dim block); +// `inverse` applies the transform then the signs. grid (T, H), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + rotate_kernel(T* __restrict__ x, const int* __restrict__ signs, + bool k_layout, bool inverse, int64_t s0, int64_t s1) { + const int lane = threadIdx.x; + T* p = x + blockIdx.x * s0 + blockIdx.y * s1 + lane * DPL; + float v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] = to_f(p[i]); + if (!inverse) { + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] *= sign_of(signs, lane * DPL + i); + } + float scale; + if (k_layout) { + fwht_wave<4>(v, lane); + scale = 0.125f; + } else { + fwht_wave<16>(v, lane); + scale = 0.0625f; + } + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float y = v[i] * scale; + if (inverse) y *= sign_of(signs, lane * DPL + i); + p[i] = from_f(y); + } +} + +// --------------------------------------------------------------------------- +// Exact window (sparse decode). Direct-mapped fp16 shadow tables keyed by the +// cache slot: entry slot & mask holds that token's exact K and V ([hkv, 2D]) +// and is valid while tags[entry] == slot. The recent table takes every stored +// token, the sink table the first kSinks positions of a sequence. A lost +// entry only means the token is read from the packed cache instead. +// --------------------------------------------------------------------------- +constexpr int kWindow = 64; +constexpr int kSinks = 4; +constexpr int kMaxWindowHits = 128; + +// One wave per token. The per-entry lock keeps two tokens that share an +// entry in the same launch from leaving a mixed row behind. +template +__global__ void __launch_bounds__(32) + window_store_kernel(const T* __restrict__ key, const T* __restrict__ value, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ positions, T* __restrict__ wtab, + int* __restrict__ wtags, int* __restrict__ wlocks, + int wmask, T* __restrict__ stab, + int* __restrict__ stags, int* __restrict__ slocks, + int smask, int hkv, int64_t sk0, int64_t sk1, + int64_t sv0, int64_t sv1) { + const int i = blockIdx.x, lane = threadIdx.x; + const int64_t slot = slot_mapping[i]; + if (slot < 0) return; + const int pos = positions[i]; + auto put = [&](T* tab, int* tags, int* locks, int e) { + if (lane == 0) { + while (atomicCAS(&locks[e], 0, 1) != 0) { + } + atomicExch(&tags[e], -1); + } + __threadfence(); + for (int h = 0; h < hkv; ++h) { + T* row = tab + ((int64_t)e * hkv + h) * (2 * D); + #pragma unroll + for (int d = 0; d < DPL; ++d) { + row[lane * DPL + d] = key[i * sk0 + h * sk1 + lane * DPL + d]; + row[D + lane * DPL + d] = value[i * sv0 + h * sv1 + lane * DPL + d]; + } + } + __threadfence(); + if (lane == 0) { + atomicExch(&tags[e], (int)slot); + __threadfence(); + atomicExch(&locks[e], 0); + } + }; + put(wtab, wtags, wlocks, (int)(slot & wmask)); + if (pos >= 0 && pos < kSinks) put(stab, stags, slocks, (int)(slot & smask)); +} + +// Attention of one query row over its selected positions that sit in the +// exact window, written as one more split of mid_o (in the rotated V space, +// like the packed-cache splits). grid (num_q, hkv), block 256. +template +__global__ void __launch_bounds__(256) window_attn_kernel( + const T* __restrict__ query, const int* __restrict__ positions, + const int* __restrict__ q_to_req, const int* __restrict__ block_table, + const int* __restrict__ indices, const T* __restrict__ wtab, + const int* __restrict__ wtags, int wmask, const T* __restrict__ stab, + const int* __restrict__ stags, int smask, float* __restrict__ mid_o, + const int* __restrict__ v_signs, float sm_scale, int split, int topk, + int64_t sidx, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int hq, int hkv, int64_t sq0, int64_t sq1, int64_t smo, + int64_t smh, int64_t sms) { + __shared__ int s_entry[kMaxWindowHits]; + __shared__ float s_score[8][kMaxWindowHits]; + __shared__ int s_count; + const int qi = blockIdx.x, kvh = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const int hpk = hq / hkv; + if (tid == 0) s_count = 0; + __syncthreads(); + const int req = q_to_req[qi]; + const int p = positions[qi]; + if (req >= 0 && req < num_reqs) { + for (int j = tid; j < topk; j += 256) { + const int t = indices[qi * sidx + j]; + const int d = p - t; + if (t < 0 || d < 0 || (d >= kWindow && t >= kSinks)) continue; + const int lb = t / block_size; + if (lb >= max_blocks) continue; + const int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) continue; + const int slot = pb * block_size + (t - lb * block_size); + int e = -1; + if (d < kWindow && wtags[slot & wmask] == slot) + e = slot & wmask; + else if (t < kSinks && stags[slot & smask] == slot) + e = (slot & smask) | (1 << 30); + if (e < 0) continue; + const int k = atomicAdd(&s_count, 1); + if (k < kMaxWindowHits) s_entry[k] = e; + } + } + __syncthreads(); + const int n = min(s_count, kMaxWindowHits); + const float sl2 = sm_scale * 1.4426950408889634f; + for (int g = w; g < hpk; g += 8) { + const int head = kvh * hpk + g; + float q[DPL]; + #pragma unroll + for (int d = 0; d < DPL; ++d) + q[d] = to_f(query[qi * sq0 + head * sq1 + lane * DPL + d]); + auto row_of = [&](int e) { + const T* tab = (e >> 30) ? stab : wtab; + return tab + ((int64_t)(e & ((1 << 30) - 1)) * hkv + kvh) * (2 * D); + }; + float m = -INFINITY; + for (int i = 0; i < n; ++i) { + const T* r = row_of(s_entry[i]); + float acc = 0.f; + #pragma unroll + for (int d = 0; d < DPL; ++d) acc += q[d] * to_f(r[lane * DPL + d]); + acc = xor_sum<16>(acc) * sl2; + if (lane == 0) s_score[w][i] = acc; + m = fmaxf(m, acc); + } + __builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront"); + __builtin_amdgcn_wave_barrier(); + __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront"); + float l = 0.f, o[DPL] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int i = 0; i < n; ++i) { + const float pr = exp2f(s_score[w][i] - m); + l += pr; + const T* r = row_of(s_entry[i]) + D; + #pragma unroll + for (int d = 0; d < DPL; ++d) o[d] += pr * to_f(r[lane * DPL + d]); + } + rotate_full(o, lane, v_signs); + float* op = mid_o + qi * smo + head * smh + (int64_t)split * sms; + #pragma unroll + for (int d = 0; d < DPL; ++d) op[lane * DPL + d] = o[d]; + if (lane == 0) { + op[D] = n > 0 ? m : -INFINITY; + op[D + 1] = l; + } + } +} + +} // namespace octave + +// --------------------------------------------------------------------------- +// Host launchers +// --------------------------------------------------------------------------- +using namespace octave; + +// fmt: V fmt (3 or 4), plus kCompactFlag for compact K (3-bit V only). +int octave_slot_bytes(int64_t fmt) { + switch (fmt) { + case 3: + return Format<3>::SLOT; + case 4: + return Format<4>::SLOT; + case 3 + kCompactFlag: + return Format<3, true>::SLOT; + } + TORCH_CHECK(false, "octave: unknown format id ", fmt); + return 0; +} + +static void check_cache(const torch::Tensor& cache, int64_t fmt) { + TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4, + "octave: cache must be uint8 [blocks, heads, block_size, slot]"); + const int slot = octave_slot_bytes(fmt); + TORCH_CHECK(cache.size(3) == slot && cache.stride(3) == 1, + "octave: slot size mismatch, expected ", slot, " got ", + cache.size(3)); +} + +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) { + const int n = slot_mapping.size(0); + if (n == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(key.size(-1) == D && key.stride(-1) == 1 && value.stride(-1) == 1, + "octave: head_size must be 256 with contiguous last dim"); + TORCH_CHECK(slot_mapping.dtype() == at::kLong); + const int hkv = cache.size(1); + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(n, hkv); + #define SQ_STORE(B, KC, T) \ + store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + (uint8_t*)cache.data_ptr(), slot_mapping.data_ptr(), \ + k_signs.data_ptr(), v_signs.data_ptr(), cache.size(2), \ + key.stride(0), key.stride(1), value.stride(0), value.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2)) + const bool bf = key.dtype() == at::kBFloat16; + if (fmt == 4) { + if (bf) + SQ_STORE(4, false, __hip_bfloat16); + else + SQ_STORE(4, false, half); + } else if (fmt == 3) { + if (bf) + SQ_STORE(3, false, __hip_bfloat16); + else + SQ_STORE(3, false, half); + } else { + if (bf) + SQ_STORE(3, true, __hip_bfloat16); + else + SQ_STORE(3, true, half); + } + #undef SQ_STORE +} + +template +static void launch_decode(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, float sm_scale_log2, int ns, + hipStream_t stream) { + const int num_q = query.size(0); + dim3 grid(num_q, cache.size(1), ns); + decode_kernel<<>>( + (const QT*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), + block_table.data_ptr(), q_to_req.data_ptr(), + q_to_klen.data_ptr(), mid_o.data_ptr(), + k_signs.data_ptr(), sm_scale_log2, num_q, cache.size(2), + block_table.size(1), block_table.size(0), cache.size(0), ns, + 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)); +} + +template +static void dispatch_decode(int g, int qg, 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, + float sl2, int ns, hipStream_t stream) { + #define SQ_DEC(GG, QQ) \ + if (g == GG && qg == QQ) { \ + launch_decode(query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, \ + stream); \ + return; \ + } + SQ_DEC(1, 1); + SQ_DEC(1, 4); + SQ_DEC(2, 1); + SQ_DEC(2, 4); + SQ_DEC(4, 1); + SQ_DEC(4, 4); + SQ_DEC(6, 1); + SQ_DEC(6, 4); + SQ_DEC(8, 1); + SQ_DEC(8, 4); + #undef SQ_DEC + TORCH_CHECK(false, "octave_decode: unsupported GQA group ", g); +} + +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& indices, + const int* qpos, const int* wtags, int wmask, + const int* stags, int smask, int64_t query_group); + +bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream); + +static void launch_reduce(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream) { + static const bool cdna3 = + std::string(at::cuda::getCurrentDeviceProperties()->gcnArchName) + .rfind("gfx942", 0) == 0; + if (cdna3 && octave_reduce_mfma(out, mid_o, v_signs, num_splits, stream)) + return; + dim3 grid(out.size(0), out.size(1)); + if (out.dtype() == at::kBFloat16) + reduce_kernel<__hip_bfloat16><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_kernel<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), + out.stride(0), out.stride(1)); +} + +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) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq % hkv == 0); + const int g = hq / hkv; + const int qg = query_group > 1 ? 4 : 1; + int ns = (int)num_kv_splits; + const float sl2 = (float)sm_scale * 1.4426950408889634f; + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const bool bf = query.dtype() == at::kBFloat16; + // Fast tier when the device and shapes allow it, else the portable kernel. + const int wmma_splits = + use_wmma ? octave_decode_wmma(query, cache, block_table, q_to_req, + q_to_klen, mid_o, k_signs, sm_scale, + num_kv_splits, fmt, std::nullopt, nullptr, + nullptr, 0, nullptr, 0, qg) + : 0; + if (wmma_splits > 0) { + ns = wmma_splits; + } else { + #define SQ_DISPATCH(B, KC) \ + if (bf) \ + dispatch_decode(g, qg, query, cache, block_table, \ + q_to_req, q_to_klen, mid_o, \ + k_signs, sl2, ns, stream); \ + else \ + dispatch_decode(g, qg, query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, stream) + if (fmt == 4) { + SQ_DISPATCH(4, false); + } else if (fmt == 3) { + SQ_DISPATCH(3, false); + } else { + SQ_DISPATCH(3, true); + } + #undef SQ_DISPATCH + } + launch_reduce(out, mid_o, v_signs, ns, stream); +} + +// Sparse decode: query row i attends to the logical positions indices[i, :] +// of request q_to_req[i] (negative entries are skipped). WMMA tier only. +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& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& stags) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(indices.dim() == 2 && indices.size(0) >= num_q && + indices.stride(1) == 1 && indices.dtype() == at::kInt); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + // With the exact window, the packed-cache splits leave the last split of + // mid_o to octave_window_attn. + const bool win = positions.has_value(); + TORCH_CHECK(!win || (wtab && wtags && stab && stags && num_kv_splits >= 2)); + const int wmask = win ? (int)wtags->size(0) - 1 : 0; + const int smask = win ? (int)stags->size(0) - 1 : 0; + TORCH_CHECK( + !win || (((wmask + 1) & wmask) == 0 && ((smask + 1) & smask) == 0), + "octave_decode_sparse: shadow tables must be powers of two"); + const int ns = octave_decode_wmma( + query, cache, block_table, q_to_req, q_to_req, mid_o, k_signs, sm_scale, + win ? num_kv_splits - 1 : num_kv_splits, fmt, indices, + win ? positions->data_ptr() : nullptr, + win ? wtags->data_ptr() : nullptr, wmask, + win ? stags->data_ptr() : nullptr, smask, 4); + TORCH_CHECK(ns > 0, + "octave_decode_sparse: needs the WMMA tier (gfx1100, head size " + "256, GQA group <= 8)"); + if (!win) { + launch_reduce(out, mid_o, v_signs, ns, stream); + return; + } + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq / hkv <= 8 && wtab->size(1) == hkv && wtab->size(2) == 2 * D); + dim3 grid(num_q, hkv); + #define SQ_WIN(T) \ + window_attn_kernel<<>>( \ + (const T*)query.data_ptr(), positions->data_ptr(), \ + q_to_req.data_ptr(), block_table.data_ptr(), \ + indices.data_ptr(), (const T*)wtab->data_ptr(), \ + wtags->data_ptr(), wmask, (const T*)stab->data_ptr(), \ + stags->data_ptr(), smask, mid_o.data_ptr(), \ + v_signs.data_ptr(), (float)sm_scale, ns, (int)indices.size(1), \ + indices.stride(0), cache.size(2), block_table.size(1), \ + block_table.size(0), cache.size(0), hq, hkv, query.stride(0), \ + query.stride(1), mid_o.stride(0), mid_o.stride(1), mid_o.stride(2)) + if (query.dtype() == at::kBFloat16) + SQ_WIN(__hip_bfloat16); + else + SQ_WIN(half); + #undef SQ_WIN + launch_reduce(out, mid_o, v_signs, ns + 1, stream); +} + +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) { + const int n = slot_mapping.size(0); + if (n == 0) return; + TORCH_CHECK(slot_mapping.dtype() == at::kLong && + positions.dtype() == at::kInt); + TORCH_CHECK(wtab.dtype() == key.dtype() && stab.dtype() == key.dtype()); + const int hkv = wtab.size(1); + const int wmask = (int)wtags.size(0) - 1, smask = (int)stags.size(0) - 1; + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + #define SQ_WSTORE(T) \ + window_store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + slot_mapping.data_ptr(), positions.data_ptr(), \ + (T*)wtab.data_ptr(), wtags.data_ptr(), wlocks.data_ptr(), \ + wmask, (T*)stab.data_ptr(), stags.data_ptr(), \ + slocks.data_ptr(), smask, hkv, key.stride(0), key.stride(1), \ + value.stride(0), value.stride(1)) + if (key.dtype() == at::kBFloat16) + SQ_WSTORE(__hip_bfloat16); + else + SQ_WSTORE(half); + #undef SQ_WSTORE +} + +void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout, + bool inverse) { + if (x.numel() == 0) return; + TORCH_CHECK(x.dim() == 3 && x.size(2) == D && x.stride(2) == 1, + "octave_rotate: x must be [T, H, 256] with contiguous last dim"); + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(x.size(0), x.size(1)); + if (x.dtype() == at::kBFloat16) + rotate_kernel<__hip_bfloat16><<>>( + (__hip_bfloat16*)x.data_ptr(), signs.data_ptr(), k_layout, inverse, + x.stride(0), x.stride(1)); + else + rotate_kernel + <<>>((half*)x.data_ptr(), signs.data_ptr(), + k_layout, inverse, x.stride(0), x.stride(1)); +} + +#else +void octave_cache_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, int64_t) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, double, int64_t, int64_t, int64_t, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode_sparse(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, double, + int64_t, int64_t, const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_window_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_rotate(torch::Tensor, torch::Tensor, bool, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +#endif diff --git a/csrc/attention/octave_decode_mfma_cdna.cu b/csrc/attention/octave_decode_mfma_cdna.cu new file mode 100644 index 00000000000..07f899600c6 --- /dev/null +++ b/csrc/attention/octave_decode_mfma_cdna.cu @@ -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 +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #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 +__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 +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(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 +__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; + 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(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(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(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 +__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><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_mfma<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + 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(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({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<<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (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 diff --git a/csrc/attention/octave_decode_wmma_rdna3.cu b/csrc/attention/octave_decode_wmma_rdna3.cu new file mode 100644 index 00000000000..6f4c6537a08 --- /dev/null +++ b/csrc/attention/octave_decode_wmma_rdna3.cu @@ -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 +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #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 +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(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 +__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; + 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(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& 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(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() : 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 \ + <<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (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 diff --git a/csrc/attention/octave_format.cuh b/csrc/attention/octave_format.cuh new file mode 100644 index 00000000000..faf8a884864 --- /dev/null +++ b/csrc/attention/octave_format.cuh @@ -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 + +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 +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 +__device__ __forceinline__ uint32_t lut(uint32_t codes) { + if constexpr (BITS == 4) + return lut4(codes); + else + return lut3(codes); +} + +} // namespace octave diff --git a/csrc/attention/octave_prefill_cdna.cu b/csrc/attention/octave_prefill_cdna.cu new file mode 100644 index 00000000000..d1638219b79 --- /dev/null +++ b/csrc/attention/octave_prefill_cdna.cu @@ -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 +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #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 +__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 +__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; + 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(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<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), 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<<>>( + (__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(), seq_lens.data_ptr(), 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 diff --git a/csrc/attention/octave_prefill_rdna3.cu b/csrc/attention/octave_prefill_rdna3.cu new file mode 100644 index 00000000000..3566bf62626 --- /dev/null +++ b/csrc/attention/octave_prefill_rdna3.cu @@ -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 +#include +#include +#include +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #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 +__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 +__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; + 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(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> kept; + std::lock_guard lock(mtx); + std::vector& 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<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), 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<<>>( + (__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(), seq_lens.data_ptr(), 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 diff --git a/csrc/attention/octave_torch_bindings.cpp b/csrc/attention/octave_torch_bindings.cpp new file mode 100644 index 00000000000..9446606df85 --- /dev/null +++ b/csrc/attention/octave_torch_bindings.cpp @@ -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 + +#include +#include + +#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& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& 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); +} diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 351bcf24f9f..dbcb6ebaef9 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -420,7 +420,8 @@ __global__ void reshape_and_cache_flash_kernel( } } -template +template __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(src[src_idx], 1.0f); } else { dst[dst_idx] = fp8::scaled_convert(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 \ +#define LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, SKIP_SCALE) \ + vllm::concat_and_cache_mla_kernel \ <<>>( \ reinterpret_cast(kv_c.data_ptr()), \ reinterpret_cast(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(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] diff --git a/csrc/cooperative_topk.cuh b/csrc/cooperative_topk.cuh index ff773904355..dbb87dc4bce 100644 --- a/csrc/cooperative_topk.cuh +++ b/csrc/cooperative_topk.cuh @@ -536,9 +536,13 @@ __device__ void cooperative_topk_body(CooperativeTopKParams 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(sl) + CS - 1) / CS; constexpr uint32_t kFusedMax = ((CS == 16) ? kFusedStagesCS16 : (CS == 8) ? kFusedStagesCS8 : kMaxSinglePassStages) * diff --git a/csrc/cpu/fp8_quant.cpp b/csrc/cpu/fp8_quant.cpp new file mode 100644 index 00000000000..dc04e2b4c2d --- /dev/null +++ b/csrc/cpu/fp8_quant.cpp @@ -0,0 +1,285 @@ +#include + +#include +#include + +#if defined(__AVX512F__) + #include +#endif + +namespace { + +template +constexpr float fp8_max() { + return std::is_same_v ? 57344.0f : 448.0f; +} + +template +inline fp8_t to_fp8(float x) { + constexpr float max = fp8_max(); + return static_cast(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(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(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(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( + (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 +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 && + !std::is_same_v) { + 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()))); + } + } +#endif + for (; i < n; ++i) { + const float v = static_cast(in[i]); + out[i] = to_fp8(kDivide ? v / scale : v * scale); + } +} + +template +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) { + __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(in[i]))); + } + return absmax; +} + +template +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(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(row_out + n, row_in + n, group_n, + 1.0f / row_scale[(n / group_n) * scale_stride_j]); + } + } + } +} + +template +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() * 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(), min_scale); + scale[m] = token_scale; + quant_row(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>& 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( + out.data_ptr(), input.const_data_ptr(), + scale.const_data_ptr(), 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& 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(); + } + + DISPATCH_FP8_QUANT(out.scalar_type(), input.scalar_type(), + "dynamic_per_token_scaled_fp8_quant", [&] { + dynamic_per_token_quant_impl( + out.data_ptr(), scale.data_ptr(), + input.const_data_ptr(), scale_ub_ptr, + num_tokens, hidden_size, input.stride(-2), + out.stride(-2)); + }); +} diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index acf75912819..37c67650e39 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -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 diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h index 13f9760a7ce..b98babea0d5 100644 --- a/csrc/cpu/sgl-kernels/gemm.h +++ b/csrc/cpu/sgl-kernels/gemm.h @@ -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(CPUActMethod::silu_and_mul) && + act_method <= static_cast(CPUActMethod::gelu_tanh_and_mul), + "Unsupported act_method: ", + act_method); + return static_cast(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 @@ -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 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); diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp index 09a170eb6d4..d0c2a46e286 100644 --- a/csrc/cpu/sgl-kernels/moe.cpp +++ b/csrc/cpu/sgl-kernels/moe.cpp @@ -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& w2_bias, const std::optional& alpha, const std::optional& 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( @@ -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( out_hidden_states.data_ptr(), diff --git a/csrc/cpu/sgl-kernels/moe.h b/csrc/cpu/sgl-kernels/moe.h index 876c2d65db6..c0b1d937511 100644 --- a/csrc/cpu/sgl-kernels/moe.h +++ b/csrc/cpu/sgl-kernels/moe.h @@ -164,6 +164,46 @@ inline void silu_and_mul_stub( } } +template +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 || std::is_same_v, + "gelu_and_mul_stub only supports input_t == float or input_t == scalar_t"); + using bVec = at::vec::Vectorized; + + // 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(x0, x1); + out_vec.store(out + d); + } +} + +template +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 || std::is_same_v, + "gelu_tanh_and_mul_stub only supports input_t == float or input_t == scalar_t"); + using bVec = at::vec::Vectorized; + + // 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(x0, x1); + out_vec.store(out + d); + } +} + template 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) { diff --git a/csrc/cpu/sgl-kernels/moe_fp8.cpp b/csrc/cpu/sgl-kernels/moe_fp8.cpp index b8fb54ac772..d9166b1b9e3 100644 --- a/csrc/cpu/sgl-kernels/moe_fp8.cpp +++ b/csrc/cpu/sgl-kernels/moe_fp8.cpp @@ -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) { diff --git a/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp index 0f393a34a85..854b317faba 100644 --- a/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp +++ b/csrc/cpu/sgl-kernels/moe_fp8_w8a8.cpp @@ -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); diff --git a/csrc/cpu/sgl-kernels/moe_int4.cpp b/csrc/cpu/sgl-kernels/moe_int4.cpp index c97784847a9..2940806c992 100644 --- a/csrc/cpu/sgl-kernels/moe_int4.cpp +++ b/csrc/cpu/sgl-kernels/moe_int4.cpp @@ -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); diff --git a/csrc/cpu/sgl-kernels/moe_int8.cpp b/csrc/cpu/sgl-kernels/moe_int8.cpp index e0b81505b72..8f15b9641df 100644 --- a/csrc/cpu/sgl-kernels/moe_int8.cpp +++ b/csrc/cpu/sgl-kernels/moe_int8.cpp @@ -9,7 +9,7 @@ namespace { -template +template 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; @@ -122,7 +126,7 @@ inline void scale_C( } /// gemm for w13 -template +template 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(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::apply( \ +#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_vnni::apply( \ A + mb_start * lda, \ B0 + nb_start * 4, \ B1 + nb_start * 4, \ @@ -248,7 +257,7 @@ struct tinygemm_kernel_vnni { ldb, \ ldc); -template +template 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( - 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( + ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N); + } else { + 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( + /* 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( + /* 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); diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index c32a01ecd50..6f86f6f2a1d 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -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 fast_sigmoid(const at::vec::Vectorized& x) { @@ -631,6 +640,26 @@ inline at::vec::Vectorized fast_silu(const at::vec::Vectorized& x) #endif } +// exact (erf) gelu, matching torch.nn.functional.gelu(approximate="none") +inline at::vec::Vectorized fast_gelu(const at::vec::Vectorized& x) { + const auto half = at::vec::Vectorized(0.5f); + const auto one = at::vec::Vectorized(1.f); + const auto inv_sqrt2 = at::vec::Vectorized(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 fast_gelu_tanh(const at::vec::Vectorized& x) { +#if defined(CPU_CAPABILITY_AVX512) + return at::vec::Vectorized(_mm512_rcp14_gelu_tanh_ps(x)); +#else + const auto c0 = at::vec::Vectorized(1.5957691216057308f); // 2 * sqrt(2/pi) + const auto c1 = at::vec::Vectorized(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 fast_sigmoid_glu(const at::vec::Vectorized& x, const at::vec::Vectorized& alpha) { #if defined(CPU_CAPABILITY_AVX512) diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index f39cf31c1f1..c68d5c6317d 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -78,7 +78,8 @@ void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states, const std::optional& w1_bias, const std::optional& w2_bias, const std::optional& alpha, - const std::optional& limit, bool is_vnni); + const std::optional& 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>& group_shape); +void dynamic_per_token_scaled_fp8_quant( + torch::Tensor& out, const torch::Tensor& input, torch::Tensor& scale, + const std::optional& scale_ub); + void init_cpu_memory_env(std::vector 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__) || \ diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp index 60d3117f24d..0f52c12a6a7 100644 --- a/csrc/fs_io.cpp +++ b/csrc/fs_io.cpp @@ -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; diff --git a/csrc/hisparse_kernels.cu b/csrc/hisparse_kernels.cu index da3857d1e97..74df6553d02 100644 --- a/csrc/hisparse_kernels.cu +++ b/csrc/hisparse_kernels.cu @@ -21,6 +21,7 @@ #include #include #include +#include #include #include @@ -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(state_row) * region_stride; int16_t* row_lru = lru_slots + static_cast(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(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(s_counters + kResidencyCounters); int16_t* s_lru_out = reinterpret_cast(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(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(row) * top_k; int32_t* row_attention = attention_indices != nullptr ? attention_indices + static_cast(row) * attention_row_stride : nullptr; - int32_t* row_miss = miss_mask != nullptr - ? miss_mask + static_cast(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(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(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(row) * top_k + i] = g; } - __syncthreads(); - if (tid == 0) { - if (valid_counts != nullptr) { - valid_counts[static_cast(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(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(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(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(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(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(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(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(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(max_smem_bytes), + STD_TORCH_CHECK(min_smem_bytes <= static_cast(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(max_union_rows) + 1, + static_cast((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 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<<>>( hot_block_table.const_data_ptr(), @@ -1260,10 +1305,10 @@ void hisparse_resolve_residency( device_global_indices.mutable_data_ptr(), lru_slots.mutable_data_ptr(), 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(source_block_size), resident_bt_stride, - resident_num_reqs, resident_num_blocks, + top_k, hot_size, hash_size, static_cast(max_union_rows), + region_stride, attention_block_stride, source_bt_stride, source_num_reqs, + source_num_blocks, static_cast(source_block_size), + resident_bt_stride, resident_num_reqs, resident_num_blocks, static_cast(resident_block_size), static_cast(resident_null_block), global_indices.stride(0), attention_row_stride, valid_count_stride); diff --git a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu index bb9485dd7d4..f6f9845dd18 100644 --- a/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu +++ b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu @@ -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(dst) + offset, &packed, + sizeof(packed)); + } else { + *reinterpret_cast(static_cast(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(h) * kV + v0 + lrow) * kK + 4 * lgrp); + const int64_t base = + h0_row + ((static_cast(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(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(p.cache_in) + offset, + sizeof(packed)); + S[kt][vt] = __builtin_convertvector(packed, f32x4); + } else { + S[kt][vt] = *reinterpret_cast( + static_cast(p.cache_in) + offset); + } + } } else if (PASS2 && (g > 0 || (from_sin && warm))) { const float* s0 = p.sin_ + (static_cast(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(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(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(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(g) * p.NH() + nh) * kK + tid] = aacc; - s1 = p.bg + - (static_cast(g) * p.NH() + nh) * - (static_cast(kV) * kK) + - static_cast(v0 / 16) * kStTile + lane * 4; - st_vt = kStTile; - st_kt = kStKt; - } - if (s1 != nullptr) + float* const s1 = p.bg + + (static_cast(g) * p.NH() + nh) * + (static_cast(kV) * kK) + + static_cast(v0 / 16) * kStTile + lane * 4; for (int vt = 0; vt < kNV; ++vt) for (int kt = 0; kt < kNKTL; ++kt) - *reinterpret_cast(s1 + vt * st_vt + (kt0 + kt) * st_kt) = + *reinterpret_cast(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(checkpoint_state->data_ptr()); + ckpt_ptr = checkpoint_state->data_ptr(); ckoff_ptr = static_cast(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(state_cache->data_ptr()); - cache_out = static_cast(state_cache->data_ptr()); + cache_in = state_cache->data_ptr(); + cache_out = state_cache->data_ptr(); sidx_ptr = static_cast(state_indices->data_ptr()); hh0_ptr = static_cast(has_initial_state->data_ptr()); } @@ -2231,6 +2258,8 @@ void fused_kda_chunk( params.NHv = static_cast(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; diff --git a/csrc/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu index 31c937b2a2b..26a2ff8171c 100644 --- a/csrc/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -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(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(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(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(topk_ids_for_sort), n_token * topk, + expert_map_ptr, n_expert, stream); sortAndScanExpert( get_ptr(topk_ids_for_sort), get_ptr(token_expert_indices), diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index 56c1e2bddf4..2b35f7405f1 100644 --- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -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); diff --git a/csrc/ops.h b/csrc/ops.h index 7c4463ac3f1..1c5ddc0beb4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -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& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& 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, diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh index 1309968d064..bd18483b48d 100644 --- a/csrc/persistent_topk.cuh +++ b/csrc/persistent_topk.cuh @@ -12,10 +12,13 @@ #include #include "topk_histogram_4096.cuh" +#include "persistent_topk_overflow.cuh" namespace vllm { namespace persistent { +using topk_histogram_4096::convert_to_uint32_v2; + // ============================================================================ // Constants // ============================================================================ @@ -25,7 +28,7 @@ constexpr int RADIX = 256; // Medium path: all shared state in dynamic smem (no static __shared__, // which would inflate the kernel's smem footprint and kill occupancy -// for the decode/trivial paths). +// for the short and trivial paths). constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 constexpr size_t kMediumHeaderSize = @@ -33,84 +36,19 @@ constexpr size_t kMediumHeaderSize = constexpr int MAX_BUFFERED_ITEMS = 4096; constexpr size_t kSmemMedium = kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968 +static_assert(topk_histogram_4096::kExactCandidateOffset + + MAX_BUFFERED_ITEMS * sizeof(int) <= + kSmemMedium); constexpr uint32_t RADIX_THRESHOLD = 32768; -// Decode path constants -constexpr int kDecodeBins = 2048; -constexpr uint32_t HIST2048_THRESHOLD = 8192; +// Short and medium path constants +constexpr int kCoarseBins = 2048; +constexpr uint32_t SHORT_EXACT_THRESHOLD = 8192; // Large path: fixed shared memory for histograms + scalars constexpr size_t kFixedSmemLarge = ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15); -// ============================================================================ -// Common helpers -// ============================================================================ - -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); -} - -// ============================================================================ -// Vectorized load helpers -// ============================================================================ - -// Unconditional float4 load with cache hint (.cg = cache at global level only). -__device__ __forceinline__ void load_float4(const float* ptr, float& v0, - float& v1, float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// Per-element predicated scalar loads with -inf default. -__device__ __forceinline__ void load_float4_predicated(const float* ptr, - int base, int seq_len, - float& v0, float& v1, - float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - int p0 = (base < seq_len); - int p1 = (base + 1 < seq_len); - int p2 = (base + 2 < seq_len); - int p3 = (base + 3 < seq_len); - asm volatile( - "{\n" - " .reg .pred pr0, pr1, pr2, pr3;\n" - " setp.ne.u32 pr0, %4, 0;\n" - " setp.ne.u32 pr1, %5, 0;\n" - " setp.ne.u32 pr2, %6, 0;\n" - " setp.ne.u32 pr3, %7, 0;\n" - " mov.u32 %0, 0xFF800000;\n" - " mov.u32 %1, 0xFF800000;\n" - " mov.u32 %2, 0xFF800000;\n" - " mov.u32 %3, 0xFF800000;\n" - " @pr0 ld.global.cg.u32 %0, [%8];\n" - " @pr1 ld.global.cg.u32 %1, [%8+4];\n" - " @pr2 ld.global.cg.u32 %2, [%8+8];\n" - " @pr3 ld.global.cg.u32 %3, [%8+12];\n" - "}\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - // ============================================================================ // Large path: inter-CTA coordination state (one per group) // ============================================================================ @@ -141,281 +79,68 @@ struct PersistentTopKParams { }; // ============================================================================ -// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) -// Uses 11-bit half-precision bins for fine granularity. -// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg. +// Coarse selection shared by the medium and filtered paths // ============================================================================ -// 11-bit bin from half-precision representation (ascending: high values -> high -// bins) -__device__ __forceinline__ uint32_t decode_bin(float x) { - __half hx = __float2half(x); - uint16_t bits = __half_as_ushort(hx); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> 5; -} +template +struct CoarseSuffixSmem { + int histogram[Bins + 1]; + uint32_t warp_sum[BlockSize / 32]; +}; -template -__device__ __noinline__ void histogram_2048_topk( - const float* __restrict__ logits, int32_t* __restrict__ output_indices, - int32_t seq_len) { - extern __shared__ int decode_smem[]; +// Rewrite the coarse counts as suffix counts (histogram[b] = keys in bins >= b) +// and call on_select(bin) from the one thread whose bin holds rank k. Uses the +// shuffle scan of histogram_4096_topk: a CUB block scan spills under the +// persistent kernel's 32-register budget. +template +__device__ __forceinline__ void select_coarse_bin( + CoarseSuffixSmem* coarse, int k, OnSelect on_select) { + namespace hist4096 = topk_histogram_4096; + constexpr int kBinsPerThread = Bins / BlockSize; + static_assert(Bins % BlockSize == 0); + static_assert(BlockSize == 32 * 32, "one warp sum per lane"); const int tx = threadIdx.x; - const int lane = tx & 31; - - // ---- Layout constants ---- - constexpr int SBASE = 8192 - 8; // 8184 - constexpr int RHIST = RADIX + 128; // 384 - constexpr int BOFF = 2 * RHIST; // 768 - constexpr int DBUF = (SBASE - BOFF) / 2; // 3708 - constexpr int MAX_ITEMS_PER_THREAD = - (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; - - enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; - - // ---- Initialize scalars (prevents stale data from prior rows) ---- - if (tx < 8) { - decode_smem[SBASE + tx] = 0; + const uint32_t lane = tx % 32; + const uint32_t warp = tx / 32; + const int first_bin = kBinsPerThread * tx; + int counts[kBinsPerThread]; + uint32_t local_sum = 0; +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + counts[i] = coarse->histogram[first_bin + i]; + local_sum += counts[i]; } - - // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- - int* histo = decode_smem; - uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; - int nitems = 0; - - for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) { - histo[i] = 0; + const uint32_t warp_inclusive = hist4096::warp_inclusive_sum(lane, local_sum); + if (lane == 31) coarse->warp_sum[warp] = warp_inclusive; + __syncthreads(); + const uint32_t warp_total = coarse->warp_sum[lane]; + const uint32_t total = hist4096::warp_reduce_sum_full(warp_total); + const uint32_t warps_before = + hist4096::warp_reduce_sum_full(lane < warp ? warp_total : 0); + // Each thread rewrites only its own bins, so no barrier is needed here. + int suffix = + static_cast(total - warps_before - warp_inclusive + local_sum); +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + coarse->histogram[first_bin + i] = suffix; + suffix -= counts[i]; } + if (tx == 0) coarse->histogram[Bins] = 0; __syncthreads(); - const int n_vec = (seq_len + 3) >> 2; - const bool row_aligned = ((reinterpret_cast(logits) & 15) == 0); - - for (int i = tx; i < n_vec; i += kThreadsPerBlock) { - const int base = i << 2; - float v0, v1, v2, v3; - - if (row_aligned && base + 3 < seq_len) { - load_float4(logits + base, v0, v1, v2, v3); - } else { - load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3); - } - - const uint16_t b0 = static_cast(decode_bin(v0)); - const uint16_t b1 = static_cast(decode_bin(v1)); - const uint16_t b2 = static_cast(decode_bin(v2)); - const uint16_t b3 = static_cast(decode_bin(v3)); - reg_bins[nitems++] = b0; - reg_bins[nitems++] = b1; - reg_bins[nitems++] = b2; - reg_bins[nitems++] = b3; - atomicAdd(&histo[b0], 1); - atomicAdd(&histo[b1], 1); - atomicAdd(&histo[b2], 1); - atomicAdd(&histo[b3], 1); - } - __syncthreads(); - - // ---- CUB suffix sum ---- - using BlockScanT = cub::BlockScan; - const int h0 = histo[2 * tx]; - const int pair_sum = h0 + histo[2 * tx + 1]; - - auto& scan_storage = *reinterpret_cast( - decode_smem + kDecodeBins); - - int pair_prefix, total; - BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total); - - // Find threshold bin purely from registers - const int pair_suffix = total - pair_prefix; - - if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx; - } - { - const int right_suf = pair_suffix - h0; - const int next_suf = pair_suffix - pair_sum; - if (right_suf >= TopK && next_suf < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx + 1; +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + const int bin = first_bin + i; + if (coarse->histogram[bin] > k && coarse->histogram[bin + 1] <= k) { + on_select(bin); } } __syncthreads(); - - const int threshold = decode_smem[SBASE + sTHR]; - - // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- - int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; - const int sOUT_abs = SBASE + sOUT; - const int sBUF0_abs = SBASE + sBUF0; - - { - const uint32_t uthr = static_cast(threshold); - int item = 0; - const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; - - for (int iter = 0; iter < n_vec_iters; iter++) { - const int i = tx + iter * kThreadsPerBlock; - const bool vec_valid = (i < n_vec); - const int base_idx = i << 2; - -#pragma unroll 4 - for (int sub = 0; sub < 4; sub++) { - const int elem_idx = base_idx + sub; - uint32_t bin = 0; - if (vec_valid) bin = reg_bins[item++]; - // Lanes past seq_len are -inf padding from the row's last float4. - const bool in_range = vec_valid && (elem_idx < seq_len); - const bool is_above = in_range && (bin > uthr); - const bool is_equal = in_range && (bin == uthr); - - const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); - if (above_mask) { - const int above_count = __popc(above_mask); - const int above_rank = __popc(above_mask & ((1u << lane) - 1)); - int above_base; - if (lane == 0) { - above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); - } - above_base = __shfl_sync(0xffffffff, above_base, 0); - if (is_above) { - output_indices[above_base + above_rank] = elem_idx; - } - } - - const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); - if (equal_mask) { - const int equal_count = __popc(equal_mask); - const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); - int equal_base; - if (lane == 0) { - equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); - } - equal_base = __shfl_sync(0xffffffff, equal_base, 0); - if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { - bufs[0][equal_base + equal_rank] = elem_idx; - } - } - } - } - } - __syncthreads(); - - int remaining_k = TopK - decode_smem[SBASE + sOUT]; - if (remaining_k <= 0) return; - - // If all buffered elements fit, output them all (common for short seqs) - const int raw_buf0 = decode_smem[SBASE + sBUF0]; - if (raw_buf0 <= remaining_k) { - const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - const int base = decode_smem[SBASE + sOUT]; - for (int i = tx; i < nb; i += kThreadsPerBlock) { - output_indices[base + i] = bufs[0][i]; - } - __syncthreads(); - return; - } - - // ---- Phase 3: Deferred refinement (rare path) ---- - int* refine[2] = {decode_smem, decode_smem + RHIST}; - const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - - for (int i = tx; i < RHIST; i += kThreadsPerBlock) { - refine[0][i] = 0; - } - __syncthreads(); - - for (int i = tx; i < num_buf0; i += kThreadsPerBlock) { - const uint32_t fp32 = convert_to_uint32_v2(logits[bufs[0][i]]); - atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1); - } - __syncthreads(); - - auto compute_suffix_sum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (tx < RADIX) { - const int stride = 1 << i; - const int s = i & 1; - const int d = s ^ 1; - int value = refine[s][tx]; - if (tx < RADIX - stride) value += refine[s][tx + stride]; - refine[d][tx] = value; - } - __syncthreads(); - } - }; - -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { - const int src = pass & 1; - const int dst = src ^ 1; - - const int raw_buf = decode_smem[SBASE + sBUF0 + src]; - const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; - - compute_suffix_sum(); - - if (tx < RADIX && refine[0][tx] > remaining_k && - refine[0][tx + 1] <= remaining_k) { - decode_smem[SBASE + sREF] = tx; - decode_smem[SBASE + sBUF0 + dst] = 0; - decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; - } - __syncthreads(); - - const int ref_thr = decode_smem[SBASE + sREF]; - remaining_k -= refine[0][ref_thr + 1]; - const int bit_offset = 24 - pass * 8; - - if (remaining_k == 0) { - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const uint32_t fp32 = convert_to_uint32_v2(logits[idx]); - if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } - } - __syncthreads(); - break; - } - - __syncthreads(); - if (tx < RADIX + 1) refine[0][tx] = 0; - __syncthreads(); - - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const float logit_val = logits[idx]; - const uint32_t fp32 = convert_to_uint32_v2(logit_val); - const int bin = (fp32 >> bit_offset) & 0xFF; - - if (bin > ref_thr) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } else if (bin == ref_thr) { - if (pass == 3) { - const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); - if (slot > 0) output_indices[TopK - slot] = idx; - } else { - const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); - if (__builtin_expect(bp < DBUF, 1)) { - bufs[dst][bp] = idx; - const int nbo = bit_offset - 8; - atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1); - } - } - } - } - __syncthreads(); - } } // ============================================================================ -// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement -// For sequences 8K < seq_len <= 64K. +// Medium path: exact 11-bit FP32 histogram + 3-pass FP32 refinement +// For sequences 8K < seq_len <= 32K. // ============================================================================ // Adapted from: @@ -423,10 +148,12 @@ __device__ __noinline__ void histogram_2048_topk( // by: DarkSharpness // which at the same time is an optimized topk kernel copied from tilelang // kernel -template -__device__ __noinline__ void histogram_256_topk( +// Returns where exact recovery resumes when the threshold bin overflows the +// candidate buffer; .needed is false otherwise. +template +__device__ __noinline__ overflow::MediumResume histogram_medium_topk( const float* __restrict__ logits, int* __restrict__ output_indices, - int logits_offset, int seq_len) { + int seq_len) { // All shared state lives in dynamic shared memory to avoid static extern __shared__ char medium_smem[]; @@ -444,16 +171,41 @@ __device__ __noinline__ void histogram_256_topk( const int thread_id = threadIdx.x; int remaining_k = TopK; - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; + // The coarse histogram overlays the refinement state until it is consumed. + struct CoarseSmem { + CoarseSuffixSmem suffix; + int threshold; + }; + static_assert(sizeof(CoarseSmem) <= kSmemMedium); + auto* coarse = reinterpret_cast(medium_smem); + int* coarse_histogram = coarse->suffix.histogram; + + for (int bin = thread_id; bin < kCoarseBins + 1; bin += kThreadsPerBlock) { + coarse_histogram[bin] = 0; } __syncthreads(); - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const auto bin = convert_to_uint8(logits[idx + logits_offset]); - atomicAdd(&shared_histogram[0][bin], 1); + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t, float score) { + atomicAdd(&coarse_histogram[convert_to_uint32_v2(score) >> 21], 1); + }); + __syncthreads(); + + select_coarse_bin(&coarse->suffix, TopK, + [&](int bin) { coarse->threshold = bin; }); + + const int threshold_bin = coarse->threshold; + const int count_above = coarse_histogram[threshold_bin + 1]; + const int threshold_bin_count = coarse_histogram[threshold_bin] - count_above; + __syncthreads(); + + if (thread_id == 0) { + shared_threshold_bin = threshold_bin; + shared_buffered_count[0] = 0; + shared_output_count = 0; } __syncthreads(); + remaining_k -= count_above; auto compute_cumulative_sum = [&]() { #pragma unroll 8 @@ -472,57 +224,50 @@ __device__ __noinline__ void histogram_256_topk( } }; - compute_cumulative_sum(); - - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[0] = 0; - shared_output_count = 0; - } - __syncthreads(); - - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - if (remaining_k == 0) { - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const int bin = convert_to_uint8(logits[idx + logits_offset]); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } - } + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t idx, float score) { + const int bin = convert_to_uint32_v2(score) >> 21; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + }); __syncthreads(); - return; + return {}; + } + + if (__builtin_expect(threshold_bin_count > MAX_BUFFERED_ITEMS, 0)) { + return {.needed = true, + .prefix = static_cast(threshold_bin) << 21, + .remaining = static_cast(remaining_k)}; } - __syncthreads(); if (thread_id < RADIX + 1) { shared_histogram[0][thread_id] = 0; } __syncthreads(); - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const float logit_value = logits[idx + logits_offset]; - const int bin = convert_to_uint8(logit_value); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); - if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { - buffered_indices[0][buffer_pos] = idx; - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int next_bin = (fp32_bits >> 24) & 0xFF; - atomicAdd(&shared_histogram[0][next_bin], 1); - } - } - } + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t idx, float score) { + const uint32_t ordered = convert_to_uint32_v2(score); + const int bin = ordered >> 21; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + const int next_bin = (ordered >> 16) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + }); __syncthreads(); -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { +#pragma unroll 3 + for (int pass = 0; pass < 3; ++pass) { const int src_buffer = pass % 2; const int dst_buffer = src_buffer ^ 1; const int raw_buffered = shared_buffered_count[src_buffer]; @@ -541,13 +286,12 @@ __device__ __noinline__ void histogram_256_topk( const int threshold_bin = shared_threshold_bin; remaining_k -= shared_histogram[0][threshold_bin + 1]; - const int bit_offset = 24 - pass * 8; + const int bit_offset = 16 - pass * 8; if (remaining_k == 0) { for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { const int idx = buffered_indices[src_buffer][i]; - const uint32_t fp32_bits = - convert_to_uint32_v2(logits[idx + logits_offset]); + const uint32_t fp32_bits = convert_to_uint32_v2(logits[idx]); const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); @@ -566,14 +310,14 @@ __device__ __noinline__ void histogram_256_topk( for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { const int idx = buffered_indices[src_buffer][i]; - const float logit_value = logits[idx + logits_offset]; + const float logit_value = logits[idx]; const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); output_indices[output_pos] = idx; } else if (bin == threshold_bin) { - if (pass == 3) { + if (pass == 2) { const int slot = atomicAdd(&shared_final_k, -1); if (slot > 0) { output_indices[TopK - slot] = idx; @@ -592,6 +336,7 @@ __device__ __noinline__ void histogram_256_topk( } __syncthreads(); } + return {}; } // ============================================================================ @@ -643,7 +388,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, } // ============================================================================ -// Large path: multi-CTA radix select for sequences > 64K +// Large path: multi-CTA radix select for sequences > 32K // // Each row is processed by a group of CTAs. Each CTA loads its chunk into // shared memory as ordered uint32, then participates in 4 rounds of @@ -754,6 +499,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, next_hist[i] = 0; } } + __syncthreads(); if (tx == 0) { red_release(&state->arrival_counter, 1); @@ -855,12 +601,12 @@ __device__ void radix_topk(const float* __restrict__ row_input, } } } + // The CTA's next row reuses this smem; finish the tie loop before it does. + __syncthreads(); } // ============================================================================ -// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK -// BS>32 uses standalone histogram_256_buffered_topk (separate kernel, -// see filtered_topk.cuh) +// Persistent kernel — short/medium/large paths with RadixTopK. // ============================================================================ template @@ -871,7 +617,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) // ======================================================================== // Group mode: multi-CTA groups with static round-robin row assignment. - // Non-large rows: CTA-0 handles trivial/decode/medium. + // Non-large rows: CTA-0 handles the trivial, short, and medium paths. // Large rows: all CTAs in the group cooperate via RadixTopK. // ======================================================================== const uint32_t ctas_per_group = params.ctas_per_group; @@ -939,10 +685,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) i += kThreadsPerBlock) { row_output[i] = (i < seq_len) ? static_cast(i) : -1; } - } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { - histogram_2048_topk(row_input, row_output, seq_len); - } else { - histogram_256_topk(row_input, row_output, 0, seq_len); + } else if (seq_len <= static_cast(SHORT_EXACT_THRESHOLD)) { + topk_histogram_4096::exact_topk_rescan_cold( + row_input, row_output, seq_len, smem_raw); + } else if (const auto resume = histogram_medium_topk( + row_input, row_output, seq_len); + __builtin_expect(resume.needed, 0)) { + overflow::recover_medium( + row_input, row_output, seq_len, smem_raw, resume.prefix, + resume.remaining); } } continue; @@ -960,19 +712,13 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) } // namespace persistent // ============================================================================ -// ============================================================================ -// Optimized FilteredTopK — single CTA per row for bs > 32. -// Kept with persistent_topk so the portable fallback owns the non-cluster path. +// FilteredTopK — one CTA per row, for the non-cluster path. +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ namespace filtered_topk { namespace hist4096 = topk_histogram_4096; -// ============================================================================ -// FilteredTopK — single CTA per row for bs > 32 -// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 -// ============================================================================ - #define FLASHINFER_CUDA_CALL(func, ...) \ { \ cudaError_t e = (func); \ @@ -1000,29 +746,26 @@ struct vec_t { #undef FLASHINFER_INLINE // FilteredTopK traits for different data types -template +template struct FilteredTopKTraits; -// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4 -// refinement rounds -template <> -struct FilteredTopKTraits { +// Both float paths use the exact ordered-FP32 key; only the coarse-bin width +// and refinement digit layout differ. +template +struct FilteredTopKTraits { using OrderedType = uint32_t; - static constexpr int NUM_REFINE_ROUNDS = 4; - static constexpr int FIRST_REFINE_SHIFT = 24; + static constexpr int COARSE_BITS = UseWideCoarse ? 12 : 11; + static constexpr int COARSE_RADIX = 1 << COARSE_BITS; + static constexpr int NUM_REFINE_ROUNDS = UseWideCoarse ? 2 : 3; + static constexpr int FIRST_REFINE_SHIFT = UseWideCoarse ? 10 : 16; + static constexpr int REFINE_BITS = UseWideCoarse ? 10 : 8; - __device__ __forceinline__ static uint8_t ToCoarseKey(float x) { - // Convert to FP16 representation and extract high 8 bits - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); + __device__ __forceinline__ static uint16_t ToCoarseKey(float x) { + return static_cast(ToOrdered(x) >> (32 - COARSE_BITS)); } __device__ __forceinline__ static OrderedType ToOrdered(float x) { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + return topk_histogram_4096::convert_to_uint32_v2(x); } }; @@ -1031,10 +774,16 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = 16 * 1024; // 16K indices per buffer constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB +constexpr uint32_t FILTERED_TOPK_WIDE_COARSE_MIN_LENGTH = 64 * 1024; +constexpr uint32_t FILTERED_TOPK_WIDE_COARSE_MAX_LENGTH = 192 * 1024; +static_assert(hist4096::kExactCandidateOffset + + FILTERED_TOPK_SMEM_INPUT_SIZE * sizeof(int) <= + FILTERED_TOPK_SMEM_DYNAMIC); -template +template struct FilteredTopKStorage { - alignas(128) int histogram[2][256 + 128]; + static constexpr int RADIX = UseWideCoarse ? 1024 : 256; + alignas(128) int histogram[2][RADIX + 128]; alignas(128) int counter; alignas(128) int threshold_bin; alignas(128) int num_input[2]; @@ -1042,16 +791,24 @@ struct FilteredTopKStorage { int last_remain; }; -// With CheckOverflow, return false before using a truncated stash so the caller -// can retry with exact full-row selection. +// With CheckOverflow, return false instead of recovering when the candidate +// stash overflows, so the caller can finish with its own exact selection. template + bool UsePredicatedShortLoads, bool CheckOverflow = false, + bool UseWideCoarse = false> __device__ __forceinline__ bool filtered_topk_row( const DType* score, IdType* dst, int length, uint32_t top_k, - FilteredTopKStorage& storage) { + FilteredTopKStorage& storage, uint32_t num_rows = 0) { + extern __shared__ __align__(16) uint8_t dynamic_smem[]; constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; - constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; + using Traits = FilteredTopKTraits; + constexpr int RADIX = 1 << Traits::REFINE_BITS; + static_assert(RADIX == FilteredTopKStorage::RADIX); + constexpr int COARSE_RADIX = Traits::COARSE_RADIX; + using CoarseSmem = persistent::CoarseSuffixSmem; + static_assert(sizeof(CoarseSmem) <= FILTERED_TOPK_SMEM_DYNAMIC); + const int tx = threadIdx.x; // Trivial case: length <= top_k @@ -1064,13 +821,14 @@ __device__ __forceinline__ bool filtered_topk_row( // Short path if (length <= 32768) { - extern __shared__ uint8_t _smem_reg[]; if constexpr (UsePredicatedShortLoads) { - hist4096::histogram_4096_topk_predicated(score, dst, length, - _smem_reg); + hist4096::histogram_4096_topk_predicated< + MAX_K, 12, 8, hist4096::OverflowRecovery::kRegisterValues>( + score, dst, length, dynamic_smem); } else { - hist4096::histogram_4096_topk(score, dst, length, - _smem_reg); + hist4096::histogram_4096_topk< + MAX_K, 12, 8, false, hist4096::OverflowRecovery::kRegisterValues>( + score, dst, length, dynamic_smem); } return true; } @@ -1080,16 +838,19 @@ __device__ __forceinline__ bool filtered_topk_row( auto& s_threshold_bin_id = storage.threshold_bin; auto& s_num_input = storage.num_input; auto& s_indices = storage.indices; + auto& s_histogram = s_histogram_buf[0]; - // Dynamic shared memory for input double buffer - extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE]; + auto* coarse_smem = reinterpret_cast(dynamic_smem); + auto* s_coarse_histogram = coarse_smem->histogram; + auto* s_input_idx = reinterpret_cast(dynamic_smem); - using Traits = FilteredTopKTraits; int topk = top_k; - // Stage 1: 8-bit coarse histogram with vectorized loads - if (tx < RADIX + 1) s_histogram[tx] = 0; + // Stage 1: high bits of the exact ordered-FP32 key. + for (int bin = tx; bin < COARSE_RADIX + 1; bin += BLOCK_SIZE) { + s_coarse_histogram[bin] = 0; + } __syncthreads(); vec_t score_vec; @@ -1102,20 +863,20 @@ __device__ __forceinline__ bool filtered_topk_row( #pragma unroll for (int j = 0; j < VEC_SIZE; ++j) { const auto bin = Traits::ToCoarseKey(score_vec[j]); - atomicAdd(&s_histogram[bin], 1); + atomicAdd(&s_coarse_histogram[bin], 1); } } // Handle tail for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { const auto bin = Traits::ToCoarseKey(score[i]); - atomicAdd(&s_histogram[bin], 1); + atomicAdd(&s_coarse_histogram[bin], 1); } __syncthreads(); - // Suffix sum - const auto run_cumsum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { + // Suffix sum for the exact-refinement histograms. + const auto run_refine_cumsum = [&]() { +#pragma unroll + for (int i = 0; i < Traits::REFINE_BITS; ++i) { if (tx < RADIX) { const auto j = 1 << i; const auto k = i & 1; @@ -1129,16 +890,15 @@ __device__ __forceinline__ bool filtered_topk_row( } }; - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { - s_threshold_bin_id = tx; + // The selecting thread also resets the refinement counters, saving a barrier. + persistent::select_coarse_bin(coarse_smem, topk, [&](int bin) { + s_threshold_bin_id = bin; s_num_input[0] = 0; s_counter = 0; - } - __syncthreads(); + }); const auto threshold_bin = s_threshold_bin_id; - topk -= s_histogram[threshold_bin + 1]; + topk -= s_coarse_histogram[threshold_bin + 1]; constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; @@ -1168,8 +928,21 @@ __device__ __forceinline__ bool filtered_topk_row( } __syncthreads(); } else { - __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; + // The threshold bin does not fit the candidate stash: finish exactly. + if (__builtin_expect(s_coarse_histogram[threshold_bin] - + s_coarse_histogram[threshold_bin + 1] > + SMEM_INPUT_SIZE, + 0)) { + if constexpr (CheckOverflow) return false; + overflow::recover_coarse( + score, dst, length, top_k, num_rows, s_coarse_histogram, + threshold_bin, dynamic_smem, s_histogram); + return true; + } + for (int bin = tx; bin < RADIX + 1; bin += BLOCK_SIZE) { + s_histogram[bin] = 0; + } __syncthreads(); // Filter + histogram for refinement @@ -1183,7 +956,7 @@ __device__ __forceinline__ bool filtered_topk_row( if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { s_input_idx[0][pos] = index; const auto ordered = Traits::ToOrdered(raw_input); - const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; + const auto sub_bin = (ordered >> FIRST_SHIFT) & (RADIX - 1); atomicAdd(&s_histogram[sub_bin], 1); } } @@ -1203,7 +976,7 @@ __device__ __forceinline__ bool filtered_topk_row( } __syncthreads(); - // Stage 2: refine with 8bit radix passes + // Stage 2: refine the selected coarse bin. #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { auto& s_last_remain = storage.last_remain; @@ -1216,7 +989,7 @@ __device__ __forceinline__ bool filtered_topk_row( const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; - run_cumsum(); + run_refine_cumsum(); if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { s_threshold_bin_id = tx; s_num_input[r_idx ^ 1] = 0; @@ -1227,13 +1000,14 @@ __device__ __forceinline__ bool filtered_topk_row( const auto threshold = s_threshold_bin_id; topk -= s_histogram[threshold + 1]; - const int offset = FIRST_SHIFT - round * 8; + const int offset = FIRST_SHIFT - round * Traits::REFINE_BITS; const bool is_last_round = (round == NUM_ROUNDS - 1); if (topk == 0) { for (int i = tx; i < num_input; i += BLOCK_SIZE) { const auto idx = s_input_idx[r_idx][i]; - const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; + const auto bin = + (Traits::ToOrdered(score[idx]) >> offset) & (RADIX - 1); if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); s_indices[pos] = idx; @@ -1243,12 +1017,15 @@ __device__ __forceinline__ bool filtered_topk_row( break; } else { __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; + for (int bin = tx; bin < RADIX + 1; bin += BLOCK_SIZE) { + s_histogram[bin] = 0; + } __syncthreads(); for (int i = tx; i < num_input; i += BLOCK_SIZE) { const auto idx = s_input_idx[r_idx][i]; const auto raw_input = score[idx]; - const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; + const auto bin = + (Traits::ToOrdered(raw_input) >> offset) & (RADIX - 1); if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); s_indices[pos] = idx; @@ -1263,7 +1040,8 @@ __device__ __forceinline__ bool filtered_topk_row( if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { s_input_idx[r_idx ^ 1][pos] = idx; const auto bin32 = Traits::ToOrdered(raw_input); - const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; + const auto sub_bin = + (bin32 >> (offset - Traits::REFINE_BITS)) & (RADIX - 1); atomicAdd(&s_histogram[sub_bin], 1); } } @@ -1291,7 +1069,7 @@ __device__ __forceinline__ bool filtered_topk_row( * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) */ template + bool UsePredicatedShortLoads = false, bool UseWideCoarse = false> __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) FilteredTopKUnifiedKernel(const DType* __restrict__ input, IdType* __restrict__ output, @@ -1301,9 +1079,11 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const uint32_t bid = blockIdx.x; if (bid >= num_rows) return; const int length = lengths ? lengths[bid] : static_cast(max_len); - __shared__ FilteredTopKStorage storage; - filtered_topk_row( - input + bid * max_len, output + bid * top_k, length, top_k, storage); + __shared__ FilteredTopKStorage storage; + filtered_topk_row(input + bid * max_len, + output + bid * top_k, length, top_k, + storage, num_rows); } // Helper to compute GCD for VEC_SIZE selection @@ -1341,17 +1121,31 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, void* args[] = {&input, &output_indices, &lengths, &num_rows, &top_k_val, &max_len}; - const int vec_size = ComputeFilteredTopKVecSize(max_len); + // Views like x[:, 1:] keep the stride but shift the base. + int vec_size = ComputeFilteredTopKVecSize(max_len); + while (vec_size > 1 && + reinterpret_cast(input) % (vec_size * sizeof(DType)) != 0) { + vec_size /= 2; + } + const bool use_wide_coarse = + MAX_K == 1024 && max_len >= FILTERED_TOPK_WIDE_COARSE_MIN_LENGTH && + max_len < FILTERED_TOPK_WIDE_COARSE_MAX_LENGTH; -#define DISPATCH_VEC_SIZE(VS) \ - if (vec_size == VS) { \ - auto kernel = \ - FilteredTopKUnifiedKernel; \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ - smem_size, stream)); \ - return cudaSuccess; \ +#define DISPATCH_VEC_SIZE(VS) \ + if (vec_size == VS) { \ + auto kernel = FilteredTopKUnifiedKernel; \ + if constexpr (MAX_K == 1024) { \ + if (use_wide_coarse) { \ + kernel = FilteredTopKUnifiedKernel; \ + } \ + } \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ + smem_size, stream)); \ + return cudaSuccess; \ } DISPATCH_VEC_SIZE(1) diff --git a/csrc/persistent_topk_overflow.cuh b/csrc/persistent_topk_overflow.cuh new file mode 100644 index 00000000000..9ae116998e8 --- /dev/null +++ b/csrc/persistent_topk_overflow.cuh @@ -0,0 +1,262 @@ +#ifndef PERSISTENT_TOPK_OVERFLOW_CUH_ +#define PERSISTENT_TOPK_OVERFLOW_CUH_ + +#include +#include +#include +#include + +#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 +__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( + scores, output, length, smem, prefix, remaining); +} + +} // namespace vllm::persistent::overflow + +namespace vllm::filtered_topk { + +template +struct vec_t; + +template +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 +__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; + using Radix = SampledRadixDigits; + constexpr uint32_t kItemsPerThread = 2; + constexpr uint32_t kSampleSize = BlockSize * kItemsPerThread; + using SampleSort = cub::BlockRadixSort; + 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( + (static_cast(2 * sample_idx + 1) * length) / + (2 * kSampleSize)); + sample_keys[i] = Traits::ToOrdered(scores[score_idx]); + } + + auto* sort_smem = + static_cast(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( + (static_cast(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(sample_keys[rank_index % kItemsPerThread]); + } + if (tx == 0) *candidate_count = 0; + __syncthreads(); + + auto* candidates = static_cast(dynamic_smem); + const uint32_t proposed_pivot = static_cast(*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(samples_at_or_above) * length > + static_cast(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(CandidateCapacity)) candidates[pos] = idx; + } + }; + + vec_t 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(top_k) || + num_candidates > static_cast(CandidateCapacity)) { + return false; + } + + auto* candidate_keys = candidates + CandidateCapacity; + for (int i = tx; i < num_candidates; i += BlockSize) { + candidate_keys[i] = + static_cast(Traits::ToOrdered(scores[candidates[i]])); + } + + if (tx == 0) { + *radix_prefix = 0; + *radix_remaining = static_cast(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(*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(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(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(prefix | (static_cast(bin) << shift)); + break; + } + count_above += count; + } + } + __syncthreads(); + } + + if (tx == 0) { + *output_counter = 0; + histogram[0] = 0; + } + __syncthreads(); + + const uint32_t pivot = static_cast(*radix_prefix); + const int equal_count = *radix_remaining; + const int equal_base = static_cast(top_k) - equal_count; + for (int i = tx; i < num_candidates; i += BlockSize) { + const int idx = candidates[i]; + const uint32_t ordered = static_cast(candidate_keys[i]); + if (ordered > pivot) { + output[atomicAdd(output_counter, 1)] = static_cast(idx); + } else if (ordered == pivot) { + const int pos = atomicAdd(&histogram[0], 1); + if (pos < equal_count) + output[equal_base + pos] = static_cast(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 +__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(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(scores, output, length, top_k, + dynamic_smem, scratch)) { + return; + } + topk_histogram_4096::exact_topk_rescan( + scores, output, length, dynamic_smem, prefix_bin << 21, remaining); +} + +} // namespace overflow +} // namespace vllm::filtered_topk + +#endif // PERSISTENT_TOPK_OVERFLOW_CUH_ diff --git a/csrc/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu index 4c6875f0623..1f9fd7c8a16 100644 --- a/csrc/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -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; 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; 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 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(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts); + n_experts, sorted_expert_ids); } else { mxfp4_cvt_fp16_to_fp4 <<>>( @@ -297,7 +308,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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<<>>( m_topk, k, reinterpret_cast(input), @@ -316,7 +328,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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 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 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( 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(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 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( 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(sorted_expert_ids->data_ptr()) + : nullptr, + m_topk, k, n_experts, stream); }); #else STD_TORCH_CHECK_NOT_IMPLEMENTED( diff --git a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu index b6043e80330..268650ae181 100644 --- a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -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& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& 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(topk_ids.data_ptr()), static_cast(input_permutation.data_ptr()), static_cast(output_permutation.data_ptr()), + sorted_expert_ids.has_value() + ? static_cast(sorted_expert_ids->data_ptr()) + : nullptr, static_cast(atomic_buffer.data_ptr()), topk_ids.numel(), topk_ids.size(1)); } diff --git a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 7e4f942f544..0fd72ebd082 100644 --- a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -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& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& 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& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& 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( diff --git a/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu new file mode 100644 index 00000000000..8cc03ea5397 --- /dev/null +++ b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu @@ -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 +#include +#include + +#include +#include +#include + +#include + +#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 +__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 +__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 +__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 +__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(v); + } +} + +#else // non-gfx11 device pass: empty kernels for symbol parity. + +template +__global__ void mxfp4_w4a8_quant_act_kernel(const T*, int8_t*, float*, int, + int) {} + +template +__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 +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<<>>( + reinterpret_cast(a.data_ptr()), a_q.data_ptr(), + a_s.data_ptr(), M, K); + + const dim3 block(32, kWarpsPerBlock); + const dim3 grid((N + kWarpsPerBlock - 1) / kWarpsPerBlock); + const int* a_q32 = reinterpret_cast(a_q.data_ptr()); + const float* a_s_ptr = a_s.data_ptr(); + const uint8_t* b_q_ptr = b_q.data_ptr(); + const uint8_t* b_s_ptr = b_scale.data_ptr(); + T* out_ptr = reinterpret_cast(out.data_ptr()); + +#define MXFP4_W4A8_LAUNCH(MM) \ + case MM: \ + mxfp4_w4a8_gemv_rdna3_kernel<<>>( \ + 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(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; +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index de963913ea5..f36cd3d9018 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -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); diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 077c3f35cef..5c34d6676b9 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -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, " diff --git a/csrc/sampled_topk.cuh b/csrc/sampled_topk.cuh index 867abbd6354..c25f8ed8e73 100644 --- a/csrc/sampled_topk.cuh +++ b/csrc/sampled_topk.cuh @@ -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 @@ -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(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 diff --git a/csrc/topk.cu b/csrc/topk.cu index fb3d946b96b..d5ff2a8386e 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -82,10 +82,13 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, uint32_t max_chunk_elements = static_cast(available_for_ordered / sizeof(uint32_t)); + // Views like x[:, 1:] keep the stride but shift the base. + const auto base = + reinterpret_cast(logits.const_data_ptr()); 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; diff --git a/csrc/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh index cf575708e8d..b5aeb2d78cd 100644 --- a/csrc/topk_histogram_4096.cuh +++ b/csrc/topk_histogram_4096.cuh @@ -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 +__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 +__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 + 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( + vecs, output, length, thr_bin, num_above, _smem)) { + exact_topk_rescan( + 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 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(scores, output, - length, _smem); + histogram_4096_topk( + scores, output, length, _smem); } } // namespace topk_histogram_4096 diff --git a/csrc/topk_histogram_4096_overflow.cuh b/csrc/topk_histogram_4096_overflow.cuh index 617d312c79d..7433a0362ac 100644 --- a/csrc/topk_histogram_4096_overflow.cuh +++ b/csrc/topk_histogram_4096_overflow.cuh @@ -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 +// 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 __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; + 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(_smem); + auto* final_candidates = reinterpret_cast( + static_cast(_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(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( + 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(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(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(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( - 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(idx); - } else if (ordered == pivot) { - const uint32_t pos = atomicAdd(&smem->histogram[0], 1); - if (pos < equal_count) { - output[equal_base + pos] = static_cast(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(idx); + } else if (ordered == pivot) { + const uint32_t pos = atomicAdd(&smem->histogram[0], 1); + if (pos < equal_count) { + output[equal_base + pos] = static_cast(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( + scores, length, [&](uint32_t idx, float score) { + if ((convert_to_uint32_v2(score) & kFinalPrefixMask) == + (pivot & kFinalPrefixMask)) { + collect(idx, score); + } + }); + } + } else { + scan_scores(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; + +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 +__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(_smem); + const uint32_t tx = threadIdx.x; + const auto range = coarse_refine_range(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(&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(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(&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(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(idx); + } + } + } + } + } + __syncthreads(); + return true; +} + } // namespace topk_histogram_4096 } // namespace vllm diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a64d844e4d7..1c84b53ece2 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -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. diff --git a/docker/Dockerfile b/docker/Dockerfile index c359a93c378..f3cac42652e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -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 diff --git a/docker/Dockerfile.efa b/docker/Dockerfile.efa new file mode 100644 index 00000000000..00c22248c13 --- /dev/null +++ b/docker/Dockerfile.efa @@ -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}" diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 029d003b554..d1a9c99803a 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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 diff --git a/docker/Dockerfile.rocm_72 b/docker/Dockerfile.rocm_72 index a24b65458ab..eeaa9d6c1b1 100644 --- a/docker/Dockerfile.rocm_72 +++ b/docker/Dockerfile.rocm_72 @@ -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 diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index c251641452f..31c8ad6405c 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -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. diff --git a/docker/versions.json b/docker/versions.json index 75c21afafc6..03a9babdc0d 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -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" diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index a38ef0ab66a..1ea961fefa4 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -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