diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 9e634ea02b4d..4ae92315b4c9 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -693,28 +693,8 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition" -# - label: ":amd: (MI300) Sequence Parallel Correctness" # TBD -# timeout_in_minutes: 180 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/" -# source_file_dependencies: -# - vllm/model_executor/layers/ -# - vllm/compilation/ -# - vllm/v1/worker/ -# - vllm/v1/cudagraph_dispatcher.py -# - tests/compile/correctness_e2e/test_sequence_parallel.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_TEST_CLEAN_GPU_MEMORY=1 -# - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - - label: ":amd: (MI300) Distributed Compile + RPC" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3180,7 +3160,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ":amd: (MI355) Fusion and Compile FP8" timeout_in_minutes: 45 @@ -3205,64 +3184,6 @@ steps: - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py -- label: ":amd: (MI355) Sequence Parallel Correctness" - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - -- label: ":amd: (MI355) AsyncTP Correctness" - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":amd: (MI355 DPX) Fusion E2E FP8 Config Sweep" timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -5916,35 +5837,6 @@ steps: - amdmi355 optional: true -- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - timeout_in_minutes: 30 - working_dir: /vllm-workspace/ - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - dind: false - agent_pool: mi355_2 - num_gpus: 2 - mirror_hardwares: - - amdexperimental - - amdproduction - - amdgfx950nightly - - amdmi355 - optional: true - label: ":amd: (MI355 DPX) MRV2 Sampler JIT Warmup" timeout_in_minutes: 10 diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index acc1bafa7f2a..4a5a4f1a3cea 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,84 +2,6 @@ group: Compile depends_on: - image-build steps: -- label: ":nvidia: (B200) Sequence Parallel Correctness" - key: sequence-parallel-correctness-tests-2xb200 - timeout_in_minutes: 45 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - mirror: - amd: - label: ":amd: (MI355) Sequence Parallel Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 180 - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - -- label: ":nvidia: (B200) AsyncTP Correctness" - key: asynctp-correctness-tests-b200 - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py - mirror: - amd: - label: ":amd: (MI355) AsyncTP Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":nvidia: (H100) Distributed Compile" key: distributed-compile-unit-tests-2xh100 timeout_in_minutes: 45 @@ -300,7 +222,6 @@ steps: - nvidia-smi # Run all models and attn backends but only Inductor partition and native custom ops - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 Quick" @@ -388,53 +309,6 @@ steps: - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and Llama-3 and ROCM_AITER_UNIFIED_ATTN and (dynamo_partition or +quant_fp8 or +rms_norm)' -- label: ":nvidia: (H100) Fusion E2E TP2 AsyncTP Config Sweep" - key: fusion-e2e-tp2-asynctp-config-sweep-h100 - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/" - device: h100 - num_devices: 2 - source_file_dependencies: - - csrc/quantization/ - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/attention/attention.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/fusions_e2e/ - commands: - - nvidia-smi - # Run just llama3 (fp8 & bf16) for all config combinations - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3" - mirror: - amd: - label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - dind: false - device: mi355_2 - num_devices: 2 - working_dir: /vllm-workspace/ - timeout_in_minutes: 30 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 timeout_in_minutes: 45 @@ -453,7 +327,6 @@ steps: # include qwen/deepseek with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported # for ar-rms-quant-fp4, also sweep llama3 - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites" @@ -480,7 +353,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)' key: fusion-e2e-tp2-ar-rms-inductor-partition-amd diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index abebdd54ae46..9d53a98f6585 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -225,9 +225,9 @@ outputs = model.generate( ) ``` -### Piecewise compilation and full graph custom passes (attention fusion, sequence parallelism) +### Piecewise compilation and full graph custom passes (attention fusion) -Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. +Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. Long term, we've added the ability to partition the graph in Inductor instead of right after Dynamo. It can be enabled with `CompilationConfig.use_inductor_graph_partition=True` but is currently experimental and only available with `torch>=2.9`. This also increases compilation time as it has to compile the whole graph and cannot reuse piecewise compilation artifacts. Once vLLM supports 2.9, we plan to make this the default approach as it will also speed up piecewise cudagraph capture. diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 002984c113a1..4095c41a7e63 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -25,8 +25,6 @@ or just on the low or high end. | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low | -| [Sequence Parallelism](#sequence-parallelism-enable_sp) | `enable_sp` | AllReduce → ReduceScatter + AllGather | Off by default | Prereq for AsyncTP | Yes | High | -| [AsyncTP GEMM + collective](#asynctp-gemm--collective-overlap-fuse_gemm_comms) | `fuse_gemm_comms` | GEMM → reduce-scatter / all-gather → GEMM | Off by default | 7-10% | Yes | High | | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always | @@ -44,9 +42,7 @@ The table below lists the quantization schemes supported by each fusion on each | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | | `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | -| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — | -| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | -| `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | +| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — | | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | @@ -56,10 +52,6 @@ The table below lists the quantization schemes supported by each fusion on each fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) for per-backend details. -† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today; -other architectures support requires setting `PassConfig.sp_min_token_num` explicitly. -SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`. - ## Enabling / Disabling Fusions Fusions are exposed through `PassConfig`, which is nested inside `CompilationConfig`: @@ -186,72 +178,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py) -### Sequence Parallelism (`enable_sp`) - -**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather, -splitting the sequence dimension across TP ranks. This restructures the graph so the subsequent AsyncTP -pass can fuse the reduce-scatter / all-gather with the surrounding GEMMs. - -Sequence Parallelism itself does not directly improve performance; it is a prerequisite for the -AsyncTP pass (`fuse_gemm_comms`). SP is only applied above a minimum token threshold that is -autoconfigured based on device capability and model `hidden_size`. Currently only active on -H100/SM90 for models with `hidden_size >= 8192`. The threshold is configurable via -`PassConfig.sp_min_token_num`. - -The general transformation: - -```text -Input → AllReduce → RMSNorm → Output -becomes: -Input → ReduceScatter → local RMSNorm → AllGather → Output -``` - -Patterns covered: - -- First block: `AllReduce → RMSNorm` → `ReduceScatter → RMSNorm → AllGather` -- Middle blocks: `AllReduce → fused_add_RMSNorm` → `ReduceScatter → fused_add_RMSNorm → AllGather` -- Both with optional `→ FP8 static quant` suffix - -Requires: `use_inductor_graph_partition=True` **or** piecewise compilation with static sizes -divisible by `tensor_parallel_size`. - -Supported hardware: Only tested on NVIDIA CUDA, possibly works on ROCm. FP8 all-gather requires sm90+. - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - -### AsyncTP GEMM + Collective Overlap (`fuse_gemm_comms`) - -!!! info - Requires `enable_sp=True` (enabled automatically). This pass is a no-op if Sequence Parallelism has not been applied. - -**What it fuses.** After Sequence Parallelism transforms the graph, fuses GEMM kernels with the -surrounding reduce-scatter (output projection) and all-gather (input projection) using -`torch.ops.symm_mem` symmetric-memory primitives, overlapping communication and computation. -This overlap is only profitable for large `num_tokens`, so the fusion (and preceding SP) -is only performed in the higher compiled range above `PassConfig.sp_min_token_num`. - -Patterns covered: - -- `GEMM → reduce-scatter` → `fused_matmul_reduce_scatter` -- `all-gather → GEMM` → `all_gather_matmul` -- FP8 scaled variants of both patterns - -Supported hardware: NVIDIA CUDA with symmetric-memory (`torch.distributed._symmetric_memory`) support. - -On B200, pattern-matching fp8 FlashInfer scaled MM is not supported, so it must be disabled -([#27893](https://github.com/vllm-project/vllm/issues/27893)) - -```shell -VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ... -``` - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/collective_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/collective_fusion.py) -- Sequence parallelism pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - ### QK Norm + RoPE (`enable_qk_norm_rope_fusion`) !!! info diff --git a/docs/design/vllm_ir.md b/docs/design/vllm_ir.md index 82628f3762fe..e324ad25d03a 100644 --- a/docs/design/vllm_ir.md +++ b/docs/design/vllm_ir.md @@ -233,19 +233,17 @@ storing this information in vLLM's `PassContext` for later use in clone eliminat ### 3. IR Fusion and Transformation Passes After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations. -These passes can perform fusion, distribute operations for sequence parallelism, and other transformations: +These passes can perform fusion and other transformations: ```python -# Example: Sequence Parallelism (see SequenceParallelismPass) -# Before SP pass +# Example: AllReduce + RMSNorm fusion (see AllReduceFusionPass) +# Before all_reduce = torch.ops.vllm.all_reduce(x, "tp:0") rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -# after SP pass -reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0") -rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -all_gather = torch.ops.vllm.all_gather(x, "tp:0") +# After +rms_norm = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm(x, weight, ...) ``` Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations, diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index e82336d24117..b40ff4da8cc6 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -148,11 +148,10 @@ When batch invariance is enabled, vLLM: 1. Uses deterministic kernel implementations for attention and other operations 2. Ensures consistent numerical behavior across different batch sizes -3. Disables certain optimizations that may introduce non-determinism (such as sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant) +3. Disables certain optimizations that may introduce non-determinism 4. Under tensor parallelism, keeps custom all-reduce on with a fixed reduction order (the 1-stage kernel is pinned, and large inputs are reduced in fixed-size chunks), and disables FlashInfer, AITER and QuickReduce all-reduce 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell), - runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out - (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes). + runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out. 6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics !!! warning diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py deleted file mode 100644 index 307f39039b39..000000000000 --- a/tests/compile/correctness_e2e/test_async_tp.py +++ /dev/null @@ -1,270 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -import logging -from contextlib import nullcontext - -import pytest - -from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS -from tests.models.registry import HF_EXAMPLE_MODELS -from tests.utils import ( - compare_two_settings, - create_new_process_for_each_test, - multi_gpu_test, -) -from vllm.config import ( - CompilationMode, -) -from vllm.platforms import current_platform -from vllm.transformers_utils.repo_utils import hf_api -from vllm.utils.flashinfer import has_flashinfer - -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_HF_OVERRIDES = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, -} - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn): - """Trained-model outputs must agree with unfused TP, including padded decode.""" - # A snapshot path pins both servers and the comparison helper's tokenizer. - model = hf_api().snapshot_download( - "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca" - ) - common_args = [ - "--dtype", - "bfloat16", - "--load-format", - "auto", - "--tensor-parallel-size", - "2", - "--distributed-executor-backend", - "mp", - "--max-model-len", - "512", - "--max-num-batched-tokens", - "512", - "--max-num-seqs", - "8", - "--kv-cache-memory-bytes", - str(64 * 1024**2), - "--gpu-memory-utilization", - "0.01", - "--attention-backend", - "TRITON_ATTN", - "--disable-custom-all-reduce", - ] - - def settings(enabled): - return [ - *common_args, - "--compilation-config", - json.dumps( - { - "mode": CompilationMode.VLLM_COMPILE, - "cudagraph_mode": "NONE", - "splitting_ops": [], - "inductor_compile_config": {"force_disable_caches": True}, - "pass_config": { - "enable_sp": enabled, - "fuse_gemm_comms": enabled, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - ), - ] - - env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"} - with caplog_mp_spawn(logging.DEBUG) as logs: - compare_two_settings( - model, - settings(False), - settings(True), - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - matches = [ - int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text) - ] - print(f"AsyncTP rewrite counts: {matches}") - assert len(matches) >= 2 and all(count > 0 for count in matches), ( - f"Expected positive AsyncTP rewrites on both TP workers, found {matches}" - ) - - -@create_new_process_for_each_test() -@pytest.mark.parametrize( - "model_id", - ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"], -) -@pytest.mark.parametrize("tp_size", [2]) -@pytest.mark.parametrize("async_tp_enabled", [True]) -@pytest.mark.parametrize("distributed_backend", ["mp"]) -def test_async_tp_pass_correctness( - model_id: str, - tp_size: int, - async_tp_enabled: bool, - distributed_backend: str, - num_gpus_available: int, - monkeypatch, -): - # Disable FlashInfer FP8 scaled_mm kernel as it is incompatible with - # async TP patterns. No-op on H100 (kernel requires CC >= 100). - monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") - - model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - model_info.check_available_online(on_fail="skip") - - pp_size = 1 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": {"fuse_gemm_comms": async_tp_enabled}, - } - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - compare_two_settings( - model_id, - async_tp_args, - tp_args, - method="generate", - force_v1_runner=True, - ) - - -@create_new_process_for_each_test() -def test_async_tp_pass_nvfp4_correctness(num_gpus_available: int, caplog_mp_spawn): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 AsyncTP requires Blackwell or ROCm CDNA") - if is_blackwell and not has_flashinfer(): - pytest.skip("FlashInfer is required for the NVFP4 AsyncTP path") - - tp_size = 2 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - "--load-format", - "dummy", - *(["--linear-backend", "flashinfer_cutlass"] if is_blackwell else []), - "--hf-overrides", - json.dumps(NVFP4_HF_OVERRIDES), - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": True, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - if is_rocm_cdna: - compilation_config["inductor_compile_config"] = {"force_disable_caches": True} - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - env = {"VLLM_DISABLE_COMPILE_CACHE": "1"} if is_rocm_cdna else None - log_context = caplog_mp_spawn(logging.DEBUG) if is_rocm_cdna else nullcontext() - with log_context as logs: - compare_two_settings( - NVFP4_MODEL_ID, - async_tp_args, - tp_args, - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - if is_rocm_cdna: - assert logs is not None - matches_by_rank = { - rank: [ - int(count) - for line in logs.text.splitlines() - if line.startswith(f"(Worker_TP{rank} pid=") - for count in FUSION_LOG_PATTERNS["async_tp"].findall(line) - ] - for rank in range(tp_size) - } - assert all( - any(count > 0 for count in matches) for matches in matches_by_rank.values() - ), ( - "Expected positive NVFP4 AsyncTP rewrites on both TP workers, " - f"found {matches_by_rank}" - ) diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py deleted file mode 100644 index 10d25320c012..000000000000 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ /dev/null @@ -1,414 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""WARNING: This test runs in both single-node (4 GPUs) and multi-node -(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is -important to set the distributed backend to "mp" to avoid Ray scheduling -all workers in a node other than the head node, which can cause the test -to fail. -""" - -import json -import os -from dataclasses import dataclass -from typing import Literal, NamedTuple - -import pytest - -from vllm.config.compilation import CompilationMode -from vllm.config.model import RunnerOption -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_torch_equal_or_newer - -from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo -from ...utils import compare_all_settings, create_new_process_for_each_test - -logger = init_logger("test_sequence_parallel") - -VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_MODEL_INFO = _HfExamplesInfo(NVFP4_MODEL_ID) - - -class ParallelSetup(NamedTuple): - tp_size: int - pp_size: int - fuse_norm_quant: bool - fuse_act_quant: bool - eager_mode: bool - - -class SPTestOptions(NamedTuple): - multi_node_only: bool - load_format: str | None = None - model_info: _HfExamplesInfo | None = None - - -@dataclass -class SPTestSettings: - parallel_setups: list[ParallelSetup] - distributed_backends: list[str] - runner: RunnerOption - test_options: SPTestOptions - - @staticmethod - def fast( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [ - ParallelSetup( - tp_size=tp_base, - pp_size=pp_multiplier * pp_base, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_multiplier in [1, 2] - ] - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - @staticmethod - def fp8_quant( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [] - for fusion_val in [False, True]: - parallel_setups.append( - ParallelSetup( - tp_size=tp_base, - pp_size=pp_base, - fuse_norm_quant=fusion_val, - fuse_act_quant=fusion_val, - eager_mode=True, - ) - ) - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - def iter_params(self, model_id: str): - opts = self.test_options - - for parallel_setup in self.parallel_setups: - for backend in self.distributed_backends: - yield ( - model_id, - parallel_setup, - backend, - self.runner, - opts, - ) - - -def _build_sp_args( - model_id: str, - parallel_setup: ParallelSetup, - distributed_backend: str, - runner: RunnerOption, - test_options: SPTestOptions, - num_gpus_available: int, - use_inductor_graph_partition: bool, - fuse_gemm_comms: bool, - enable_prompt_embeds: bool, - *, - is_multimodal: bool, -) -> tuple[list[str], list[str]] | None: - ( - tp_size, - pp_size, - fuse_norm_quant, - fuse_act_quant, - eager_mode, - ) = parallel_setup - - multi_node_only = test_options.multi_node_only - load_format = test_options.load_format - - model_info = test_options.model_info or HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - - trust_remote_code = model_info.trust_remote_code - tokenizer_mode = model_info.tokenizer_mode - hf_overrides = dict(model_info.hf_overrides) - require_embed_inputs = model_info.require_embed_inputs - - if load_format == "dummy": - # Avoid OOM - text_overrides = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, - } - - if is_multimodal: - hf_overrides.update({"text_config": text_overrides}) - else: - hf_overrides.update(text_overrides) - else: - model_info.check_available_online(on_fail="skip") - - if num_gpus_available < tp_size * pp_size: - return None - if VLLM_MULTI_NODE and distributed_backend == "mp": - return None - if multi_node_only and not VLLM_MULTI_NODE: - return None - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - if eager_mode: - common_args.append("-cc.cudagraph_mode=none") - if runner != "auto": - common_args.extend(["--runner", runner]) - if trust_remote_code: - common_args.append("--trust-remote-code") - if tokenizer_mode: - common_args.extend(["--tokenizer-mode", tokenizer_mode]) - if load_format: - common_args.extend(["--load-format", load_format]) - if hf_overrides: - common_args.extend(["--hf-overrides", json.dumps(hf_overrides)]) - if require_embed_inputs: - common_args.extend( - [ - "--skip-tokenizer-init", - "--enable-prompt-embeds", - "--enable-mm-embeds", - ] - ) - elif enable_prompt_embeds: - common_args.append("--enable-prompt-embeds") - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [4, 8], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": fuse_gemm_comms, - "fuse_norm_quant": fuse_norm_quant, - "fuse_act_quant": fuse_act_quant, - "fuse_allreduce_rms": False, - "eliminate_noops": True, - "sp_min_token_num": 0, - }, - "use_inductor_graph_partition": use_inductor_graph_partition, - } - if not use_inductor_graph_partition: - compilation_config["splitting_ops"] = [] - - tp_sp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--pipeline-parallel-size", - str(pp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - return tp_args, tp_sp_args - - -def _compare_sp_settings( - model_id: str, - settings: list[tuple[list[str], list[str]]], - *, - method: Literal["generate", "encode"], -) -> None: - if not settings: - pytest.skip("No supported sequence-parallel configurations") - - settings_by_tp_args: dict[tuple[str, ...], list[list[str]]] = {} - for tp_args, tp_sp_args in settings: - settings_by_tp_args.setdefault(tuple(tp_args), []).append(tp_sp_args) - - for tp_args_key, grouped_tp_sp_args in settings_by_tp_args.items(): - all_args = [list(tp_args_key), *grouped_tp_sp_args] - compare_all_settings( - model_id, - all_args, - [None] * len(all_args), - method=method, - force_v1_runner=True, - ) - - -SP_TEXT_GENERATION_MODELS = { - # [Decoder-only] - "hmellor/tiny-random-LlamaForCausalLM": SPTestSettings.fast(), - "RedHatAI/Llama-3.2-1B-Instruct-FP8": SPTestSettings.fp8_quant(), -} - - -@pytest.mark.parametrize( - ("model_id", "settings"), - [ - pytest.param(model_id, settings, id=model_id) - for model_id, settings in SP_TEXT_GENERATION_MODELS.items() - ], -) -@create_new_process_for_each_test() -def test_tp_sp_generation( - model_id: str, - settings: SPTestSettings, - num_gpus_available, -): - # Skip FP8 SP-only test on sm89 (compute capability 8.9). - # An unknown capability (None) is left to fail in the test body rather than - # being silently skipped here. - capability = current_platform.get_device_capability() - if "fp8" in model_id.lower() and capability is not None and capability < (9, 0): - pytest.skip("FP8 reduction support begins with sm90 capable devices.") - - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for _, parallel_setup, backend, runner, test_options in settings.iter_params( - model_id - ): - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - backend, - runner, - test_options, - num_gpus_available, - use_inductor_graph_partition, - fuse_gemm_comms=False, # TODO: enable async TP - enable_prompt_embeds=False, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -# Focused regression test for the SP + prompt_embeds graph-rewrite path. -# Covers pp_size=1 (SP only) and pp_size=2 (SP + PP); kept small on purpose so -# we don't double the matrix of `test_tp_sp_generation` above. -SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [ - ParallelSetup( - tp_size=2, - pp_size=pp_size, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_size in [1, 2] -] - - -@create_new_process_for_each_test() -def test_tp_sp_generation_prompt_embeds( - num_gpus_available, -): - model_id = "hmellor/tiny-random-LlamaForCausalLM" - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for parallel_setup in SP_PROMPT_EMBEDS_PARALLEL_SETUPS: - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - distributed_backend="mp", - runner="auto", - test_options=SPTestOptions(multi_node_only=False, load_format=None), - num_gpus_available=num_gpus_available, - use_inductor_graph_partition=use_inductor_graph_partition, - fuse_gemm_comms=False, - enable_prompt_embeds=True, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -@create_new_process_for_each_test() -def test_tp_sp_nvfp4_generation(num_gpus_available: int): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 sequence parallelism requires Blackwell or ROCm CDNA") - - comparison = _build_sp_args( - NVFP4_MODEL_ID, - ParallelSetup( - tp_size=2, - pp_size=1, - fuse_norm_quant=True, - fuse_act_quant=True, - eager_mode=True, - ), - "mp", - "auto", - SPTestOptions( - multi_node_only=False, - load_format="dummy", - model_info=NVFP4_MODEL_INFO, - ), - num_gpus_available, - use_inductor_graph_partition=False, - fuse_gemm_comms=False, - enable_prompt_embeds=False, - is_multimodal=False, - ) - _compare_sp_settings( - NVFP4_MODEL_ID, - [] if comparison is None else [comparison], - method="generate", - ) diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 1f92ff5a1988..a4d996b48ddd 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/tests/compile/fusions_e2e/common.py @@ -21,8 +21,6 @@ class Matches(NamedTuple): # distributed ar_rms_fusion: int = 0 aiter_ar_rms_fusion: int = 0 - sequence_parallel: int = 0 - async_tp: int = 0 class ModelFusionInfo(NamedTuple): @@ -143,8 +141,4 @@ def has_cuda_graph_wrapper_metadata() -> bool: "aiter_ar_rms_fusion": re.compile( r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns" ), - "sequence_parallel": re.compile( - r"sequence_parallelism.py:\d+] Replaced (\d+) patterns" - ), - "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"), } diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index db7ba05a5734..8da7aa2a5a94 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -170,26 +170,8 @@ def run( for match_name in matches_check: log_matches = list(int(ms) for ms in log_matches_dict[match_name]) - # AR+RMS skips the largest range; SP skips the smallest. - # When both are enabled, AR+RMS activation count is - # model-dependent (hidden_size affects threshold), so derive - # from log data. - if ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - assert ( - len(log_matches) >= tp_size and len(log_matches) % tp_size == 0 - ), ( - f"Expected multiple of {tp_size} ar_rms log entries, " - f"found {len(log_matches)}" - ) - num_ranges_activated = len(log_matches) // tp_size - elif ( - match_name in ("ar_rms_fusion", "sequence_parallel") - and num_compile_ranges >= 2 - ): + # AR+RMS skips the largest range. + if match_name == "ar_rms_fusion" and num_compile_ranges >= 2: num_ranges_activated = num_compile_ranges - 1 else: num_ranges_activated = num_compile_ranges @@ -233,44 +215,6 @@ def run( f"Expecting at least {expected_matches - matches.ar_rms_fusion} " f"where ar+rms+quant was activated" ) - elif ( - match_name == "async_tp" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # AsyncTP only finds patterns on ranges where SP ran. - n_sp_ranges = num_compile_ranges - 1 - assert ( - sum(m == expected_matches for m in log_matches) - == tp_size * n_sp_ranges - ), ( - f"Expecting {expected_matches} async_tp on " - f"{tp_size * n_sp_ranges} SP-range entries, " - f"found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size, ( - f"Expecting 0 async_tp on {tp_size} small-range entries " - f"(no SP), found: {log_matches}" - ) - elif ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # SP consumes allreduce patterns first, so AR+RMS finds - # full matches only on the smallest range (no SP). - assert sum(m == expected_matches for m in log_matches) == tp_size, ( - f"Expecting {expected_matches} ar_rms on " - f"{tp_size} small-range entries, found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size * ( - num_ranges_activated - 1 - ), ( - f"Expecting 0 ar_rms on " - f"{tp_size * (num_ranges_activated - 1)} large-range " - f"entries (SP took precedence), found: {log_matches}" - ) - elif match_name == "act_quant_fusion": actual_match = match_table.get("activation_quant_fusion_pass", 0) assert actual_match == expected_matches * n_expected, ( @@ -305,17 +249,4 @@ def run( f"(found {len(log_matches)}) in:\n {log_holder.text}" ) - if match_name == "sequence_parallel" and num_compile_ranges >= 2: - log_matches = re.findall( - r"pass_manager.py:\d+] Skipping " - r".*SequenceParallelismPass.* with compile range", - log_holder.text, - ) - - n_expected = tp_size * (num_compile_ranges - num_ranges_activated) - assert len(log_matches) == n_expected, ( - f'Could not find {n_expected} "Skipping SequenceParallelismPass" ' - f"(found {len(log_matches)}) in:\n {log_holder.text}" - ) - return run diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 082ab0bf5296..d80d171984c3 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -73,8 +73,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -85,8 +83,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -96,8 +92,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -127,8 +121,6 @@ rms_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -138,8 +130,6 @@ matches=lambda n_layers: Matches( attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -149,8 +139,6 @@ norm_rope_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -161,8 +149,6 @@ norm_rope_fusion=n_layers, attn_quant_fusion=0, # attn + group quant not supported ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -192,9 +178,6 @@ # MLA attn + per-group FP8 quant attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - # TODO - # sequence_parallel= n_layers * 2 + 1, - # async_tp=n_layers * 2, ), ) @@ -213,8 +196,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), model_kwargs=( {"quantization_config": {"moe": {"activation": "mxfp8"}}} diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py deleted file mode 100644 index b0d32dad1c17..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_async_tp.py +++ /dev/null @@ -1,339 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import pytest - -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from ...utils import multi_gpu_test -from .common import ( - INDUCTOR_GRAPH_PARTITION, - AttentionBackendCase, - Matches, - custom_ops_combos, - is_blackwell, - nvfp4_kernel_exposes_input_quant_key, -) -from .models import ( - FLASHINFER_ATTN, - TRITON_ATTN, - llama3_8b, - llama3_8b_fp4, - llama3_8b_fp8, - llama4_scout_fp8, - qwen3_a3b, -) - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4], -) -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") -def test_tp2_async_tp_nvfp4_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - if nvfp4_kernel_exposes_input_quant_key(): - pytest.skip( - "NVFP4 kernel exposes input_quant_key; manual fusion fires " - "instead of compiler pass-based fusion" - ) - - # NVFP4 currently wires the all-gather + GEMM path only. - matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_act_quant=False, - fuse_attn_quant=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py deleted file mode 100644 index 3143143bec86..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos -from .models import TRITON_ATTN, llama3_8b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_async_tp_bf16_fusions( - inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test -): - n_layers = 4 - model_kwargs = dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ) - run_e2e_fusion_test( - llama3_8b.model_name, - llama3_8b.matches(n_layers), - model_kwargs, - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py b/tests/compile/fusions_e2e/test_tp2_rocm_moe.py deleted file mode 100644 index 9a228ca76c50..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm MoE coverage for the TP2 GEMM/collective fusion contract.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION -from .models import TRITON_ATTN, qwen3_a3b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION[:1]) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_moe_async_tp_bf16_fusions( - inductor_graph_partition, attn_backend, run_e2e_fusion_test -): - """MoE dense projections retain the Qwen SP/AsyncTP rewrite counts.""" - n_layers = 4 - run_e2e_fusion_test( - qwen3_a3b.model_name, - qwen3_a3b.matches(n_layers), - dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - moe_intermediate_size=128, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - num_experts=8, - num_experts_per_tok=2, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ), - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=["-rms_norm"], - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py deleted file mode 100644 index c0fd4c00a8f7..000000000000 --- a/tests/compile/passes/distributed/test_async_tp.py +++ /dev/null @@ -1,394 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import ( - multi_gpu_test, -) -from vllm.compilation.passes.fusion.collective_fusion import AsyncTPPass -from vllm.config import ( - CompilationConfig, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type -FP8_DTYPE = current_platform.fp8_dtype() - -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestMMRSModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.gate_proj = torch.nn.Parameter( - torch.empty((self.hidden_size * 2, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.gate_proj, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + reduce scatter in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - - # matrix multiplication - permute = self.gate_proj.permute(1, 0) - mm = torch.mm(view, permute) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_matmul_reduce_scatter.default] - - -class TestAGMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = torch.nn.Parameter( - torch.empty((hidden_size, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + all gather in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - all_gather = tensor_model_parallel_all_gather(view, dim=0) - permute = self.weight.permute(1, 0) - mm = torch.mm(all_gather, permute) - return mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_matmul.default] - - -class _BaseScaledMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = ( - torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - # Initialize scale_b for _scaled_mm. - self.scale_b = torch.ones(1, self.hidden_size, dtype=torch.float32) - - -class TestScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the scaled_mm + reduce scatter in the FX graph.""" - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - fp8_input, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(scaled_mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + scaled_mm in the FX graph.""" - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - all_gather, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - return scaled_mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -class TestCutlassScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the cutlass_scaled_mm + reduce scatter - in the FX graph - - """ - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - mm_out = torch.empty( - (fp8_input.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=input.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, fp8_input, self.weight, scale_a, self.scale_b, None - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm_out, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGCutlassScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + cutlass_scaled_mm - in the FX graph - """ - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - - mm_out = torch.empty( - (all_gather.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=all_gather.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, all_gather, self.weight, scale_a, self.scale_b, None - ) - return mm_out - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model", - [ - TestMMRSModel, - TestAGMMModel, - TestScaledMMRSModel, - TestAGScaledMMModel, - pytest.param( - TestCutlassScaledMMRSModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - pytest.param( - TestAGCutlassScaledMMModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("dynamic", [True, False]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_async_tp_pass_replace( - test_model: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, -): - num_processes = 2 - distributed_init_method = get_file_store_init_method() - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model, - batch_size, - seq_len, - hidden_size, - dtype, - dynamic, - distributed_init_method, - ), - nprocs=nprocs, - ) - - run_torch_spawn(async_tp_pass_on_test_model, num_processes) - - -def test_async_tp_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - async_tp_pass = object.__new__(AsyncTPPass) - async_tp_pass.compilation_config = vllm_config.compilation_config - - with pytest.raises( - AssertionError, match="AsyncTPPass requires full-graph compilation" - ): - async_tp_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def async_tp_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: torch.nn.Module, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, - distributed_init_method: str, - check_numerics: bool = False, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - } - ) - - # initialize distributed - init_distributed_environment( - world_size=world_size, - rank=local_rank, - distributed_init_method=distributed_init_method, - local_rank=local_rank, - ) - - # configure vllm config for SequenceParallelismPass - vllm_config = VllmConfig() - vllm_config.compilation_config = CompilationConfig( - pass_config=PassConfig( - fuse_gemm_comms=True, - ), - ) - vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - vllm_config.model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - - async_tp_pass = AsyncTPPass(vllm_config) - backend = TestBackend(async_tp_pass) - - assert ( - async_tp_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - async_tp_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - - model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor - - hidden_states = torch.randn( - (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False - ) - if check_numerics: - hidden_states += local_rank * 0.125 - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - actual = compiled_model(hidden_states) - - if check_numerics: - # These are new checks, so no existing tolerance is relaxed. - # Allow BF16 rounding vs. eager, but keep the lifetime check exact. - expected = model(hidden_states) - torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) - snapshot = actual.clone() - changed = compiled_model(-hidden_states) - torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(actual, snapshot, rtol=0, atol=0) - - assert async_tp_pass.matched_count == 1 - - # In pre-nodes, all gather or reduce scatter should exist, - # fused_matmul_reduce_scatter or fused_all_gather_matmul should not - backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) - - # In post-nodes, fused_matmul_reduce_scatter or \ - # fused_all_gather_matmul should exist - backend.check_after_ops(model.ops_in_model_after()) diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py deleted file mode 100644 index c40d75f6754a..000000000000 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ /dev/null @@ -1,343 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import TestFP8Layer, multi_gpu_test -from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass -from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass -from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass -from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass -from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass -from vllm.config import ( - CompilationConfig, - CUDAGraphMode, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - get_current_vllm_config, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) -from vllm.platforms import current_platform -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - -FP8_DTYPE = current_platform.fp8_dtype() -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestAllReduceRMSNormModel(torch.nn.Module): - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)] - - def forward(self, x): - z = torch.relu(x) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = torch.mm(y, self.w[0]) - x2 = tensor_model_parallel_all_reduce(z2) - - y2, resid = self.norm[1](x2, resid) - - z3 = torch.mm(y2, self.w[1]) - x3 = tensor_model_parallel_all_reduce(z3) - - y3, resid = self.norm[2](x3, resid) - - z4 = torch.mm(y3, self.w[2]) - x4 = tensor_model_parallel_all_reduce(z4) - - y4, resid = self.norm[3](x4, resid) - return y4 - - def ops_in_model_before(self): - return [torch.ops.vllm.all_reduce.default] - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model(self): - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - ] - - -class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module): - quant_key = kFp8StaticTensorSym - - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.vllm_config = get_current_vllm_config() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.fp8_linear_layers = [ - TestFP8Layer( - weight_shape=(hidden_size, hidden_size), - activation_quant_key=self.quant_key, - weight_quant_key=self.quant_key, - input_dtype=self.vllm_config.model_config.dtype, - ) - for i in range(3) - ] - - def forward(self, hidden_states): - # avoid having graph input be an arg to a pattern directly - z = torch.relu(hidden_states) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = self.fp8_linear_layers[0](y) - - x2 = tensor_model_parallel_all_reduce(z2) - y2, resid = self.norm[1](x2, resid) - - z3 = self.fp8_linear_layers[1](y2) - - x3 = tensor_model_parallel_all_reduce(z3) - y3, resid = self.norm[2](x3, resid) # use resid here - - z4 = self.fp8_linear_layers[2](y3) - x4 = tensor_model_parallel_all_reduce(z4) - y4, resid = self.norm[3](x4, resid) # use resid here - return y4 - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model_before(self): - return [ - torch.ops.vllm.all_reduce.default, - ] - - def ops_in_model(self): - if self.vllm_config.compilation_config.pass_config.fuse_norm_quant: - return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default] - else: - quant_ops = ( - [torch.ops._C.static_scaled_fp8_quant.default] - if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers) - else [torch.ops.aten.reciprocal] - ) - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - *quant_ops, - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model_cls, custom_ops", - [ - (TestAllReduceRMSNormModel, "+rms_norm"), - (TestAllReduceRMSNormModel, "-rms_norm"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,-quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,-quant_fp8"), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("fuse_norm_quant", [True, False]) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_sequence_parallelism_pass( - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - num_processes = 2 - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model_cls, - custom_ops, - batch_size, - seq_len, - hidden_size, - dtype, - fuse_norm_quant, - dynamic, - ), - nprocs=nprocs, - ) - - run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes) - - -def test_sequence_parallelism_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - sequence_parallelism_pass = object.__new__(SequenceParallelismPass) - sequence_parallelism_pass.compilation_config = vllm_config.compilation_config - sequence_parallelism_pass.min_token_num = 1 - - with pytest.raises( - AssertionError, - match="SequenceParallelismPass requires full-graph compilation", - ): - sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def sequence_parallelism_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "12345", - } - ) - - # initialize distributed - init_distributed_environment() - - # configure vllm config for SequenceParallelismPass - custom_ops_list = custom_ops.split(",") if custom_ops else [] - compilation_config = CompilationConfig( - splitting_ops=[], # avoid automatic rms_norm enablement - cudagraph_mode=CUDAGraphMode.NONE, # avoid piecewise warnings - custom_ops=custom_ops_list, - pass_config=PassConfig( - enable_sp=True, - fuse_norm_quant=fuse_norm_quant, - eliminate_noops=True, - ), - ) # NoOp needed for fusion - device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - vllm_config = VllmConfig( - model_config=model_config, - device_config=device_config, - compilation_config=compilation_config, - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - noop_pass = NoOpEliminationPass(vllm_config) - sequence_parallelism_pass = SequenceParallelismPass(vllm_config) - cleanup_pass = PostCleanupPass(vllm_config) - assert ( - sequence_parallelism_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - sequence_parallelism_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - passes_for_backend: list[VllmInductorPass] = [ - noop_pass, - sequence_parallelism_pass, - ] - - if fuse_norm_quant: - fusion_pass = RMSNormQuantFusionPass(vllm_config) - passes_for_backend.append(fusion_pass) - - passes_for_backend.append(cleanup_pass) - - backend = TestBackend(*passes_for_backend) - - model = test_model_cls(hidden_size) - - hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype) - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - compiled_model(hidden_states) - - assert sequence_parallelism_pass.matched_count == 4 - - # In pre-nodes, all reduce should be there, - # reduce scatter and all gather should not - for op in model.ops_in_model_before(): - assert backend.op_count(op, before=True) == 4 - - # In post-nodes, reduce scatter and all gather should be there, - # all reduce should not - for op in model.ops_in_model_after(): - assert backend.op_count(op, before=False) == 4 - - for op in model.ops_in_model(): - assert backend.op_count(op, before=False) > 0 diff --git a/tests/compile/rocm/__init__.py b/tests/compile/rocm/__init__.py deleted file mode 100644 index 208f01a7cb5e..000000000000 --- a/tests/compile/rocm/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/compile/rocm/test_async_tp.py b/tests/compile/rocm/test_async_tp.py deleted file mode 100644 index a4462265f8e7..000000000000 --- a/tests/compile/rocm/test_async_tp.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite.""" - -import pytest -import torch - -from tests.compile.passes.distributed.test_async_tp import ( - TestAGMMModel, - TestMMRSModel, - async_tp_pass_on_test_model, -) -from tests.utils import multi_gpu_test -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method - - -class TestRocmMMRSModel(TestMMRSModel): - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None) - return tensor_model_parallel_reduce_scatter(mm, dim=0) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -class TestRocmAGMMModel(TestAGMMModel): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__(hidden_size, dtype) - self.weight = torch.nn.Parameter( - torch.empty((hidden_size * 2, hidden_size)), requires_grad=False - ) - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - gathered = tensor_model_parallel_all_gather(view, dim=0) - return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract") -@pytest.mark.parametrize( - "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel] -) -@pytest.mark.parametrize("dynamic", [False, True]) -def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic): - """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives.""" - torch.multiprocessing.spawn( - async_tp_pass_on_test_model, - args=( - 2, - test_model, - 8, - 16, - 128, - torch.bfloat16, - dynamic, - get_file_store_init_method(), - True, - ), - nprocs=2, - ) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5a1ba13a57f1..4e7a756f8cbe 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -464,49 +464,41 @@ def test_should_split(): ( "cudagraph_capture_sizes", "max_cudagraph_capture_size", - "tp_size", - "enable_sp", "max_num_batched_tokens", "cudagraph_mode", "expected_max_size", ), [ - (None, None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([1, 2, 4], 4, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + (None, None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([1, 2, 4], 4, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), ( [1, 2, 4], 8, - 1, - False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError, ), - ([1, 256], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([], None, 1, False, 2048, CUDAGraphMode.NONE, 0), - (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0), + ([1, 256], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([], None, 2048, CUDAGraphMode.NONE, 0), + (None, 0, 2048, CUDAGraphMode.NONE, 0), # truncated to nearest multiple of 8 or 16 - (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + (None, 257, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), # max_num_batched_tokens <= max_cudagraph_capture_size should always be # captured even if not landing on a 16-stride step - (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), + (None, 2048, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), # max from list - ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), - # SP forces full-graph compilation, sizes are filtered by TP - ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), # limited by the max_tokens - ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), # the list should contain at least 1 element when use cudagraph - ([], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + ([], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), # the max capturing size should be >= 1 when use cudagraph - (None, 0, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + (None, 0, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), ], ) def test_cudagraph_sizes_post_init( cudagraph_capture_sizes, max_cudagraph_capture_size, - tp_size, - enable_sp, max_num_batched_tokens, cudagraph_mode, expected_max_size, @@ -515,10 +507,7 @@ def test_cudagraph_sizes_post_init( if expected_max_size == ValidationError: ctx = pytest.raises(expected_max_size) - with ( - ctx, - patch.object(current_platform, "device_count", return_value=tp_size), - ): + with ctx: kwargs = {} if cudagraph_capture_sizes is not None: kwargs["cudagraph_capture_sizes"] = cudagraph_capture_sizes @@ -526,18 +515,15 @@ def test_cudagraph_sizes_post_init( kwargs["max_cudagraph_capture_size"] = max_cudagraph_capture_size compilation_config = CompilationConfig( pass_config=PassConfig( - enable_sp=enable_sp, fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True, - sp_min_token_num=512 if enable_sp else None, ), cudagraph_mode=cudagraph_mode, **kwargs, ) engine_args = EngineArgs( model="facebook/opt-125m", - tensor_parallel_size=tp_size, max_num_seqs=min(max_num_batched_tokens, 128), max_num_batched_tokens=max_num_batched_tokens, compilation_config=compilation_config, @@ -728,35 +714,6 @@ def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded(): ) -def test_cudagraph_capture_sizes_respect_sequence_parallelism(): - """Sequence-parallel capture sizes stay divisible by tensor parallel size.""" - compilation_config = CompilationConfig( - cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE - ) - compilation_config.pass_config.enable_sp = True - config = _mock_config_for_cudagraph_sizes( - max_num_seqs=32, - num_speculative_tokens=16, - max_num_batched_tokens=32768, - compilation_config=compilation_config, - ) - config.parallel_config = SimpleNamespace(tensor_parallel_size=2) - config.update_sizes_for_sequence_parallelism = lambda sizes: ( - VllmConfig.update_sizes_for_sequence_parallelism(config, sizes) - ) - - with patch.object( - current_platform, - "is_device_capability_family", - return_value=False, - ): - VllmConfig._set_cudagraph_sizes(config) - - assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes) - assert 544 not in compilation_config.cudagraph_capture_sizes - assert compilation_config.max_cudagraph_capture_size == 512 - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") def test_real_vllm_config_caps_widest_ngram_decode_batch(): """Real config post-init preserves the default capture ceiling.""" @@ -1046,123 +1003,6 @@ def test_blackwell_cudagraph_default(is_blackwell, expected_max_size): ) -@pytest.mark.skipif( - not current_platform.support_static_graph_mode(), - reason="Skip if not cudagraph mode supported", -) -@pytest.mark.parametrize( - ( - "cudagraph_mode", - "use_inductor_graph_partition", - "expected_enable_sp", - "expected_cudagraph_mode", - "expected_piecewise_compile", - "expected_capture_sizes", - "expected_max_size", - ), - [ - (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4), - ( - CUDAGraphMode.FULL_DECODE_ONLY, - False, - True, - CUDAGraphMode.FULL_DECODE_ONLY, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - False, - True, - CUDAGraphMode.FULL, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - True, - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - [2, 4], - 4, - ), - ], -) -def test_sequence_parallelism_requires_full_graph_compilation( - cudagraph_mode: CUDAGraphMode, - use_inductor_graph_partition: bool, - expected_enable_sp: bool, - expected_cudagraph_mode: CUDAGraphMode, - expected_piecewise_compile: bool, - expected_capture_sizes: list[int], - expected_max_size: int, -): - with patch.object(current_platform, "device_count", return_value=2): - vllm_config = VllmConfig( - parallel_config=ParallelConfig(tensor_parallel_size=2), - scheduler_config=SchedulerConfig( - max_num_seqs=128, - max_num_batched_tokens=2048, - max_model_len=2048, - is_encoder_decoder=False, - ), - ) - vllm_config.model_config = MagicMock( - dtype=torch.float16, - enforce_eager=False, - is_moe=False, - disable_cascade_attn=False, - get_hidden_size=MagicMock(return_value=4096), - ) - vllm_config.compilation_config = CompilationConfig( - mode=CompilationMode.VLLM_COMPILE, - cudagraph_capture_sizes=[1, 2, 4, 15], - max_cudagraph_capture_size=None, - compile_sizes=["cudagraph_capture_sizes"], - use_inductor_graph_partition=use_inductor_graph_partition, - pass_config=PassConfig( - enable_sp=True, - fuse_gemm_comms=True, - fuse_norm_quant=True, - fuse_act_quant=True, - eliminate_noops=True, - sp_min_token_num=512, - ), - cudagraph_mode=cudagraph_mode, - ) - vllm_config.compilation_config.set_splitting_ops_for_v1( - all2all_backend=vllm_config.parallel_config.all2all_backend, - data_parallel_size=1, - ) - vllm_config._set_compile_ranges() - vllm_config._set_cudagraph_sizes() - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - == use_inductor_graph_partition - ) - assert ( - bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile - ) - assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp - assert ( - vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp - ) - assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode - assert ( - vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes - ) - assert ( - vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size - ) - assert ( - 511 in vllm_config.compilation_config.compile_ranges_endpoints - ) == expected_enable_sp - - def test_cached_compilation_config(default_vllm_config): import torch from torch._inductor.utils import run_and_get_code diff --git a/tests/compile/test_sequence_parallelism_threshold.py b/tests/compile/test_sequence_parallelism_threshold.py deleted file mode 100644 index 090b77b330ab..000000000000 --- a/tests/compile/test_sequence_parallelism_threshold.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm.compilation.passes.fusion.sequence_parallelism import ( - SP_MIN_HIDDEN_SIZE, - SP_MIN_PER_GPU_SIZE_MB, - get_sequence_parallelism_threshold, -) - - -class TestGetSequenceParallelismThreshold: - """Tests for get_sequence_parallelism_threshold function.""" - - def test_non_cuda_returns_none(self, mock_cuda_platform): - """Non-CUDA platforms should return None.""" - with mock_cuda_platform(is_cuda=False): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_unsupported_device_capability_returns_none(self, mock_cuda_platform): - """Unsupported device capabilities (e.g., sm80) should return None.""" - with mock_cuda_platform(capability=(8, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_small_hidden_size_returns_none(self, mock_cuda_platform): - """H100 with hidden_size below threshold should return None.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=4096, - tp_size=2, - element_size=2, # 4096 < 8192 - ) - assert result is None - - def test_h100_large_model_returns_threshold(self, mock_cuda_platform): - """H100 with large enough hidden_size should return calculated threshold.""" - with mock_cuda_platform(capability=(9, 0)): - hidden_size = 8192 - tp_size = 2 - element_size = 2 # float16/bfloat16 - - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - - # Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - MiB = 1024 * 1024 - expected = int( - (SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB) - // (hidden_size * element_size) - ) - assert result == expected - assert result == 1024 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # Boundary: exactly at min hidden size threshold, tp_size=1 - # (8 * 1 * 1024 * 1024) // (8192 * 2) = 512 - (8192, 1, 2, 512), - # Larger hidden size reduces token threshold - # (8 * 1 * 1024 * 1024) // (16384 * 2) = 256 - (16384, 1, 2, 256), - # Larger tp_size increases token threshold - # (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048 - (8192, 4, 2, 2048), - # Larger element_size (fp32) reduces token threshold - # (8 * 2 * 1024 * 1024) // (8192 * 4) = 512 - (8192, 2, 4, 512), - ], - ) - def test_threshold_calculation_variations( - self, mock_cuda_platform, hidden_size, tp_size, element_size, expected - ): - """Test threshold calculation with various parameter combinations.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_hidden_size_boundary(self, mock_cuda_platform): - """Test behavior at the exact hidden_size boundary.""" - with mock_cuda_platform(capability=(9, 0)): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90], - tp_size=2, - element_size=2, - ) - assert result is not None - - -# XPU-specific constants (must match sequence_parallelism.py values) -_XPU_MIN_HIDDEN_SIZE = 4096 -_XPU_MIN_PER_GPU_SIZE_MB = 8.0 - - -class TestGetSequenceParallelismThresholdXPU: - """Tests for get_sequence_parallelism_threshold on XPU platform.""" - - def test_xpu_small_hidden_size_returns_none(self, mock_xpu_platform): - """XPU with hidden_size below threshold should return None.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - def test_xpu_large_model_returns_threshold(self, mock_xpu_platform): - """XPU with hidden_size >= threshold should return calculated value.""" - with mock_xpu_platform(): - hidden_size = _XPU_MIN_HIDDEN_SIZE - tp_size = 2 - element_size = 2 - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - # (8 * 2 * 1024 * 1024) // (4096 * 2) = 2048 - MiB = 1024 * 1024 - expected = int( - (_XPU_MIN_PER_GPU_SIZE_MB * tp_size * MiB) // (hidden_size * element_size) - ) - assert result == expected - assert result == 2048 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # (8 * 1 * 1024 * 1024) // (4096 * 2) = 1024 - (4096, 1, 2, 1024), - # (8 * 4 * 1024 * 1024) // (4096 * 2) = 4096 - (4096, 4, 2, 4096), - # (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - (8192, 2, 2, 1024), - # (8 * 2 * 1024 * 1024) // (4096 * 4) = 1024 - (4096, 2, 4, 1024), - ], - ) - def test_xpu_threshold_calculation_variations( - self, mock_xpu_platform, hidden_size, tp_size, element_size, expected - ): - """Test XPU threshold calculation with various parameter combinations.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_xpu_hidden_size_boundary(self, mock_xpu_platform): - """Test behavior at the exact XPU hidden_size boundary.""" - with mock_xpu_platform(): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE, - tp_size=2, - element_size=2, - ) - assert result is not None diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 2d57236f8efa..5416fdc8ea6d 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -582,10 +582,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch): use_ubatching=False, enable_elastic_ep=False, ), - compilation_config=SimpleNamespace( - mode=CompilationMode.NONE, - pass_config=SimpleNamespace(enable_sp=False), - ), + compilation_config=SimpleNamespace(mode=CompilationMode.NONE), cache_config=SimpleNamespace( kv_sharing_fast_prefill=False, mamba_cache_mode="none", diff --git a/tests/test_config.py b/tests/test_config.py index 0098ecdc3dac..ebf76c76a0c6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -843,7 +843,6 @@ def test_breakable_cudagraph_platform_default( ("bi-off", 2048, (16, 8), 6144, 1, False), ("opt-out", 2048, (16, 8), 6144, 1, False), ("eager", 2048, (16, 8), 6144, 1, False), - ("sp", 2048, (16, 8), 6144, 1, False), ("fp16", 2048, (16, 8), 6144, 1, False), ("quantized", 2048, (16, 8), 6144, 1, False), ("untuned", 4096, (32, 32), 11008, 1, False), @@ -886,11 +885,7 @@ def test_batch_invariant_breakable_cudagraph( get_num_kv_heads=lambda pc: heads[1], ) config.parallel_config = SimpleNamespace(tensor_parallel_size=tp) - config.compilation_config = ( - CompilationConfig(pass_config=PassConfig(enable_sp=True)) - if case == "sp" - else CompilationConfig() - ) + config.compilation_config = CompilationConfig() try: assert config._maybe_enable_breakable_cudagraph() is expected if expected: @@ -1021,7 +1016,6 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( "FakeAttentionBackend", uniform_decode_query_len=4, use_v2_model_runner=use_v2_model_runner, - tensor_parallel_size=1, ) assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE @@ -1104,7 +1098,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, ) @@ -1117,7 +1110,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, is_profiling=True, @@ -3359,8 +3351,6 @@ def test_vllm_config_explicit_overrides(): take precedence over callable defaults, across different models and optimization levels. """ - from vllm.config.compilation import PassConfig - quantized_model = ModelConfig("RedHatAI/Llama-3.2-1B-FP8") moe_model = ModelConfig("deepseek-ai/DeepSeek-V2-Lite") regular_model = ModelConfig("Qwen/Qwen1.5-7B") diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 90adce187448..6dc20a817794 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -378,9 +378,6 @@ def run_model(scheduler_output, intermediate_tensors): worker = SimpleNamespace( vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace( - pass_config=SimpleNamespace(enable_sp=False) - ), parallel_config=SimpleNamespace( pipeline_parallel_size=2, distributed_executor_backend="mp" ), diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py deleted file mode 100644 index f2aa6d7d2f1d..000000000000 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ /dev/null @@ -1,1013 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Callable -from contextlib import suppress - -import torch -import torch._inductor.pattern_matcher as pm -import torch.distributed.distributed_c10d as c10d -import torch.fx as fx -from torch._inductor.pattern_matcher import PatternMatcherPass -from torch.distributed._symmetric_memory import enable_symm_mem_for_group - -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op - -from ..inductor_pass import enable_fake_mode -from ..vllm_inductor_pass import ( - VllmFusionPatternMatcherPass, - VllmInductorPass, - VllmPatternMatcherPass, - VllmPatternReplacement, -) - -FP8_DTYPE = current_platform.fp8_dtype() - -logger = init_logger(__name__) - - -def _flashinfer_scaled_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - bias: torch.Tensor | None = None, - scale_result: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, - use_fast_accum: bool = False, -) -> None: - # Import lazily to avoid a circular import during module initialization - # when docs or other tooling import the pass without FlashInfer. - from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm_out - - assert bias is None, "FlashInfer symm_mem adapter does not support bias" - assert scale_result is None, ( - "FlashInfer symm_mem adapter does not support result scaling" - ) - assert not use_fast_accum, ( - "FlashInfer symm_mem adapter does not support use_fast_accum" - ) - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer symm_mem adapter expects 2D inputs and output" - ) - assert scale_a.numel() == 1 and scale_b.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - - flashinfer_scaled_fp8_mm_out( - A, - B, - scale_a, - scale_b, - out=out, - out_dtype=out_dtype or out.dtype, - ) - - -def _flashinfer_fp4_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - alpha: torch.Tensor, - out_dtype: torch.dtype | None = None, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> None: - from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm_out - - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs and output" - ) - flashinfer_scaled_fp4_mm_out( - A, - B, - scale_a, - scale_b, - alpha, - out=out, - out_dtype=out_dtype or out.dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter_fake( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - result_shape = list(output_shape) - result_shape[orig_scatter_dim] //= world_size - return torch.empty( - result_shape, - dtype=out_dtype or torch.bfloat16, - device=A.device, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert orig_scatter_dim == 0 and scatter_dim_after_maybe_reshape == 0, ( - "FlashInfer symm_mem adapter currently only supports scatter_dim=0" - ) - world_size = c10d._resolve_process_group(group_name).size() - assert A.ndim == 2 and B.ndim == 2, "FlashInfer symm_mem adapter expects 2D inputs" - assert A.is_contiguous(), "FlashInfer symm_mem adapter expects contiguous A" - assert A_scale.numel() == 1 and B_scale.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - assert A.shape[0] % world_size == 0, ( - "FlashInfer symm_mem adapter expects M divisible by world size" - ) - - kwargs = { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - return torch.distributed._symmetric_memory._fused_scaled_matmul_reduce_scatter_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A=A, - B=B, - A_scale=A_scale, - kwargs=kwargs, - out_dtype=out_dtype, - reduce_op=reduce_op, - orig_scatter_dim=orig_scatter_dim, - scatter_dim_after_maybe_reshape=scatter_dim_after_maybe_reshape, - group_name=group_name, - output_shape=output_shape, - ) - - -def fused_all_gather_flashinfer_scaled_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_scaled_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer symm_mem adapter currently only supports gather_dim=0" - ) - _, outputs = torch.distributed._symmetric_memory._fused_all_gather_matmul_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A_shard=A_shard, - Bs=[B], - A_scale=A_scale, - kwargs_list=[ - { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - ], - out_dtypes=[out_dtype], - gather_dim=gather_dim, - group_name=group_name, - return_A=False, - ) - return outputs[0] - - -def fused_all_gather_flashinfer_fp4_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_fp4_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer FP4 symm_mem adapter currently only supports gather_dim=0" - ) - assert A_shard.ndim == 2 and A_scale_shard.ndim == 2 and B.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs" - ) - if view_a_scale_as_fp8: - A_scale_shard = A_scale_shard.view(torch.float8_e4m3fn) - - group = c10d._resolve_process_group(group_name) - world_size = group.size() - output = A_shard.new_empty( - A_shard.shape[0] * world_size, - B.shape[1], - dtype=out_dtype or torch.bfloat16, - ) - output_shards = output.chunk(world_size) - - A = A_shard.new_empty(A_shard.shape[0] * world_size, A_shard.shape[1]) - A_scale = A_scale_shard.new_empty( - A_scale_shard.shape[0] * world_size, - A_scale_shard.shape[1], - ) - - def fp4_shard_consumer(shards: list[torch.Tensor], rank: int) -> None: - _flashinfer_fp4_mm_out( - shards[0], - B, - scale_a=shards[1], - scale_b=B_scale, - alpha=alpha, - out=output_shards[rank], - out_dtype=out_dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - torch.distributed._symmetric_memory._pipelined_multi_all_gather_and_consume( - [A_shard, A_scale_shard], - fp4_shard_consumer, - [A, A_scale], - group_name, - False, - ) - return output - - -direct_register_custom_op( - op_name="fused_flashinfer_scaled_matmul_reduce_scatter", - op_func=fused_flashinfer_scaled_matmul_reduce_scatter, - fake_impl=fused_flashinfer_scaled_matmul_reduce_scatter_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_scaled_matmul", - op_func=fused_all_gather_flashinfer_scaled_matmul, - fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_fp4_matmul", - op_func=fused_all_gather_flashinfer_fp4_matmul, - fake_impl=fused_all_gather_flashinfer_fp4_matmul_fake, -) - - -class BasePattern: - def __init__(self, dtype: torch.dtype, device: str | None) -> None: - self.dtype = dtype - self.device = device - self.tp = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - - -class GEMMReduceScatterPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - mul = torch.empty([16, 4], device=self.device, dtype=self.dtype) - mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - return [mul, mm_weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - if self.rocm_gemm: - mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None) - else: - mm = torch.ops.aten.mm.default(mul, mm_weight) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( - mul, - mm_weight.t() if self.rocm_gemm else mm_weight, - "sum", - scatter_dim=0, - group_name=self.tp.device_group.group_name, - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern): - rocm_gemm = True - - -class AllGatherGEMMPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([4, 4], device=self.device, dtype=self.dtype) - weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [x, weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - if self.rocm_gemm: - return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None) - return torch.ops.aten.mm.default(all_gather, weight) - - def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul( - x, - [weight.t() if self.rocm_gemm else weight], - gather_dim=0, - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmAllGatherGEMMPattern(AllGatherGEMMPattern): - rocm_gemm = True - - -class ScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - return [input, mm_weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - scaled_mm = torch.ops.aten._scaled_mm.default( - input, - mat2=mat2, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - scaled_mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - return [x, weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - return torch.ops.aten._scaled_mm.default( - all_gather, - mat2=weight, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class CutlassScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype) - return [input, mm_weight, scale_a, scale_b, cutlass_mm_output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=cutlass_mm_output, - a=input, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - cutlass_scaled_mm[1], - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherCutlassScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - s2 = weight.shape[1] - output = torch.empty([s1, s2], device=self.device, dtype=self.dtype) - - return [x, weight, scale_a, scale_b, output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=output, - a=all_gather, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - return cutlass_scaled_mm[1] - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class FlashInferBMMFP8ReduceScatterPattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_2d = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - bmm = torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(a_2d, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - output = torch.ops.aten.reshape.default(bmm, list(bmm.shape[1:])) - return torch.ops.vllm.reduce_scatter.default( - output, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_flashinfer_scaled_matmul_reduce_scatter.default( - a_2d, - b_2d, - a_scale, - b_scale, - "sum", - 0, - 0, - self.tp.device_group.group_name, - [a_2d.shape[0], b_2d.shape[1]], - self.dtype, - ) - - return _replacement - - -class FlashInferAllGatherBMMFP8Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_shard_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(all_gather, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - fused = torch.ops.vllm.fused_all_gather_flashinfer_scaled_matmul.default( - a_shard_2d, - b_2d, - a_scale, - b_scale, - 0, - self.tp.device_group.group_name, - self.dtype, - ) - return torch.ops.aten.unsqueeze.default(fused, 0) - - return _replacement - - -class FlashInferAllGatherFP4Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def __init__( - self, - dtype: torch.dtype, - device: str | None, - backend: str, - use_8x4_sf_layout: bool, - a_scale_view: str, - ) -> None: - super().__init__(dtype, device) - self.backend = backend - self.use_8x4_sf_layout = use_8x4_sf_layout - self.a_scale_view = a_scale_view - - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - b_2d = torch.empty([8, 16], device=self.device, dtype=torch.uint8) - a_scale_shard = torch.empty([128, 4], device=self.device, dtype=torch.int32) - b_scale = torch.empty([4, 128], device=self.device, dtype=torch.uint8) - alpha = torch.empty([], device=self.device, dtype=torch.float32) - return [ - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - ] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - all_gather_a = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - all_gather_a_scale = torch.ops.vllm.all_gather.default( - a_scale_shard, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - a_scale = all_gather_a_scale - if self.a_scale_view in ("float8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.float8_e4m3fn) - if self.a_scale_view in ("uint8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.uint8) - return torch.ops.vllm.flashinfer_mm_fp4.default( - all_gather_a, - b_2d, - a_scale, - b_scale, - alpha, - self.dtype, - self.use_8x4_sf_layout, - self.backend, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_all_gather_flashinfer_fp4_matmul.default( - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - 0, - self.tp.device_group.group_name, - self.dtype, - self.a_scale_view in ("float8", "float8_uint8"), - self.use_8x4_sf_layout, - self.backend, - ) - - return _replacement - - -class AsyncTPPass(VllmFusionPatternMatcherPass): - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config, pass_name="async_tp_pass") - - enable_symm_mem_for_group(get_tp_group().device_group.group_name) - GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass) - - AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass) - if current_platform.is_rocm() and self.model_dtype == torch.bfloat16: - RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - RocmAllGatherGEMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - - # These fusions are enabled only for bfloat16 models because - # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling - # only supports bfloat16 as the output dtype. - if self.model_dtype == torch.bfloat16: - ScaledMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - AllGatherScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - if hasattr(torch.ops._C, "cutlass_scaled_mm"): - CutlassScaledMMReduceScatterPattern( - self.model_dtype, self.device - ).register(self.pm_pass) - AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - with suppress(ImportError): - import vllm.utils.flashinfer # noqa: F401 - if hasattr(torch.ops.vllm, "bmm_fp8"): - self.register( - FlashInferAllGatherBMMFP8Pattern(self.model_dtype, self.device) - ) - self.register( - FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device) - ) - if hasattr(torch.ops.vllm, "flashinfer_mm_fp4"): - for backend in ("cutlass", "cudnn"): - for a_scale_view in ("float8_uint8", "uint8"): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - backend, - use_8x4_sf_layout=False, - a_scale_view=a_scale_view, - ) - ) - for use_8x4_sf_layout in (False, True): - for a_scale_view in ("float8",): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "trtllm", - use_8x4_sf_layout=use_8x4_sf_layout, - a_scale_view=a_scale_view, - ) - ) - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "cute-dsl", - use_8x4_sf_layout=False, - a_scale_view="float8", - ) - ) - # NVFP4 reduce-scatter does not need scale communication: FP4 - # scales are consumed by the local GEMM and only BF16 partial - # outputs are reduced. Keep this PR scoped to the all-gather - # path; reduce-scatter needs a dedicated FP4 producer rather - # than the existing FP8-style helper. - - self.dump_patterns(config, self.pm_pass) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - # This pass is applied on top of the sequence parallelism pass, - # which is only supported in fullgraph compilation mode. - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "AsyncTPPass requires full-graph compilation" - return True - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.pm_pass.apply(graph) - VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count - logger.debug("Replaced %s patterns", self.matched_count) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py deleted file mode 100644 index b1a39b1e0146..000000000000 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ /dev/null @@ -1,620 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import functools -from collections.abc import Callable, Sequence -from typing import Any - -import torch -import torch._inductor.pattern_matcher as pm -import torch.fx as fx -from torch._higher_order_ops.auto_functionalize import auto_functionalized -from torch._inductor.pattern_matcher import PatternMatcherPass - -import vllm.ir.ops -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) - -from ..inductor_pass import enable_fake_mode -from ..utility.noop_elimination import NoOpEliminationPass -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass -from .matcher_utils import MatcherQuantFP8 - -logger = init_logger(__name__) - -if hasattr(torch.ops._C, "scaled_fp4_quant"): - SCALED_FP4_QUANT_OUT_OVERLOAD = torch.ops._C.scaled_fp4_quant.out - SCALED_FP4_QUANT_DEFAULT_OVERLOAD = torch.ops._C.scaled_fp4_quant.default - -# Min hidden size per device capability for sequence parallelism -# Only apply sequence parallelism for models with hidden_size >= threshold -SP_MIN_HIDDEN_SIZE: dict[int, int] = { - 90: 8192, # H100: only for models with hidden_size >= 8192 - 100: 8192, # Blackwell family: only for models with hidden_size >= 8192 -} - -# Min size per GPU per device capability for sequence parallelism -# Total min size = min_per_gpu_size * tp_size -# This ensures the threshold scales appropriately with tensor parallelism -SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = { - 90: 8, # 8MB per GPU for H100 - # Use a more conservative threshold on Blackwell so TP8 starts later. - 100: 32, -} - - -def get_sequence_parallelism_threshold( - hidden_size: int, - tp_size: int, - element_size: int, -) -> int | None: - """Calculate the minimum token threshold for applying sequence parallelism. - - Returns None if sequence parallelism should not be applied based on model size. - - Branching logic based on device capability: - - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability] - - If not, returns None (SP disabled for small models on this device) - - If yes, calculates threshold based on per-GPU size - - Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) // - (hidden_size * element_size) - """ - from vllm.platforms import current_platform - - if current_platform.is_xpu(): - min_hidden_size = 4096 - min_per_gpu_size_mb = 8.0 - elif current_platform.is_cuda(): - capability = current_platform.get_device_capability() - if capability is None: - return None - - # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket. - if current_platform.is_device_capability_family(100): - device_capability = 100 - else: - device_capability = capability.to_int() - - # Check if device has configured thresholds - _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) - _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) - if _hidden is None or _gpu_mb is None: - return None - min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb - else: - return None - - # Only apply sequence parallelism for models meeting the size threshold - if hidden_size < min_hidden_size: - return None - - MiB = 1024 * 1024 - min_size = min_per_gpu_size_mb * MiB * tp_size - return int(min_size // (hidden_size * element_size)) - - -def get_first_out_wrapper( - fn: Callable[..., Sequence[torch.Tensor]], -) -> Callable[..., torch.Tensor]: - @functools.wraps(fn) - def wrapper(*args: Any) -> torch.Tensor: - return fn(*args)[0] - - return wrapper - - -class _SequenceParallelPatternHelper: - """Helper for sequence parallelism patterns.""" - - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - self.epsilon = epsilon - self.dtype = dtype - self.device = device - self.tp_group = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() - - def _all_reduce(self, x: torch.Tensor) -> torch.Tensor: - return tensor_model_parallel_all_reduce(x) - - def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.reduce_scatter.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def _all_gather(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs) - - def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) - - -class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight - return [self.empty([1, 8, 4]), self.empty([4])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rmsnorm = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - - return rmsnorm, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - - rmsnorm = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - all_gather = self._all_gather(rmsnorm) - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [ - residual, - mm_1, - rms_norm_weights, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - return rmsnorm[0], rmsnorm[1] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # The pattern matcher replaces from the end of the graph - # (last layer first). At the time each match is replaced, - # the preceding layer has NOT been replaced yet, so - # `residual` is still full-size and the slice below is - # correct. Once the preceding layer IS replaced, its - # residual output shrinks to [local_len, H], and this - # slice becomes semantically incorrect (e.g. for rank > 0, - # the indices would be out of bounds). However, since the - # symbolic output shape equals the input shape, - # NoOpEliminationPass (called at the end of - # SequenceParallelismPass.__call__) removes these slices - # before the graph is ever executed or compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - all_gather = self._all_gather(rmsnorm[0]) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, rmsnorm[1] - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight, scale - return [self.empty([1, 8, 4]), self.empty([4]), self.empty_f32([1, 1])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - return quant, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - scale = torch.empty([1, 1], device=self.device, dtype=torch.float32) - - return [residual, mm_1, rms_norm_weights, scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - return quant, residual_out - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # See MiddleAllReduceRMSNormPattern.replacement for a - # detailed explanation of the temporary slice below: - # it is correct when first inserted, becomes semantically - # incorrect after the preceding layer is replaced, and is - # removed by NoOpEliminationPass before the graph is compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, residual_out - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - input = self.empty([8, 16]) - weight = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [input, weight, input_global_scale, quant_output, output_scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], all_reduce, quant[2] - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return ( - self._all_gather(quant[0]), - reduce_scatter, - self._all_gather(quant[1]), - ) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = self.empty([8, 16]) - residual = self.empty([8, 16]) - rms_norm_weights = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [ - residual, - mm_1, - rms_norm_weights, - input_global_scale, - quant_output, - output_scale, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], residual_out, quant[2] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # Keep this slice in sync with the non-quantized SP replacement: - # once the previous SP pattern fires, it becomes a no-op. - reduce_scatter = self._reduce_scatter(mm_1) - residual = residual[0 : reduce_scatter.size(0), ...] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return self._all_gather(quant[0]), residual_out, self._all_gather(quant[1]) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class SequenceParallelismPass(VllmPatternMatcherPass): - """This pass enables sequence parallelism for models. - It identifies patterns where an AllReduce operation is followed by - an RMSNorm (or RMSNorm and then Quantization) operation. - These patterns are replaced with a ReduceScatter operation, followed by - a local RMSNorm/Quantization, and then an AllGather operation. - - The general transformation is: - Input -> AllReduce -> RMSNorm -> Output - becomes - Input -> ReduceScatter -> RMSNorm -> AllGather -> Output - - While this pass itself does not directly yield performance improvements, - it lays the groundwork for subsequent fusion passes, such as - GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can - significantly reduce communication overhead and improve overall model - performance. - - This pass is only supported when compiling the whole graph (fullgraph - mode, i.e. using Inductor graph partition or empty splitting_ops). - Piecewise compilation is not supported because the residual tensor - gets split across TP ranks, causing size mismatches at subgraph - boundaries. - - This pass splits up the residual tensor across TP ranks and hence divides - its size. The pattern matcher starts at the end of the graph (last layer - first), so when each replacement inserts a residual slice, the preceding - layer has not been replaced yet and the slice is correct. Once the - preceding layer IS replaced, its residual output shrinks and the slice - becomes semantically incorrect (out-of-bounds indices for rank > 0). - The graph is never executed in this intermediate state — - NoOpEliminationPass removes these slices based on symbolic shape equality - (input shape == output shape) before the graph is compiled. - """ - - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config) - - # Get min_token_num threshold - # Read min_token_num from config (calculated during config init) - self.min_token_num = None - if config.model_config is not None: - pass_config = config.compilation_config.pass_config - self.min_token_num = pass_config.sp_min_token_num - - if self.min_token_num is not None: - # Take the min to avoid exceeding max_num_batched_tokens - max_batched = config.scheduler_config.max_num_batched_tokens - if max_batched is not None: - self.min_token_num = min(self.min_token_num, max_batched) - logger.debug_once( - f"Sequence parallelism min token threshold: {self.min_token_num}", - scope="global", - ) - - # Used to clean up redundant views created temporarily - # to circumvent residual shape change issues - self.noop_cleanup = NoOpEliminationPass(config) - self.noop_cleanup.pass_name = f"{self.pass_name}.{self.noop_cleanup.pass_name}" - - self.patterns: PatternMatcherPass = PatternMatcherPass( - pass_name="sequence_parallelism_pass" - ) - - for epsilon in [1e-5, 1e-6]: - # RMSNorm + Static FP8 quantization patterns - FirstAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - if "SCALED_FP4_QUANT_OUT_OVERLOAD" in globals(): - FirstAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - # Normal RMSNorm patterns - FirstAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - MiddleAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - self.dump_patterns(config, self.patterns) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - """Determines if sequence parallelism should be applied for the given - compile range. - - SP is only beneficial for larger batch sizes where the communication - overhead is amortized. For small batches, the overhead of splitting - and gathering tensors across TP ranks outweighs the benefits. - - Returns False (SP disabled) when: - - min_token_num is None (SP disabled for this device/config) - - The compile range starts below the minimum token threshold - """ - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "SequenceParallelismPass requires full-graph compilation" - - # min_token_num is None when SP is disabled for this device/config - # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size) - if self.min_token_num is None: - return False - - # Only apply SP when batch size meets the minimum threshold - return compile_range.start >= self.min_token_num - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.patterns.apply(graph) - logger.debug("Replaced %s patterns", self.matched_count) - # Clean up reshape nodes - self.noop_cleanup(graph) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index df0e58b21081..197f2208a15c 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -35,13 +35,11 @@ RMSNormReshapeFusionPass, ) from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass - from .fusion.sequence_parallelism import SequenceParallelismPass from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass - from .fusion.collective_fusion import AsyncTPPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass @@ -156,11 +154,6 @@ def configure(self, config: VllmConfig) -> None: if self.pass_config.eliminate_noops: self.passes += [NoOpEliminationPass(config)] - if self.pass_config.enable_sp: - self.passes += [SequenceParallelismPass(config)] - if self.pass_config.fuse_gemm_comms: - self.passes += [AsyncTPPass(config)] - if enable_transformers_norm_canonicalization: self.passes += [AddRMSNormFusionPass(config)] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index cbd132ab5b68..fc371d94b493 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -133,12 +133,6 @@ class PassConfig: """Fuse the custom Attention and MLAAttention + quant ops.""" eliminate_noops: bool = Field(default=True) """Eliminate no-op ops.""" - enable_sp: bool = None # type: ignore[assignment] - """Enable sequence parallelism. Requires TP>1. Automatically disabled - if the model's hidden_size is too small for SP to be beneficial - (threshold is device-capability dependent).""" - fuse_gemm_comms: bool = None # type: ignore[assignment] - """Enable async TP.""" fuse_allreduce_rms: bool = None # type: ignore[assignment] """Enable flashinfer allreduce fusion.""" enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment] @@ -183,12 +177,6 @@ class PassConfig: 8: 1, # 1MB }, }, where key is the device capability""" - sp_min_token_num: int | None = None - """The minimum number of tokens above which vllm should use - sequence parallelism. Specified as an integer token count. - Unspecified will fallback to default values which are compute - capability and world size dependent.""" - # TODO(luka) better pass enabling system. def flashinfer_max_size(self, world_size: int) -> int | None: @@ -231,8 +219,6 @@ def compute_hash(self) -> str: "fuse_norm_quant", "fuse_act_quant", "fuse_attn_quant", - "enable_sp", - "fuse_gemm_comms", "fuse_allreduce_rms", "fuse_act_padding", "fuse_mla_dual_rms_norm", @@ -1244,25 +1230,6 @@ def set_splitting_ops_for_v1( self.cudagraph_mode = CUDAGraphMode.FULL self.splitting_ops = [] - if ( - not self.use_inductor_graph_partition - and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms) - and self.splitting_ops - ): - logger.warning_once( - "Sequence parallelism requires full-graph compilation when " - "use_inductor_graph_partition is off. Setting splitting_ops " - "to an empty list to preserve SP and async TP." - ) - self.splitting_ops = [] - if self.cudagraph_mode.has_piecewise_cudagraphs(): - logger.warning_once( - "Sequence parallelism is incompatible with piecewise " - "cudagraph when use_inductor_graph_partition is off. " - "Setting cudagraph_mode to FULL." - ) - self.cudagraph_mode = CUDAGraphMode.FULL - # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible if ( all2all_backend == "deepep_high_throughput" @@ -1403,7 +1370,6 @@ def resolve_cudagraph_mode_and_sizes( min_cg_attn_backend: str | None, uniform_decode_query_len: int = 1, use_v2_model_runner: bool = False, - tensor_parallel_size: int = 1, kv_cache_config: "KVCacheConfig | None" = None, max_num_reqs: int | None = None, is_profiling: bool = False, @@ -1530,10 +1496,7 @@ def resolve_cudagraph_mode_and_sizes( and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and uniform_decode_query_len > 1 ): - self.adjust_cudagraph_sizes_for_spec_decode( - uniform_decode_query_len, - tensor_parallel_size, - ) + self.adjust_cudagraph_sizes_for_spec_decode(uniform_decode_query_len) # For Mamba models with FULL decode cudagraphs, each decode # sequence needs one Mamba cache block. The decode cudagraph @@ -1562,25 +1525,8 @@ def resolve_cudagraph_mode_and_sizes( self.cudagraph_mode = cudagraph_mode return cudagraph_mode - def adjust_cudagraph_sizes_for_spec_decode( - self, uniform_decode_query_len: int, tensor_parallel_size: int - ): + def adjust_cudagraph_sizes_for_spec_decode(self, uniform_decode_query_len: int): multiple_of = uniform_decode_query_len - if tensor_parallel_size > 1 and self.pass_config.enable_sp: - multiple_of = max(uniform_decode_query_len, tensor_parallel_size) - if ( - multiple_of % uniform_decode_query_len != 0 - or multiple_of % tensor_parallel_size != 0 - ): - raise ValueError( - f"Can't determine cudagraph shapes that are both a " - f"multiple of {uniform_decode_query_len} " - f"(num_speculative_tokens + 1) required by spec-decode " - f"and {tensor_parallel_size} (tensor_parallel_size) " - f"required by sequence parallelism please adjust " - f"num_speculative_tokens or disable sequence parallelism" - ) - if not self.cudagraph_capture_sizes or multiple_of <= 1: return @@ -1600,7 +1546,7 @@ def adjust_cudagraph_sizes_for_spec_decode( if len(rounded_sizes) == 0: raise ValueError( f"No valid cudagraph sizes after rounding to multiple of {multiple_of} " - f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)" + f"(num_speculative_tokens + 1)" f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}" f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})" f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8fc89f1591cd..3fee812700b0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -145,12 +145,10 @@ class OptimizationLevel(IntEnum): PerformanceMode = Literal["balanced", "interactivity", "throughput"] IS_QUANTIZED = False -IS_DENSE = False # The optimizations that depend on these properties currently set to False # in all cases. # if model_config is not None: # IS_QUANTIZED = lambda c: c.model_config.is_quantized() -# IS_DENSE = lambda c: not c.model_config.is_model_moe() # See https://github.com/vllm-project/vllm/issues/25689. @@ -259,8 +257,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": False, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": False, "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, @@ -282,8 +278,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, @@ -305,8 +299,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -328,8 +320,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -828,9 +818,6 @@ def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool: or model.dtype != torch.bfloat16 or model.quantization is not None or not current_platform.is_cuda() - # Sequence parallelism / async TP are torch.compile passes. - or self.compilation_config.pass_config.enable_sp - or self.compilation_config.pass_config.fuse_gemm_comms ): return False family = bi._get_tuned_matmul_arch_family( @@ -1770,26 +1757,6 @@ def __post_init__(self): ) self.compilation_config.mode = CompilationMode.NONE - # TODO: This is a stopgap that turns sequence parallelism / async TP - # off under batch invariance. Make the sequence-parallel reduce-scatter - # path batch-invariant instead so they can stay enabled (#56370). - # Sequence parallelism / async TP rewrite all_reduce + rms_norm into a - # reduce-scatter whose reduction order depends on the batch, so they - # are not batch-invariant. Decide this before the breakable-CUDA-graph - # default below, which declines to auto-enable when they are on. - pass_config = self.compilation_config.pass_config - if envs.VLLM_BATCH_INVARIANT and ( - pass_config.enable_sp or pass_config.fuse_gemm_comms - ): - logger.warning_once( - "Disabling sequence parallelism and async TP " - "(pass_config.enable_sp / fuse_gemm_comms) when " - "VLLM_BATCH_INVARIANT is enabled: the reduce-scatter path " - "is not batch-invariant (see vllm-project/vllm#56370)." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph() if not breakable_cudagraph_enabled and ( @@ -1920,38 +1887,6 @@ def has_blocked_weights(): breakable_cudagraph_enabled=breakable_cudagraph_enabled ) - # async tp is built on top of sequence parallelism and requires it. - pass_config = self.compilation_config.pass_config - if pass_config.fuse_gemm_comms: - pass_config.enable_sp = True - if pass_config.enable_sp: - if self.parallel_config.tensor_parallel_size == 1: - logger.warning_once("Sequence Parallelism requires TP>1, disabling") - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - else: - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - if pass_config.sp_min_token_num is None: - logger.warning_once( - "Model hidden_size too small for the SP " - "threshold heuristic, disabling. To force SP, " - "set pass_config.sp_min_token_num manually." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - from vllm.utils.torch_utils import HAS_OPAQUE_TYPE if HAS_OPAQUE_TYPE: @@ -2184,30 +2119,6 @@ def has_blocked_weights(): data_parallel_size=effective_dp_size, ) - if self.compilation_config.pass_config.enable_sp: - # With pipeline parallelism, native rms norm tracing errors due to - # incorrect residual shape. - # Use custom rms norm to unblock. In the future, - # the pass will operate on higher-level IR to avoid the issue. - # TODO: https://github.com/vllm-project/vllm/issues/27894 - if self.compilation_config.mode != CompilationMode.VLLM_COMPILE: - logger.warning_once( - "Sequence parallelism is enabled, but running in wrong " - "vllm compile mode: %s.", - self.compilation_config.mode, - ) - - if self.parallel_config.pipeline_parallel_size > 1: - if "-rms_norm" not in self.compilation_config.custom_ops: - self.compilation_config.custom_ops.append("+rms_norm") - else: - logger.warning_once( - "Sequence parallelism not supported with " - "native rms_norm when using %s, " - "this will likely lead to an error.", - "pipeline parallelism", - ) - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( @@ -2395,29 +2306,6 @@ def has_blocked_weights(): # Log the custom passes that are enabled self.compilation_config.pass_config.log_enabled_passes() - def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list: - # remove the sizes that not multiple of tp_size when - # enable sequence parallelism - removed_sizes = [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size != 0 - ] - if removed_sizes: - logger.warning( - "Batch sizes %s are removed because they are not " - "multiple of tp_size %d when " - "sequence parallelism is enabled", - removed_sizes, - self.parallel_config.tensor_parallel_size, - ) - - return [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size == 0 - ] - def _set_max_num_scheduled_tokens(self): """In most cases, the scheduler may schedule a batch with as many tokens as the worker is configured to handle. @@ -2665,16 +2553,6 @@ def request_counts(max_reqs: int) -> list[int]: # de-duplicate and sort the sizes cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes)) - if ( - self.parallel_config.tensor_parallel_size > 1 - and self.compilation_config.pass_config.enable_sp - ): - # Sequence parallelism only captures TP-divisible sizes, so a - # wider non-divisible decode batch cannot be captured under SP. - cudagraph_capture_sizes = self.update_sizes_for_sequence_parallelism( - cudagraph_capture_sizes - ) - # user-specific compilation_config.max_cudagraph_capture_size get # truncated to valid_max_size when they are inconsistent. valid_max_size = ( @@ -2764,37 +2642,6 @@ def _set_compile_ranges(self): "allreduce-rms fusion will be enabled for all num_tokens." ) - # Add the compile ranges for sequence parallelism - if compilation_config.pass_config.enable_sp: - pass_config = compilation_config.pass_config - - # Calculate min_token_num if not explicitly provided - # User override works regardless of hidden_size - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - min_token_num = pass_config.sp_min_token_num - max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens - if min_token_num is not None and ( - max_num_batched_tokens is not None - and min_token_num < max_num_batched_tokens - and min_token_num > 1 - ): - # Add endpoint at min_token_num - 1 to ensure SP applies - # starting from min_token_num - # This creates ranges: [1, min-1] (no SP), [min, max] (SP applies) - computed_compile_ranges_endpoints.append(min_token_num - 1) - if compilation_config.pass_config.fuse_rope_kvcache: max_token_num = ( compilation_config.pass_config.rope_kvcache_fusion_max_token_num @@ -3124,12 +2971,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported: list[str] = [] speculative_config = self.speculative_config - if ( - self.compilation_config.pass_config.enable_sp - and self.parallel_config.tensor_parallel_size > 1 - ): - unsupported.append("sequence parallelism") - # V2 does not implement the external_launcher (torchrun) PP-output # broadcast that V1 uses to keep all ranks in sync (broadcast_pp_output). if ( diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 6b7db6d0dc55..355a15094db0 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -352,7 +352,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: pass_config = compilation_config.pass_config fusion_passes_to_disable = { - "fuse_gemm_comms": "Async TP", "fuse_allreduce_rms": "AllReduce + RMSNorm fusion", "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index d7fa807c33b3..4ce055d3f865 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -1091,47 +1091,6 @@ def flashinfer_scaled_fp4_mm( ) -def flashinfer_scaled_fp4_mm_out( - a: torch.Tensor, - b: torch.Tensor, - block_scale_a: torch.Tensor, - block_scale_b: torch.Tensor, - alpha: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None, - use_8x4_sf_layout: bool, - backend: str, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2 - assert a.stride(-1) == 1 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert out.device.type == "cuda" - - if backend in ("cutlass", "cudnn"): - if block_scale_a.dtype != torch.uint8: - block_scale_a = block_scale_a.view(torch.uint8) - if block_scale_b.dtype != torch.uint8: - block_scale_b = block_scale_b.view(torch.uint8) - - from flashinfer import mm_fp4 as flashinfer_mm_fp4_ - - flashinfer_mm_fp4_( - a, - b, - block_scale_a, - block_scale_b, - alpha, - out_dtype or out.dtype, - out=out, - block_size=16, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - return out - - def flashinfer_scaled_fp8_mm( a: torch.Tensor, b: torch.Tensor, @@ -1162,38 +1121,6 @@ def flashinfer_scaled_fp8_mm( return output -def flashinfer_scaled_fp8_mm_out( - a: torch.Tensor, - b: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert scale_a.numel() == 1 and scale_b.numel() == 1 - assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn - assert out.device.type == "cuda" - assert a.is_contiguous() - - from flashinfer import bmm_fp8 as bmm_fp8_ - - bmm_fp8_( - a.unsqueeze(0), - # FlashInfer expects the weight in the same column-major view layout - # consumed by flashinfer_scaled_fp8_mm, so keep the transposed view. - b.unsqueeze(0), - scale_a, - scale_b, - out_dtype or out.dtype, - out.unsqueeze(0), - "auto", - ) - return out - - def flashinfer_quant_nvfp4_8x4_sf_layout( a: torch.Tensor, a_global_sf: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1337,9 +1264,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "use_trtllm_attention", "flashinfer_mxfp4_quantize", "flashinfer_scaled_fp4_mm", - "flashinfer_scaled_fp4_mm_out", "flashinfer_scaled_fp8_mm", - "flashinfer_scaled_fp8_mm_out", "flashinfer_quant_nvfp4_8x4_sf_layout", "flashinfer_fp8_blockscale_gemm", "should_use_flashinfer_for_blockscale_fp8_gemm", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 373042a41502..61d6197ceb9d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -730,7 +730,6 @@ def initialize_kv_cache( attn_cg_support.min_cg_attn_backend, self.decode_query_len, use_v2_model_runner=True, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0cc3fcc29a8f..21b0f67e2fe8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -116,7 +116,7 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks from vllm.utils.platform_utils import num_compute_units @@ -222,7 +222,6 @@ ) from vllm.v1.worker.utils import ( EncoderTimingStats, - is_residual_scattered_for_sp, raise_if_nan_logits, ) from vllm.v1.worker.workspace import lock_workspace @@ -3340,20 +3339,9 @@ def sync_and_gather_intermediate_tensors( ) -> IntermediateTensors: assert self.intermediate_tensors is not None - tp = self.vllm_config.parallel_config.tensor_parallel_size - is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens) - - # When sequence parallelism is enabled, the "residual" tensor is - # sharded across TP ranks. All-gather it here because downstream - # QKV + Attention needs the full residual before the SP split point. if sync_self: assert intermediate_tensors is not None for k, v in intermediate_tensors.items(): - is_scattered = k == "residual" and is_rs - if is_scattered: - local_len = num_tokens // tp - v = get_tp_group().all_gather(v[:local_len], dim=0) - self.intermediate_tensors[k][:num_tokens].copy_( v[:num_tokens], non_blocking=True ) @@ -3448,14 +3436,6 @@ def _pool( async_output_copy_stream=self._get_or_create_async_output_copy_stream(), ) - def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: - # Pad tokens to multiple of tensor_parallel_size when - # enabled collective fusion for SP - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if self.compilation_config.pass_config.enable_sp and tp_size > 1: - return round_up(num_scheduled_tokens, tp_size) - return num_scheduled_tokens - def _prepare_padding_mask( self, num_tokens_unpadded: int, num_tokens_padded: int ) -> torch.Tensor: @@ -3944,8 +3924,6 @@ def _determine_batch_execution_and_padding( ) has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora - num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens) - def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): return self.cudagraph_dispatcher.dispatch( num_tokens=num_tokens, @@ -3957,18 +3935,9 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) cudagraph_mode, batch_descriptor = dispatch_cudagraph( - num_tokens_padded, disable_full=use_cascade_attn or has_encoder_output + num_tokens, disable_full=use_cascade_attn or has_encoder_output ) num_tokens_padded = batch_descriptor.num_tokens - if self.compilation_config.pass_config.enable_sp: - assert ( - batch_descriptor.num_tokens - % self.vllm_config.parallel_config.tensor_parallel_size - == 0 - ), ( - "Sequence parallelism requires num_tokens to be " - "a multiple of tensor parallel size" - ) # Extra coordination when running data-parallel since we need to coordinate # across ranks @@ -4453,15 +4422,9 @@ def execute_model( sample_hidden_states = hidden_states[logits_indices] if not get_pp_group().is_last_rank: - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, num_tokens_padded - ) - } get_pp_group().send_tensor_dict( hidden_states.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) logits = None else: @@ -7135,7 +7098,6 @@ def _check_and_update_cudagraph_mode( min_cg_attn_backend, self.uniform_decode_query_len, use_v2_model_runner=False, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d4c64c554273..70427d85a103 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -12,7 +12,6 @@ from types import NoneType from typing import TYPE_CHECKING, Any, cast -import numpy as np import torch import torch.nn as nn @@ -98,7 +97,6 @@ maybe_apply_startup_plan, maybe_save_startup_plan, ) -from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -1273,45 +1271,10 @@ def execute_model( intermediate_tensors = None forward_pass = scheduler_output.total_num_scheduled_tokens > 0 - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - all_gather_tensors = {} - compilation_config = self.vllm_config.compilation_config - parallel_config = self.vllm_config.parallel_config - - if ( - parallel_config.pipeline_parallel_size > 1 - and compilation_config.pass_config.enable_sp - and forward_pass - ): - # currently only supported by V1 GPUModelRunner - assert not self.use_v2_model_runner - num_scheduled_tokens_np = np.array( - list(scheduler_output.num_scheduled_tokens.values()), - dtype=np.int32, - ) - # TODO(lucas): This is pretty gross; ideally we should only ever call - # `_determine_batch_execution_and_padding` once (will get called again - # in `execute_model`) but this requires a larger refactor of PP. - _, batch_desc, _, _, _ = ( - self.model_runner._determine_batch_execution_and_padding( - num_tokens=num_scheduled_tokens, - num_reqs=len(num_scheduled_tokens_np), - num_scheduled_tokens_np=num_scheduled_tokens_np, - max_num_scheduled_tokens=num_scheduled_tokens_np.max(), - use_cascade_attn=False, # TODO(lucas): Handle cascade attention - ) - ) - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, batch_desc.num_tokens - ) - } - if forward_pass and not get_pp_group().is_first_rank: tensor_dict, comm_handles, comm_postprocess = ( get_pp_group().irecv_tensor_dict( all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) ) assert tensor_dict is not None @@ -1349,7 +1312,6 @@ def execute_model( handles = get_pp_group().isend_tensor_dict( output.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) self._pp_send_work = handles[1:] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 422ba8885fd7..b2bf6543a07c 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -10,7 +10,7 @@ import numpy as np import torch -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -794,35 +794,6 @@ def get_uniform_decode_token_count( return None -def is_residual_scattered_for_sp( - vllm_config: VllmConfig, num_input_tokens: int -) -> bool: - """Check if the residual tensor is scattered for sequence parallelism. - - The residual tensor is scattered across tensor parallel ranks when sequence - parallelism and tensor parallelism is enabled. SP is only supported in - full-graph compilation mode. - """ - if not vllm_config.compilation_config.pass_config.enable_sp: - return False - - tp = vllm_config.parallel_config.tensor_parallel_size - - if tp == 1: - return False - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - or not vllm_config.compilation_config.splitting_ops - ), "Sequence parallelism requires full-graph compilation" - - # When sequence parallelism is enabled, we always pad num_input_tokens - # to be a multiple of tensor_parallel_size (tp) earlier. - assert num_input_tokens % tp == 0 - - return True - - @dataclass class EncoderTimingStats: """Per-request timing statistics for encoder forward pass."""