From 3b6c0bdae9ef0a923dfbd545e3633b95db36a13d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 1 Sep 2026 15:58:27 -0700 Subject: [PATCH] [CI] Shard long kernel test groups (#54754) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 92 +++++++++++++++++------------- 1 file changed, 53 insertions(+), 39 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 37fe5489e6..35cd72bff3 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -106,7 +106,7 @@ steps: key: kernels-attention-test device: l4 num_devices: 1 - timeout_in_minutes: 65 + timeout_in_minutes: 30 source_file_dependencies: - csrc/attention/ - vllm/v1/attention @@ -116,7 +116,7 @@ steps: - tests/kernels/attention commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + parallelism: 7 mirror: amd: label: ":amd: (MI300) Attention Kernels Shard %N" @@ -169,19 +169,22 @@ steps: key: kernels-quantization-test device: l4 num_devices: 1 - timeout_in_minutes: 60 + timeout_in_minutes: 30 source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization - tests/kernels/quantization commands: - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + parallelism: 6 mirror: amd: label: ":amd: (MI300) Quantization Kernels Shard %N" dind: false device: mi300_1 + # Pin the AMD mirror at its current 2 shards so raising NVIDIA + # parallelism does not multiply AMD copies (needs ci-infra #473). + parallelism: 2 timeout_in_minutes: 120 source_file_dependencies: - csrc/quantization/ @@ -277,9 +280,10 @@ steps: - pytest -v -s kernels/attention/test_deepgemm_attention.py - pytest -v -s quantization/test_cutlass_w4a16.py -- label: ":nvidia: (B200) Kernels" +- label: ":nvidia: (B200) Kernels Shard %N" key: kernels-b200 - timeout_in_minutes: 80 + timeout_in_minutes: 30 + parallelism: 3 working_dir: "/vllm-workspace/" device: b200-k8s # optional: true @@ -323,39 +327,49 @@ steps: - tests/kernels/test_top_k_per_row.py commands: - nvidia-smi - - python3 examples/basic/offline_inference/chat.py - # Attention - # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353 - - pytest -v -s tests/kernels/attention/test_attention_selector.py - - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2' - - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py - - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py - - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py - - pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py - - pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - - pytest -v -s tests/kernels/test_top_k_per_row.py - # Quantization - - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8' - - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py - - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py - - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py - - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py - - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py - - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py - - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py - - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py - - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py - - pytest -v -s tests/kernels/moe/test_flashinfer.py - - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py - - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py - - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py - - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py - - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py - - pytest -v -s tests/kernels/attention/test_minimax_m3.py - - pytest -v -s tests/kernels/test_ll_bf16_gemm.py - # e2e - - pytest -v -s tests/models/quantization/test_nvfp4.py + - | + set -eu + case "$$BUILDKITE_PARALLEL_JOB" in + 0) + pytest -v -s tests/kernels/attention/test_attention_selector.py + pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py + pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py + pytest -v -s tests/kernels/test_top_k_per_row.py + pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py + pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py + pytest -v -s tests/kernels/moe/test_nvfp4_moe.py + pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py + pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py + pytest -v -s tests/kernels/test_ll_bf16_gemm.py + ;; + 1) + python3 examples/basic/offline_inference/chat.py + pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py + pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py + pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py + pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py + pytest -v -s tests/kernels/moe/test_mxfp4_moe.py + pytest -v -s tests/kernels/moe/test_flashinfer_moe.py + pytest -v -s tests/kernels/moe/test_cutedsl_moe.py + pytest -v -s tests/models/quantization/test_nvfp4.py + ;; + 2) + pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2' + pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py + pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8' + pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py + pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py + pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py + pytest -v -s tests/kernels/moe/test_flashinfer.py + pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py + pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py + pytest -v -s tests/kernels/attention/test_minimax_m3.py + ;; + *) + echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" >&2 + exit 1 + ;; + esac - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly