@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bd7157a071

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[torch.compile] Enable attention and allreduce fusion without custom ops enabled (#24604)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Luka Govedič committed Oct 17, 2025 at 10:10 UTC bd7157a071e28521979255e21526d887a87fc1a7

28 files changed +1499 -701

.buildkite/test-pipeline.yaml

+30 -12

@@ -416,8 +416,8 @@ steps:

416 - pytest -v -s compile/test_basic_correctness.py

417 - pytest -v -s compile/piecewise/

418

419 -- label: PyTorch Fullgraph Test # 20min

420 - timeout_in_minutes: 30

419 +- label: PyTorch Fullgraph Test # 22min

420 + timeout_in_minutes: 35

421 mirror_hardwares: [amdexperimental]

422 torch_nightly: true

423 source_file_dependencies:

@@ -425,6 +425,7 @@ steps:

425 - tests/compile

426 commands:

427 - pytest -v -s compile/test_full_graph.py

428 + - pytest -v -s compile/test_fusions_e2e.py

429

430 - label: Kernels Core Operation Test # 48min

431 timeout_in_minutes: 75

@@ -807,8 +808,8 @@ steps:

808 # Whisper needs spawn method to avoid deadlock

809 - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

810

810 -- label: Blackwell Test # 38 min

811 - timeout_in_minutes: 60

811 +- label: Blackwell Test # 21 min

812 + timeout_in_minutes: 30

813 working_dir: "/vllm-workspace/"

814 gpu: b200

815 # optional: true

@@ -821,8 +822,6 @@ steps:

822 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

823 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

824 - vllm/v1/attention/backends/flashinfer.py

824 - - vllm/compilation/fusion.py

825 - - vllm/compilation/fusion_attn.py

825 commands:

826 - nvidia-smi

827 - python3 examples/offline_inference/basic/chat.py

@@ -839,15 +838,32 @@ steps:

838 - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py

839 - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py

840 - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

841 + - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py

842 + - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py

843 - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py

844 - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py

844 - # Fusion

845 - - pytest -v -s tests/compile/test_fusion_all_reduce.py

846 - - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern

845 - pytest -v -s tests/kernels/moe/test_flashinfer.py

846 +

847 +- label: Blackwell Fusion Tests # 30 min

848 + timeout_in_minutes: 40

849 + working_dir: "/vllm-workspace/"

850 + gpu: b200

851 + source_file_dependencies:

852 + - csrc/quantization/fp4/

853 + - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

854 + - vllm/v1/attention/backends/flashinfer.py

855 + - vllm/compilation/

856 + # can affect pattern matching

857 + - vllm/model_executor/layers/layernorm.py

858 + - vllm/model_executor/layers/activation.py

859 + - vllm/model_executor/layers/quantization/input_quant_fp8.py

860 + commands:

861 + - nvidia-smi

862 + - pytest -v -s tests/compile/test_fusion_attn.py

863 - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py

849 - - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py

850 - - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py

864 + # this runner has 2 GPUs available even though num_gpus=2 is not set

865 + - pytest -v -s tests/compile/test_fusion_all_reduce.py

866 + - pytest -v -s tests/compile/test_fusions_e2e.py

867

868 - label: Blackwell GPT-OSS Eval

869 timeout_in_minutes: 60

@@ -1100,7 +1116,7 @@ steps:

1116 - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

1117

1118 ##### H200 test #####

1103 -- label: Distrubted Tests (H200) # optional

1119 +- label: Distributed Tests (H200) # optional

1120 gpu: h200

1121 optional: true

1122 working_dir: "/vllm-workspace/"

@@ -1108,6 +1124,8 @@ steps:

1124 commands:

1125 - pytest -v -s tests/compile/test_async_tp.py

1126 - pytest -v -s tests/compile/test_sequence_parallelism.py

1127 + - pytest -v -s tests/compile/test_fusion_all_reduce.py

1128 + - pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm

1129 - pytest -v -s tests/distributed/test_context_parallel.py

1130 - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1 --dp-size=2 --max-model-len 2048

1131

csrc/layernorm_kernels.cu

@@ -392,6 +392,8 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]

392 torch::Tensor& residual, // [..., hidden_size]

393 torch::Tensor& weight, // [hidden_size]

394 double epsilon) {

395 + TORCH_CHECK(weight.scalar_type() == input.scalar_type());

396 + TORCH_CHECK(input.scalar_type() == residual.scalar_type());

397 TORCH_CHECK(residual.is_contiguous());

398 TORCH_CHECK(weight.is_contiguous());

399 int hidden_size = input.size(-1);

csrc/layernorm_quant_kernels.cu

@@ -229,6 +229,8 @@ void fused_add_rms_norm_static_fp8_quant(

229 double epsilon) {

230 TORCH_CHECK(out.is_contiguous());

231 TORCH_CHECK(residual.is_contiguous());

232 + TORCH_CHECK(residual.scalar_type() == input.scalar_type());

233 + TORCH_CHECK(weight.scalar_type() == input.scalar_type());

234 int hidden_size = input.size(-1);

235 int input_stride = input.stride(-2);

236 int num_tokens = input.numel() / hidden_size;

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -145,7 +145,11 @@ void rms_norm_dynamic_per_token_quant(

145 if (scale_ub.has_value()) {

146 TORCH_CHECK(out.dtype() == kFp8Type);

147 }

148 + TORCH_CHECK(weight.dtype() == input.dtype());

149 TORCH_CHECK(scales.dtype() == torch::kFloat32);

150 + if (residual) {

151 + TORCH_CHECK(residual->scalar_type() == input.scalar_type());

152 + }

153

154 VLLM_DISPATCH_FLOATING_TYPES(

155 input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {

tests/compile/backend.py

+22 -3

@@ -3,16 +3,22 @@

4 import weakref

5 from collections.abc import Callable, Sequence

6 +from contextlib import nullcontext

7 from copy import deepcopy

9 +import depyf

10 from torch import fx

11 from torch._ops import OpOverload

12 +from torch.fx._utils import lazy_format_graph_code

14 from vllm.compilation.fx_utils import find_op_nodes

15 from vllm.compilation.inductor_pass import InductorPass

16 from vllm.compilation.pass_manager import with_pattern_match_debug

17 from vllm.compilation.vllm_inductor_pass import VllmInductorPass

18 from vllm.config import VllmConfig, get_current_vllm_config

19 +from vllm.logger import init_logger

20 +

21 +logger = init_logger("vllm.tests.compile.backend")

24 class LazyInitPass(InductorPass):

@@ -45,20 +51,32 @@ class TestBackend:

52 def __init__(self, *passes: InductorPass | Callable[[fx.Graph], None]):

53 self.custom_passes = list(passes)

48 - compile_config = get_current_vllm_config().compilation_config

49 - self.inductor_config = compile_config.inductor_compile_config

54 + vllm_config = get_current_vllm_config()

55 + compile_config = vllm_config.compilation_config

56 + # Deepcopy to allow multiple TestBackend instances to use the same VllmConfig

57 + self.inductor_config = deepcopy(compile_config.inductor_compile_config)

58 self.inductor_config["force_disable_caches"] = True

59 self.inductor_config["post_grad_custom_post_pass"] = self.post_pass

61 + if debug_dump_path := vllm_config.compile_debug_dump_path():

62 + logger.debug("Dumping depyf output to %s", debug_dump_path)

63 + self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix())

64 + else:

65 + self.debug_ctx = nullcontext()

66 +

67 def __call__(self, graph: fx.GraphModule, example_inputs):

68 self.graph_pre_compile = deepcopy(graph)

69 from torch._inductor.compile_fx import compile_fx

57 - return compile_fx(graph, example_inputs, config_patches=self.inductor_config)

71 + with self.debug_ctx:

72 + return compile_fx(

73 + graph, example_inputs, config_patches=self.inductor_config

74 + )

76 @with_pattern_match_debug

77 def post_pass(self, graph: fx.Graph):

78 self.graph_pre_pass = deepcopy(graph)

79 + lazy_format_graph_code("graph_pre_pass", graph.owning_module)

81 VllmInductorPass.dump_prefix = 0

82 for pass_ in self.custom_passes:

@@ -68,6 +86,7 @@ class TestBackend:

86 VllmInductorPass.dump_prefix = None

88 self.graph_post_pass = deepcopy(graph)

89 + lazy_format_graph_code("graph_post_pass", graph.owning_module)

90 # assign by reference, will reflect the final state of the graph

91 self.final_graph = graph

tests/compile/test_full_graph.py

+59 -73

@@ -1,8 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import logging

4 import tempfile

5 +from pathlib import Path

6 from typing import Any

8 import pytest

@@ -10,8 +10,6 @@ import torch

11 from tests.quantization.utils import is_quant_method_supported

12 from vllm import LLM, SamplingParams

13 -from vllm.attention.backends.registry import _Backend

14 -from vllm.attention.selector import global_force_attn_backend_context_manager

13 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig

14 from vllm.platforms import current_platform

15 from vllm.utils import is_torch_equal_or_newer

@@ -22,23 +20,24 @@ from ..utils import create_new_process_for_each_test

20 def models_list(*, all: bool = True, keywords: list[str] | None = None):

21 TEST_MODELS: list[tuple[str, dict[str, Any]]] = [

22 ("facebook/opt-125m", {}),

25 - (

26 - "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",

27 - {

28 - "dtype": torch.float16,

29 - },

30 - ),

23 (

24 "neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic",

33 - {

34 - "dtype": torch.float16,

35 - },

25 + {"dtype": torch.float16},

26 ),

37 - ("neuralmagic/Llama-3.2-1B-Instruct-quantized.w8a8", {}),

27 ("meta-llama/Llama-3.2-1B-Instruct", {}),

28 ]

30 if all:

31 + TEST_MODELS.extend(

32 + [

33 + ("neuralmagic/Llama-3.2-1B-Instruct-quantized.w8a8", {}),

34 + (

35 + "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change",

36 + {"dtype": torch.float16},

37 + ),

38 + ]

39 + )

40 +

41 # TODO: figure out why this fails.

42 if False and is_quant_method_supported("gguf"): # noqa: SIM223

43 TEST_MODELS.append(

@@ -83,31 +82,38 @@ def models_list(*, all: bool = True, keywords: list[str] | None = None):

82 "compilation_mode",

83 [CompilationMode.DYNAMO_TRACE_ONCE, CompilationMode.VLLM_COMPILE],

84 )

86 -@pytest.mark.parametrize("model_info", models_list(all=True))

85 +@pytest.mark.parametrize("model, model_kwargs", models_list(all=True))

86 @create_new_process_for_each_test()

87 def test_full_graph(

88 monkeypatch: pytest.MonkeyPatch,

90 - model_info: tuple[str, dict[str, Any]],

89 + model: str,

90 + model_kwargs: dict[str, Any],

91 compilation_mode: int,

92 ):

93 - model, model_kwargs = model_info

93 + if (

94 + "w8a8" in model

95 + or "w8w8" in model

96 + and current_platform.has_device_capability((10, 0))

97 + ):

98 + # int8 removed on Blackwell:

99 + pytest.skip("int8 support removed on Blackwell")

100

101 with monkeypatch.context():

102 print(f"MODEL={model}")

103

98 - run_model(compilation_mode, model, model_kwargs)

104 + run_model(compilation_mode, model, **model_kwargs)

105

106

107 # TODO(luka) add other supported compilation config scenarios here

108 @pytest.mark.parametrize(

103 - "compilation_config, model_info",

109 + "compilation_config, model, model_kwargs",

110 [

111 # additional compile sizes, only some of the models

112 (

113 CompilationConfig(mode=CompilationMode.VLLM_COMPILE, compile_sizes=[1, 2]),

108 - model,

114 + *model_info,

115 )

110 - for model in models_list(all=False)

116 + for model_info in models_list(all=False)

117 ]

118 + [

119 # RMSNorm + quant fusion, only 8-bit quant models

@@ -117,18 +123,19 @@ def test_full_graph(

123 custom_ops=["+rms_norm"],

124 pass_config=PassConfig(enable_fusion=True, enable_noop=True),

125 ),

120 - model,

126 + *model_info,

127 )

122 - for model in models_list(keywords=["FP8-dynamic", "quantized.w8a8"])

128 + for model_info in models_list(keywords=["FP8-dynamic", "quantized.w8a8"])

129 ]

130 + [

131 # Test depyf integration works

132 (

133 CompilationConfig(

134 mode=CompilationMode.VLLM_COMPILE,

129 - debug_dump_path=tempfile.gettempdir(),

135 + debug_dump_path=Path(tempfile.gettempdir()),

136 ),

131 - ("facebook/opt-125m", {}),

137 + "facebook/opt-125m",

138 + {},

139 ),

140 ]

141 + [

@@ -142,9 +149,9 @@ def test_full_graph(

149 cudagraph_mode=CUDAGraphMode.PIECEWISE,

150 compile_sizes=[1, 2],

151 ),

145 - model,

152 + *model_info,

153 )

147 - for model in models_list(all=False)

154 + for model_info in models_list(all=False)

155 if is_torch_equal_or_newer("2.9.0.dev")

156 ],

157 )

@@ -152,16 +159,24 @@ def test_full_graph(

159 @create_new_process_for_each_test()

160 def test_custom_compile_config(

161 compilation_config: CompilationConfig,

155 - model_info: tuple[str, dict[str, Any]],

162 + model: str,

163 + model_kwargs: dict[str, Any],

164 ):

165 + if (

166 + "w8a8" in model

167 + or "w8w8" in model

168 + and current_platform.has_device_capability((10, 0))

169 + ):

170 + # int8 removed on Blackwell:

171 + pytest.skip("int8 support removed on Blackwell")

172 +

173 if compilation_config.use_inductor_graph_partition and not is_torch_equal_or_newer(

174 "2.9.0.dev"

175 ):

176 pytest.skip("inductor graph partition is only available in PyTorch 2.9+")

177

162 - model, model_kwargs = model_info

178 print(f"MODEL={model}")

164 - run_model(compilation_config, model, model_kwargs)

179 + run_model(compilation_config, model, **model_kwargs)

180

181

182 @pytest.mark.parametrize(

@@ -176,50 +191,16 @@ def test_fp8_kv_scale_compile(compilation_mode: int):

191 "calculate_kv_scales": True,

192 "max_model_len": 512,

193 }

179 - run_model(compilation_mode, model, model_kwargs)

194 + run_model(compilation_mode, model, **model_kwargs)

195

196

182 -def test_inductor_graph_partition_attn_fusion(caplog_vllm):

183 - if not is_torch_equal_or_newer("2.9.0.dev"):

184 - pytest.skip("inductor graph partition is only available in PyTorch 2.9+")

185 -

186 - model = "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"

187 - compilation_config = CompilationConfig(

188 - mode=CompilationMode.VLLM_COMPILE,

189 - use_inductor_graph_partition=True,

190 - cudagraph_mode=CUDAGraphMode.PIECEWISE,

191 - custom_ops=["+quant_fp8"],

192 - pass_config=PassConfig(enable_attn_fusion=True, enable_noop=True),

197 +def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):

198 + compilation_config = (

199 + compile_config

200 + if isinstance(compile_config, CompilationConfig)

201 + else CompilationConfig(level=compile_config)

202 )

194 - model_kwargs = {

195 - "kv_cache_dtype": "fp8",

196 - "max_model_len": 1024,

197 - }

198 - with (

199 - caplog_vllm.at_level(logging.DEBUG),

200 - global_force_attn_backend_context_manager(_Backend.FLASHINFER),

201 - ):

202 - run_model(compilation_config, model, model_kwargs)

203

204 - try:

205 - assert "Fused quantization onto 48 attention nodes" in caplog_vllm.text, (

206 - caplog_vllm.text

207 - )

208 - except AssertionError:

209 - # Note: this message is only triggered when the compilation goes

210 - # through the custom pass. Due to multiple layers of cache on

211 - # PyTorch side, the compilation of a graph may be cached such

212 - # that custom pass directly goes through cache. In this case,

213 - # we go through this branch and assert that the pass is not

214 - # triggered.

215 - assert "Fused quantization" not in caplog_vllm.text

216 -

217 -

218 -def run_model(

219 - compile_config: int | CompilationConfig,

220 - model: str,

221 - model_kwargs: dict[str, Any],

222 -):

204 prompts = [

205 "Hello, my name is",

206 "The president of the United States is",

@@ -227,12 +208,17 @@ def run_model(

208 "The future of AI is",

209 ]

210 sampling_params = SamplingParams(temperature=0)

211 + # Allow override from model_kwargs

212 + model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}

213 + model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}

214 +

215 + # No cudagraphs by default

216 + if compilation_config.cudagraph_mode is None:

217 + compilation_config.cudagraph_mode = CUDAGraphMode.NONE

218 +

219 llm = LLM(

220 model=model,

232 - enforce_eager=True,

233 - tensor_parallel_size=1,

234 - disable_custom_all_reduce=True,

235 - compilation_config=compile_config,

221 + compilation_config=compilation_config,

222 **model_kwargs,

223 )

224 outputs = llm.generate(prompts, sampling_params)

tests/compile/test_functionalization.py

+61 -53

@@ -11,7 +11,13 @@ from vllm.compilation.fusion import RMSNormQuantFusionPass

11 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func

12 from vllm.compilation.noop_elimination import NoOpEliminationPass

13 from vllm.compilation.post_cleanup import PostCleanupPass

14 -from vllm.config import CompilationConfig, PassConfig, VllmConfig

14 +from vllm.config import (

15 + CompilationConfig,

16 + ModelConfig,

17 + PassConfig,

18 + VllmConfig,

19 + set_current_vllm_config,

20 +)

21 from vllm.model_executor.layers.activation import SiluAndMul

22 from vllm.model_executor.layers.layernorm import RMSNorm

23 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

@@ -48,8 +54,7 @@ class TestSiluMul(torch.nn.Module):

54 return y

56 def example_inputs(self, num_tokens=32, hidden_size=128):

51 - dtype = torch.float16 if TEST_FP8 else torch.float32

52 - return (torch.rand(num_tokens, hidden_size * 2, dtype=dtype),)

57 + return (torch.rand(num_tokens, hidden_size * 2),)

59 def ops_in_model(self, do_fusion):

60 if TEST_FP8 and do_fusion:

@@ -67,15 +72,11 @@ class TestFusedAddRMSNorm(torch.nn.Module):

72 self.hidden_size = hidden_size

73 self.intermediate_size = intermediate_size

70 - dtype = torch.float16 if TEST_FP8 else torch.float32

71 -

75 self.gate_proj = torch.nn.Parameter(

73 - torch.empty((intermediate_size, hidden_size), dtype=dtype)

76 + torch.empty((intermediate_size, hidden_size))

77 )

78 self.norm = RMSNorm(intermediate_size, 1e-05)

76 - self.norm.weight = torch.nn.Parameter(

77 - torch.ones(intermediate_size, dtype=dtype)

78 - )

79 + self.norm.weight = torch.nn.Parameter(torch.ones(intermediate_size))

81 torch.nn.init.normal_(self.gate_proj, std=0.02)

@@ -112,9 +113,8 @@ class TestFusedAddRMSNorm(torch.nn.Module):

113 return norm_output, residual_output

114

115 def example_inputs(self, batch_size=8, hidden_size=16, seq_len=16):

115 - dtype = torch.float16 if TEST_FP8 else torch.float32

116 - hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

117 - residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

116 + hidden_states = torch.randn((batch_size * seq_len, hidden_size))

117 + residual = torch.randn((batch_size * seq_len, hidden_size))

118 return (hidden_states, residual)

119

120 def ops_in_model(self, do_fusion):

@@ -145,10 +145,9 @@ class TestRotaryEmbedding(torch.nn.Module):

145 return q_rotated, k_rotated

146

147 def example_inputs(self, num_tokens=32, head_dim=64):

148 - dtype = torch.float16

148 positions = torch.arange(num_tokens, dtype=torch.long)

150 - q = torch.randn(num_tokens, head_dim, dtype=dtype)

151 - k = torch.randn(num_tokens, head_dim, dtype=dtype)

149 + q = torch.randn(num_tokens, head_dim)

150 + k = torch.randn(num_tokens, head_dim)

151 return (positions, q, k)

152

153 def ops_in_model(self, do_fusion):

@@ -166,7 +165,7 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):

165 self.hidden_size = head_dim * num_heads

166

167 self.qkv_proj = torch.nn.Linear(

169 - self.hidden_size, self.hidden_size * 3, bias=False, dtype=torch.float16

168 + self.hidden_size, self.hidden_size * 3, bias=False

169 )

170

171 self.rotary_emb = get_rope(

@@ -190,10 +189,9 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):

189 return qkv_updated

190

191 def example_inputs(self, num_tokens=32, head_dim=64, num_heads=4):

193 - dtype = torch.float16

192 hidden_size = head_dim * num_heads

193 positions = torch.arange(num_tokens, dtype=torch.long)

196 - hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)

194 + hidden_states = torch.randn(num_tokens, hidden_size)

195 return (positions, hidden_states)

196

197 def ops_in_model(self, do_fusion):

@@ -211,48 +209,58 @@ MODELS = [

209 ]

210

211

212 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

213 @pytest.mark.parametrize("model_class", MODELS)

214 @pytest.mark.parametrize("do_fusion", [True, False])

215 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")

217 -def test_fix_functionalization(model_class: torch.nn.Module, do_fusion: bool):

216 +def test_fix_functionalization(

217 + model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype

218 +):

219 torch.set_default_device("cuda")

219 -

220 - vllm_config = VllmConfig()

221 - vllm_config.compilation_config = CompilationConfig(

222 - pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True)

220 + torch.set_default_dtype(dtype)

221 +

222 + vllm_config = VllmConfig(

223 + model_config=ModelConfig(dtype=dtype),

224 + compilation_config=CompilationConfig(

225 + custom_ops=["all"],

226 + pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True),

227 + ),

228 )

224 - noop_pass = NoOpEliminationPass(vllm_config)

225 - fusion_pass = RMSNormQuantFusionPass(vllm_config)

226 - cleanup_pass = PostCleanupPass(vllm_config)

227 - act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)

228 -

229 - passes = (

230 - [noop_pass, fusion_pass, act_quant_fusion_pass, cleanup_pass]

231 - if do_fusion

232 - else [noop_pass, cleanup_pass]

233 - )

234 - func_pass = FixFunctionalizationPass(vllm_config)

229

236 - backend_func = TestBackend(*passes, func_pass)

237 - backend_no_func = TestBackend(*passes)

230 + with set_current_vllm_config(vllm_config):

231 + assert RMSNorm.enabled()

232 + noop_pass = NoOpEliminationPass(vllm_config)

233 + fusion_pass = RMSNormQuantFusionPass(vllm_config)

234 + cleanup_pass = PostCleanupPass(vllm_config)

235 + act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)

236 +

237 + passes = (

238 + [noop_pass, fusion_pass, act_quant_fusion_pass, cleanup_pass]

239 + if do_fusion

240 + else [noop_pass, cleanup_pass]

241 + )

242 + func_pass = FixFunctionalizationPass(vllm_config)

243

239 - model = model_class()

240 - torch.compile(model, backend=backend_func)(*model.example_inputs())

241 - torch.compile(model, backend=backend_no_func)(*model.example_inputs())

244 + backend_func = TestBackend(*passes, func_pass)

245 + backend_no_func = TestBackend(*passes)

246

243 - # check if the functionalization pass is applied

244 - for op in model.ops_in_model(do_fusion):

245 - find_auto_fn(backend_no_func.graph_post_pass.nodes, op)

246 - assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None

247 + model = model_class()

248 + torch.compile(model, backend=backend_func)(*model.example_inputs())

249 + torch.compile(model, backend=backend_no_func)(*model.example_inputs())

250

248 - # make sure the ops were all de-functionalized

249 - found = dict()

250 - for node in backend_func.graph_post_pass.nodes:

251 + # check if the functionalization pass is applied

252 for op in model.ops_in_model(do_fusion):

252 - if is_func(node, op):

253 - found[op] = True

254 - for op in model.ops_not_in_model():

255 - if is_func(node, op):

256 - found[op] = True

257 - assert all(found[op] for op in model.ops_in_model(do_fusion))

258 - assert all(not found.get(op) for op in model.ops_not_in_model())

253 + find_auto_fn(backend_no_func.graph_post_pass.nodes, op)

254 + assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None

255 +

256 + # make sure the ops were all de-functionalized

257 + found = dict()

258 + for node in backend_func.graph_post_pass.nodes:

259 + for op in model.ops_in_model(do_fusion):

260 + if is_func(node, op):

261 + found[op] = True

262 + for op in model.ops_not_in_model():

263 + if is_func(node, op):

264 + found[op] = True

265 + assert all(found[op] for op in model.ops_in_model(do_fusion))

266 + assert all(not found.get(op) for op in model.ops_not_in_model())

tests/compile/test_fusion.py

+86 -35

@@ -5,15 +5,18 @@ import pytest

5 import torch

7 import vllm.plugins

8 -from vllm.compilation.fusion import (

9 - FUSED_OPS,

10 - QUANT_OPS,

11 - FusedRMSQuantKey,

12 - RMSNormQuantFusionPass,

13 -)

8 +from vllm.compilation.fusion import FUSED_OPS, FusedRMSQuantKey, RMSNormQuantFusionPass

9 +from vllm.compilation.fx_utils import find_op_nodes

10 +from vllm.compilation.matcher_utils import QUANT_OPS

11 from vllm.compilation.noop_elimination import NoOpEliminationPass

12 from vllm.compilation.post_cleanup import PostCleanupPass

16 -from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig

13 +from vllm.config import (

14 + CompilationConfig,

15 + CompilationMode,

16 + ModelConfig,

17 + PassConfig,

18 + VllmConfig,

19 +)

20 from vllm.model_executor.layers.layernorm import RMSNorm

21 from vllm.model_executor.layers.quantization.utils.quant_utils import (

22 GroupShape,

@@ -32,6 +35,9 @@ from .backend import TestBackend

36 FP8_DTYPE = current_platform.fp8_dtype()

38 +RMS_OP = torch.ops._C.rms_norm.default

39 +RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default

40 +

42 class TestModel(torch.nn.Module):

43 def __init__(

@@ -45,18 +51,18 @@ class TestModel(torch.nn.Module):

51 ):

52 super().__init__(*args, **kwargs)

53 self.cuda_force_torch = cuda_force_torch

48 - self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]

49 - self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]

54 + self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]

55 + self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

56 group_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN

57 quant_scale = ScaleDesc(torch.float32, static, group_shape)

52 - self.key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)

58 + self.quant_key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)

59 if static:

54 - self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]

60 + self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

61 else:

56 - self.scale = [None for _ in range(2)]

62 + self.scale = [None for _ in range(3)]

63 self.w = [

64 torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()

59 - for _ in range(2)

65 + for _ in range(3)

66 ]

68 with override_cutlass_fp8_supported(not cuda_force_torch):

@@ -65,8 +71,12 @@ class TestModel(torch.nn.Module):

71 act_quant_group_shape=group_shape,

72 )

74 + self.enable_rms_norm_custom_op = self.norm[0].enabled()

75 + self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()

76 +

77 def forward(self, x):

69 - resid = torch.sqrt(x)

78 + # avoid having graph input be an arg to a pattern directly

79 + x = resid = torch.relu(x)

80 y = self.norm[0](x)

82 x2 = self.fp8_linear.apply(

@@ -78,24 +88,44 @@ class TestModel(torch.nn.Module):

88 x3 = self.fp8_linear.apply(

89 y2, self.w[1], self.wscale[1], input_scale=self.scale[1]

90 )

91 +

92 y3, resid = self.norm[2](x3, resid) # use resid here

82 - return y3

84 - def ops_in_model_before(self):

85 - return [QUANT_OPS[self.key]]

94 + x4 = self.fp8_linear.apply(

95 + y3, self.w[2], self.wscale[2], input_scale=self.scale[2]

96 + )

97 +

98 + y4, resid = self.norm[3](x4, resid) # use resid here

99 + return y4

100

101 def ops_in_model_after(self):

102 return [

89 - FUSED_OPS[FusedRMSQuantKey(self.key, False)],

90 - FUSED_OPS[FusedRMSQuantKey(self.key, True)],

103 + FUSED_OPS[FusedRMSQuantKey(self.quant_key, True)],

104 + FUSED_OPS[FusedRMSQuantKey(self.quant_key, False)],

105 ]

106

107 + def ops_in_model_before(self):

108 + return (

109 + [QUANT_OPS[self.quant_key]]

110 + if self.enable_quant_fp8_custom_op

111 + else [torch.ops.aten.reciprocal]

112 + )

113 +

114 + def ops_in_model_before_partial(self):

115 + return (

116 + [RMS_OP, RMS_ADD_OP]

117 + if self.enable_rms_norm_custom_op

118 + else [torch.ops.aten.rsqrt]

119 + )

120 +

121

122 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

123 @pytest.mark.parametrize("hidden_size", [64])

124 @pytest.mark.parametrize("num_tokens", [257])

125 @pytest.mark.parametrize("eps", [1e-5, 1e-6])

126 @pytest.mark.parametrize("static", [True, False])

127 +@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])

128 +@pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])

129 # cuda_force_torch used to test torch code path on platforms that

130 # cutlass_fp8_supported() == True.

131 @pytest.mark.parametrize(

@@ -105,19 +135,32 @@ class TestModel(torch.nn.Module):

135 not current_platform.is_cuda_alike(), reason="Only test on CUDA and ROCm"

136 )

137 def test_fusion_rmsnorm_quant(

108 - dtype, hidden_size, num_tokens, eps, static, cuda_force_torch

138 + dtype,

139 + hidden_size,

140 + num_tokens,

141 + eps,

142 + static,

143 + enable_rms_norm_custom_op,

144 + enable_quant_fp8_custom_op,

145 + cuda_force_torch,

146 ):

147 torch.set_default_device("cuda")

148 torch.set_default_dtype(dtype)

149 torch.manual_seed(1)

150 maybe_create_device_identity() # needed for certain non-cutlass fp8 paths

151

152 + custom_ops = []

153 + if enable_rms_norm_custom_op:

154 + custom_ops.append("+rms_norm")

155 + if enable_quant_fp8_custom_op:

156 + custom_ops.append("+quant_fp8")

157 vllm_config = VllmConfig(

158 + model_config=ModelConfig(dtype=dtype),

159 compilation_config=CompilationConfig(

160 mode=CompilationMode.VLLM_COMPILE,

118 - custom_ops=["+rms_norm", "+quant_fp8"],

161 + custom_ops=custom_ops,

162 pass_config=PassConfig(enable_fusion=True, enable_noop=True),

120 - )

163 + ),

164 )

165 with vllm.config.set_current_vllm_config(vllm_config):

166 # Reshape pass is needed for the fusion pass to work

@@ -126,31 +169,39 @@ def test_fusion_rmsnorm_quant(

169 cleanup_pass = PostCleanupPass(vllm_config)

170

171 backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)

172 + backend2 = TestBackend(noop_pass, cleanup_pass)

173 model = TestModel(hidden_size, eps, static, cuda_force_torch)

174

175 # First dimension dynamic

176 x = torch.rand(num_tokens, hidden_size)

177 torch._dynamo.mark_dynamic(x, 0)

178

135 - result = model(x)

179 + model_fused = torch.compile(model, backend=backend)

180 + result_fused = model_fused(x)

181

137 - model2 = torch.compile(model, backend=backend)

138 - result2 = model2(x)

182 + model_unfused = torch.compile(model, backend=backend2)

183 + result_unfused = model_unfused(x)

184

140 - # Higher tol for dynamic, even higher for bfloat16

141 - if static:

142 - ATOL, RTOL = (1e-3, 1e-3)

143 - elif dtype == torch.float16:

185 + if dtype == torch.float16:

186 ATOL, RTOL = (2e-3, 2e-3)

187 else:

188 ATOL, RTOL = (1e-2, 1e-2)

189

148 - torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)

190 + torch.testing.assert_close(result_fused, result_unfused, atol=ATOL, rtol=RTOL)

191

150 - assert fusion_pass.matched_count == 2

151 -

152 - # In pre-nodes, fp8 quant should be there and fused kernels should not

192 + assert fusion_pass.matched_count == 3

193 backend.check_before_ops(model.ops_in_model_before())

154 -

155 - # In post-nodes, fused kernels should be there and fp8 quant should not

194 + backend.check_before_ops(

195 + model.ops_in_model_before_partial(), fully_replaced=False

196 + )

197 backend.check_after_ops(model.ops_in_model_after())

198 +

199 + # If RMSNorm custom op is disabled (native/torch impl used),

200 + # there's a risk that the fused add doesn't get included in the

201 + # replacement and only the rms part gets fused with quant.

202 + # Hence, we check only 2 add nodes are left (final fused rmsnorm add).

203 + if not enable_rms_norm_custom_op:

204 + n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))

205 + # 7 = 1 (RMS) + 3x2 (3xRMS_ADD, 2 each)

206 + assert n_add_nodes(backend.graph_pre_pass) == 7

207 + assert n_add_nodes(backend.graph_post_pass) == 2

tests/compile/test_fusion_all_reduce.py

+156 -80

@@ -6,6 +6,7 @@ import pytest

6 import torch

8 import vllm.envs as envs

9 +from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

10 from vllm.compilation.collective_fusion import AllReduceFusionPass

11 from vllm.compilation.fix_functionalization import FixFunctionalizationPass

12 from vllm.compilation.noop_elimination import NoOpEliminationPass

@@ -17,6 +18,7 @@ from vllm.config import (

18 ModelConfig,

19 PassConfig,

20 VllmConfig,

21 + set_current_vllm_config,

22 )

23 from vllm.distributed import tensor_model_parallel_all_reduce

24 from vllm.distributed.parallel_state import (

@@ -25,8 +27,8 @@ from vllm.distributed.parallel_state import (

27 )

28 from vllm.model_executor.layers.layernorm import RMSNorm

29 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

30 + Fp8LinearOp,

31 GroupShape,

29 - QuantFP8,

32 )

33 from vllm.platforms import current_platform

34 from vllm.utils import update_environment_variables

@@ -40,33 +42,30 @@ class TestAllReduceRMSNormModel(torch.nn.Module):

42 super().__init__()

43 self.hidden_size = hidden_size

44 self.eps = eps

43 - self.norm = RMSNorm(hidden_size, eps)

45 + self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]

46 + self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]

45 - def forward(self, hidden_states, residual):

46 - view = hidden_states.reshape(-1, self.hidden_size)

47 - all_reduce = tensor_model_parallel_all_reduce(view)

48 - norm = self.norm(all_reduce)

49 - return norm

48 + def forward(self, x):

49 + # avoid having graph input be an arg to a pattern directly

50 + z = torch.relu(x)

51 + x = resid = tensor_model_parallel_all_reduce(z)

52 + y = self.norm[0](x)

51 - def ops_in_model_before(self):

52 - return [torch.ops.vllm.all_reduce.default]

54 + z2 = torch.mm(y, self.w[0])

55 + x2 = tensor_model_parallel_all_reduce(z2)

54 - def ops_in_model_after(self):

55 - return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]

57 + y2, resid = self.norm[1](x2, resid)

59 + z3 = torch.mm(y2, self.w[1])

60 + x3 = tensor_model_parallel_all_reduce(z3)

58 -class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):

59 - def __init__(self, hidden_size=16, token_num=16, eps=1e-6):

60 - super().__init__()

61 - self.hidden_size = hidden_size

62 - self.eps = eps

63 - self.norm = RMSNorm(hidden_size, eps)

62 + y3, resid = self.norm[2](x3, resid)

63 +

64 + z4 = torch.mm(y3, self.w[2])

65 + x4 = tensor_model_parallel_all_reduce(z4)

65 - def forward(self, hidden_states, residual):

66 - view = hidden_states.reshape(-1, self.hidden_size)

67 - all_reduce = tensor_model_parallel_all_reduce(view)

68 - norm, _ = self.norm(all_reduce, residual)

69 - return norm

67 + y4, resid = self.norm[3](x4, resid)

68 + return y4

70 def ops_in_model_before(self):

71 return [torch.ops.vllm.all_reduce.default]

@@ -75,24 +74,53 @@ class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):

74 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]

78 -class TestAllReduceFusedAddRMSNormStaticQuantFP8Model(torch.nn.Module):

77 +class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):

78 def __init__(self, hidden_size=16, token_num=16, eps=1e-6):

79 super().__init__()

80 self.hidden_size = hidden_size

81 self.eps = eps

83 - self.norm = RMSNorm(hidden_size, eps)

84 - self.quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)

85 - self.scale = torch.rand(1, dtype=torch.float32)

86 - self.output = torch.empty((token_num, hidden_size), dtype=torch.float32)

87 -

88 - def forward(self, hidden_states, residual):

89 - view = hidden_states.reshape(-1, self.hidden_size)

90 - all_reduce = tensor_model_parallel_all_reduce(view)

91 - norm_output, residual_output = self.norm(all_reduce, residual)

92 - torch.ops._C.static_scaled_fp8_quant(

93 - self.output, norm_output.contiguous(), self.scale

82 + self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]

83 + self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

84 + self.w = [

85 + torch.rand(hidden_size, hidden_size)

86 + .to(dtype=current_platform.fp8_dtype())

87 + .t()

88 + for _ in range(3)

89 + ]

90 +

91 + self.fp8_linear = Fp8LinearOp(

92 + act_quant_static=True,

93 + act_quant_group_shape=GroupShape.PER_TENSOR,

94 + )

95 +

96 + self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

97 +

98 + def forward(self, hidden_states):

99 + # avoid having graph input be an arg to a pattern directly

100 + z = torch.relu(hidden_states)

101 + x = resid = tensor_model_parallel_all_reduce(z)

102 + y = self.norm[0](x)

103 +

104 + z2 = self.fp8_linear.apply(

105 + y, self.w[0], self.wscale[0], input_scale=self.scale[0]

106 + )

107 +

108 + x2 = tensor_model_parallel_all_reduce(z2)

109 + y2, resid = self.norm[1](x2, resid)

110 +

111 + z3 = self.fp8_linear.apply(

112 + y2, self.w[1], self.wscale[1], input_scale=self.scale[1]

113 )

95 - return self.output, residual_output

114 +

115 + x3 = tensor_model_parallel_all_reduce(z3)

116 + y3, resid = self.norm[2](x3, resid) # use resid here

117 +

118 + z4 = self.fp8_linear.apply(

119 + y3, self.w[2], self.wscale[2], input_scale=self.scale[2]

120 + )

121 + x4 = tensor_model_parallel_all_reduce(z4)

122 + y4, resid = self.norm[3](x4, resid) # use resid here

123 + return y4

124

125 def ops_in_model_after(self):

126 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]

@@ -100,7 +128,9 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP8Model(torch.nn.Module):

128 def ops_in_model_before(self):

129 return [

130 torch.ops.vllm.all_reduce.default,

103 - torch.ops._C.static_scaled_fp8_quant.default,

131 + torch.ops._C.static_scaled_fp8_quant.default

132 + if self.fp8_linear.quant_fp8.enabled()

133 + else torch.ops.aten.reciprocal.default,

134 ]

135

136

@@ -109,25 +139,48 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):

139 super().__init__()

140 self.hidden_size = hidden_size

141 self.eps = eps

112 - self.norm = RMSNorm(hidden_size, eps)

113 - self.scale = torch.rand(1, dtype=torch.float32)

114 - self.output = torch.empty((token_num, hidden_size), dtype=torch.float32)

115 -

116 - round_up = lambda x, y: (x + y - 1) // y * y

117 - rounded_m = round_up(token_num, 128)

118 - scale_n = hidden_size // 16

119 - rounded_n = round_up(scale_n, 4)

120 - self.output_scale = torch.empty((rounded_m, rounded_n // 4), dtype=torch.int32)

121 -

122 - def forward(self, hidden_states, residual):

123 - view = hidden_states.reshape(-1, self.hidden_size)

124 - all_reduce = tensor_model_parallel_all_reduce(view)

125 - norm_output, residual_output = self.norm(all_reduce, residual)

126 - norm_output = norm_output.reshape(-1, norm_output.shape[-1])

127 - torch.ops._C.scaled_fp4_quant(

128 - self.output, norm_output, self.output_scale, self.scale

142 + self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]

143 +

144 + self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]

145 + self.agscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

146 + wgscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]

147 + self.alpha = [1 / (w * a) for w, a in zip(wgscale, self.agscale)]

148 +

149 + wq_gen, wscale_gen = zip(

150 + *(scaled_fp4_quant(w, wg) for w, wg in zip(self.w, wgscale))

151 + )

152 + self.wq, self.wscale = list(wq_gen), list(wscale_gen)

153 + print(f"{self.wq=}, {self.wscale=}")

154 +

155 + def forward(self, hidden_states):

156 + # avoid having graph input be an arg to a pattern directly

157 + z = torch.relu(hidden_states)

158 + x = resid = tensor_model_parallel_all_reduce(z)

159 + y = self.norm[0](x)

160 +

161 + yq, y_scale = scaled_fp4_quant(y, self.agscale[0])

162 + z2 = cutlass_scaled_fp4_mm(

163 + yq, self.wq[0], y_scale, self.wscale[0], self.alpha[0], out_dtype=y.dtype

164 + )

165 +

166 + x2 = tensor_model_parallel_all_reduce(z2)

167 + y2, resid = self.norm[1](x2, resid)

168 +

169 + yq2, y_scale2 = scaled_fp4_quant(y2, self.agscale[1])

170 + z3 = cutlass_scaled_fp4_mm(

171 + yq2, self.wq[1], y_scale2, self.wscale[1], self.alpha[1], out_dtype=y2.dtype

172 )

130 - return self.output, residual_output, self.output_scale

173 +

174 + x3 = tensor_model_parallel_all_reduce(z3)

175 + y3, resid = self.norm[2](x3, resid) # use resid here

176 +

177 + yq3, y_scale3 = scaled_fp4_quant(y3, self.agscale[2])

178 + z4 = cutlass_scaled_fp4_mm(

179 + yq3, self.wq[2], y_scale3, self.wscale[2], self.alpha[2], out_dtype=y3.dtype

180 + )

181 + x4 = tensor_model_parallel_all_reduce(z4)

182 + y4, resid = self.norm[3](x4, resid) # use resid here

183 + return y4

184

185 def ops_in_model_after(self):

186 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]

@@ -141,19 +194,19 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):

194

195 @multi_gpu_test(num_gpus=2)

196 @pytest.mark.parametrize(

144 - "test_model",

197 + "test_model, enable_quant_fp8_custom_op",

198 [

146 - TestAllReduceRMSNormModel,

147 - TestAllReduceFusedAddRMSNormModel,

148 - TestAllReduceFusedAddRMSNormStaticQuantFP8Model,

149 - # TODO: Enable with torch==2.8.0

150 - # TestAllReduceFusedAddRMSNormStaticQuantFP4Model,

199 + (TestAllReduceRMSNormModel, False),

200 + (TestAllReduceRMSNormStaticQuantFP8Model, True),

201 + (TestAllReduceRMSNormStaticQuantFP8Model, False),

202 + (TestAllReduceFusedAddRMSNormStaticQuantFP4Model, False),

203 ],

204 )

205 @pytest.mark.parametrize("batch_size", [8])

206 @pytest.mark.parametrize("seq_len", [8])

155 -@pytest.mark.parametrize("hidden_size", [16])

207 +@pytest.mark.parametrize("hidden_size", [64])

208 @pytest.mark.parametrize("dtype", [torch.bfloat16])

209 +@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])

210 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")

211 @pytest.mark.skipif(

212 not find_spec("flashinfer")

@@ -167,6 +220,8 @@ def test_all_reduce_fusion_pass_replace(

220 seq_len: int,

221 hidden_size: int,

222 dtype: torch.dtype,

223 + enable_rms_norm_custom_op,

224 + enable_quant_fp8_custom_op,

225 ):

226 num_processes = 2

227 if (

@@ -181,7 +236,16 @@ def test_all_reduce_fusion_pass_replace(

236 def run_torch_spawn(fn, nprocs):

237 torch.multiprocessing.spawn(

238 fn,

184 - args=(num_processes, test_model, batch_size, seq_len, hidden_size, dtype),

239 + args=(

240 + num_processes,

241 + test_model,

242 + batch_size,

243 + seq_len,

244 + hidden_size,

245 + dtype,

246 + enable_rms_norm_custom_op,

247 + enable_quant_fp8_custom_op,

248 + ),

249 nprocs=nprocs,

250 )

251

@@ -196,6 +260,8 @@ def all_reduce_fusion_pass_on_test_model(

260 seq_len: int,

261 hidden_size: int,

262 dtype: torch.dtype,

263 + enable_rms_norm_custom_op,

264 + enable_quant_fp8_custom_op,

265 ):

266 current_platform.seed_everything(0)

267

@@ -217,15 +283,22 @@ def all_reduce_fusion_pass_on_test_model(

283 init_distributed_environment()

284 initialize_model_parallel(tensor_model_parallel_size=world_size)

285

286 + custom_ops = []

287 + if enable_rms_norm_custom_op:

288 + custom_ops.append("+rms_norm")

289 + if enable_quant_fp8_custom_op:

290 + custom_ops.append("+quant_fp8")

291 +

292 vllm_config = VllmConfig(

293 compilation_config=CompilationConfig(

222 - mode=CompilationMode.VLLM_COMPILE, custom_ops=["+rms_norm", "+quant_fp8"]

294 + mode=CompilationMode.VLLM_COMPILE, custom_ops=custom_ops

295 )

296 )

297 vllm_config.compilation_config.pass_config = PassConfig(

298 enable_fi_allreduce_fusion=True, enable_noop=True

299 )

300 vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

301 + vllm_config.parallel_config.rank = local_rank # Setup rank for debug path

302

303 # this is a fake model name to construct the model config

304 # in the vllm_config, it's not really used.

@@ -233,24 +306,27 @@ def all_reduce_fusion_pass_on_test_model(

306 vllm_config.model_config = ModelConfig(

307 model=model_name, trust_remote_code=True, dtype=dtype, seed=42

308 )

309 + with set_current_vllm_config(vllm_config):

310 + all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)

311 + noop_pass = NoOpEliminationPass(vllm_config)

312 + func_pass = FixFunctionalizationPass(vllm_config)

313 + cleanup_pass = PostCleanupPass(vllm_config)

314 +

315 + backend = TestBackend(

316 + noop_pass, all_reduce_fusion_pass, func_pass, cleanup_pass

317 + )

318

237 - all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)

238 - noop_pass = NoOpEliminationPass(vllm_config)

239 - func_pass = FixFunctionalizationPass(vllm_config)

240 - cleanup_pass = PostCleanupPass(vllm_config)

241 -

242 - backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass, cleanup_pass)

243 -

244 - token_num = batch_size * seq_len

245 - model = test_model_cls(hidden_size, token_num)

319 + token_num = batch_size * seq_len

320 + model = test_model_cls(hidden_size, token_num)

321

247 - hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)

248 - residual = torch.randn((token_num, hidden_size), requires_grad=False)

322 + hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)

323

250 - compiled_model = torch.compile(model, backend=backend)

251 - compiled_model(hidden_states, residual)

324 + compiled_model = torch.compile(model, backend=backend)

325 + compiled_model(hidden_states)

326

253 - assert all_reduce_fusion_pass.matched_count == 1

254 - backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)

255 - backend.check_after_ops(model.ops_in_model_after())

256 - del all_reduce_fusion_pass

327 + assert all_reduce_fusion_pass.matched_count == 4, (

328 + f"{all_reduce_fusion_pass.matched_count=}"

329 + )

330 + backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)

331 + backend.check_after_ops(model.ops_in_model_after())

332 + del all_reduce_fusion_pass

tests/compile/test_fusion_attn.py

+58 -46

@@ -6,14 +6,15 @@ import pytest

6 import torch._dynamo

8 from tests.compile.backend import LazyInitPass, TestBackend

9 +from tests.utils import flat_product

10 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata

11 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

12 from vllm.attention import Attention, AttentionMetadata

13 from vllm.attention.backends.registry import _Backend

14 from vllm.attention.selector import global_force_attn_backend_context_manager

14 -from vllm.compilation.fusion import QUANT_OPS

15 from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass

16 from vllm.compilation.fx_utils import find_op_nodes

17 +from vllm.compilation.matcher_utils import QUANT_OPS

18 from vllm.compilation.noop_elimination import NoOpEliminationPass

19 from vllm.compilation.post_cleanup import PostCleanupPass

20 from vllm.config import (

@@ -28,21 +29,18 @@ from vllm.config import (

29 )

30 from vllm.forward_context import get_forward_context, set_forward_context

31 from vllm.model_executor.layers.quantization.utils.quant_utils import (

32 + QuantKey,

33 kFp8StaticTensorSym,

34 kNvfp4Quant,

35 )

36 from vllm.model_executor.layers.quantization.utils.w8a8_utils import Fp8LinearOp

37 from vllm.platforms import current_platform

36 -from vllm.utils import is_torch_equal_or_newer

38 +from vllm.utils.flashinfer import has_flashinfer

39 from vllm.v1.kv_cache_interface import AttentionSpec

41 FP8_DTYPE = current_platform.fp8_dtype()

42 FP4_DTYPE = torch.uint8

42 -# globals needed for string-import custom Dynamo backend field

43 -backend: TestBackend | None = None

44 -backend_unfused: TestBackend | None = None

45 -

45 class AttentionQuantPatternModel(torch.nn.Module):

46 """Base model for AttentionQuantPattern fusion."""

@@ -104,6 +102,7 @@ class AttentionQuantPatternModel(torch.nn.Module):

102 num_blocks = batch_size * max_blocks

103 backend = self.attn.backend

104

105 + # TODO(luka) use get_kv_cache_stride_order

106 # Create dummy KV cache for the selected backend

107 if backend == _Backend.ROCM_ATTN:

108 # k/v as 1st dimention

@@ -241,26 +240,40 @@ class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):

240 )

241

242

243 +MODELS_FP8: list[tuple[str, type]] = []

244 +MODELS_FP4: list[tuple[str, type]] = []

245 +HEADS: list[tuple[int, int]] = []

246 +SPLIT_ATTENTION: list[bool] = []

247 +BACKENDS_FP8: list[_Backend] = []

248 +BACKENDS_FP4: list[_Backend] = []

249 +

250 if current_platform.is_cuda():

245 - MODELS = [

251 + HEADS = [(64, 8), (40, 8)]

252 + MODELS_FP8 = [

253 (

254 "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",

255 TestAttentionFp8StaticQuantPatternModel,

249 - ),

256 + )

257 + ]

258 + MODELS_FP4 = [

259 (

260 "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",

261 TestAttentionNvfp4QuantPatternModel,

253 - ),

262 + )

263 ]

255 - HEADS = [(64, 8), (40, 8)]

264 + BACKENDS_FP8 = [_Backend.TRITON_ATTN, _Backend.FLASHINFER]

265 + BACKENDS_FP4 = [_Backend.FLASHINFER]

266 +

267 elif current_platform.is_rocm():

257 - MODELS = [

268 + HEADS = [(32, 8), (40, 8)]

269 + MODELS_FP8 = [

270 ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)

271 ]

260 - HEADS = [(32, 8), (40, 8)]

261 -else:

262 - MODELS = []

263 - HEADS = []

272 + BACKENDS = [

273 + _Backend.ROCM_AITER_UNIFIED_ATTN,

274 + _Backend.ROCM_ATTN,

275 + _Backend.TRITON_ATTN,

276 + ]

277

278

279 @pytest.mark.parametrize("num_qo_heads, num_kv_heads", HEADS)

@@ -269,46 +282,36 @@ else:

282 "batch_size", [7, 256, 533] if current_platform.is_cuda() else [8]

283 )

284 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])

272 -@pytest.mark.parametrize("model_name, model_class", MODELS)

285 @pytest.mark.parametrize(

274 - "backend",

275 - [_Backend.FLASHINFER]

276 - if current_platform.is_cuda()

277 - else [_Backend.ROCM_AITER_UNIFIED_ATTN, _Backend.ROCM_ATTN, _Backend.TRITON_ATTN],

278 -)

279 -# TODO(boyuan): test inductor graph partition on rocm

280 -@pytest.mark.parametrize(

281 - "use_inductor_graph_partition",

282 - [False] if current_platform.is_rocm() else [False, True],

286 + "backend, model_name, model_class, custom_ops",

287 + # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8

288 + list(flat_product(BACKENDS_FP8, MODELS_FP8, ["+quant_fp8", "-quant_fp8"]))

289 + # quant_fp4 only has the custom impl

290 + + list(flat_product(BACKENDS_FP4, MODELS_FP4, [""])),

291 )

292 @pytest.mark.skipif(

293 not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"

294 )

295 @pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")

288 -@pytest.mark.skipif(

289 - current_platform.is_cuda() and not current_platform.is_device_capability((10, 0)),

290 - reason="On CUDA only test on SM100(Blackwell)",

291 -)

292 -@pytest.mark.skipif(

293 - not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"

294 -)

296 def test_attention_quant_pattern(

297 num_qo_heads: int,

298 num_kv_heads: int,

299 head_size: int,

300 batch_size: int,

301 dtype: torch.dtype,

302 + custom_ops: str,

303 model_name: str,

304 model_class: type[AttentionQuantPatternModel],

305 backend: _Backend,

304 - use_inductor_graph_partition: bool,

306 dist_init,

306 - caplog_vllm,

307 ):

308 """Test AttentionStaticQuantPattern fusion pass"""

309 + if backend == _Backend.FLASHINFER and (

310 + not current_platform.is_device_capability((10, 0)) or not has_flashinfer()

311 + ):

312 + pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")

313

310 - if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

311 - pytest.skip("inductor graph partition is only available in PyTorch 2.9+")

314 + custom_ops_list = custom_ops.split(",") if custom_ops else []

315

316 device = torch.device("cuda:0")

317 torch.manual_seed(42)

@@ -322,8 +325,7 @@ def test_attention_quant_pattern(

325 scheduler_config=SchedulerConfig(max_num_seqs=1024),

326 compilation_config=CompilationConfig(

327 mode=CompilationMode.VLLM_COMPILE,

325 - custom_ops=["+quant_fp8"],

326 - use_inductor_graph_partition=use_inductor_graph_partition,

328 + custom_ops=custom_ops_list,

329 ),

330 cache_config=CacheConfig(cache_dtype="fp8"),

331 )

@@ -358,8 +360,9 @@ def test_attention_quant_pattern(

360 forward_ctx = get_forward_context()

361 forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)

362

361 - # Run model directly without compilation and fusion

362 - result_unfused = model_unfused(q, k, v)

363 + # Run model directly without fusion

364 + # Still compile so query QuantFP8 has closer numerics

365 + result_unfused = torch.compile(model_unfused, fullgraph=True)(q, k, v)

366

367 # Run model with attn fusion enabled

368 vllm_config.compilation_config.pass_config = PassConfig(

@@ -414,16 +417,25 @@ def test_attention_quant_pattern(

417 )

418

419 # Check attn fusion support

417 - quant_key = model_class.quant_key

420 + quant_key: QuantKey = model_class.quant_key

421 attn_fusion_supported = [

422 layer.impl.fused_output_quant_supported(quant_key)

423 for key, layer in vllm_config.compilation_config.static_forward_context.items()

424 ]

422 - if any(attn_fusion_supported):

423 - # Check quantization ops in the graph before and after fusion

424 - # Note: fully_replaced=False because query quant ops remain in graph.

425 - # Only output quant ops are fused into attention.

426 - test_backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=False)

425 + assert sum(attn_fusion_supported) == len(attn_fusion_supported), (

426 + "All layers should support attention fusion"

427 + )

428 +

429 + # Check quantization ops in the graph before and after fusion

430 + quant_op = (

431 + torch.ops.aten.reciprocal

432 + if "-quant_fp8" in custom_ops_list

433 + else QUANT_OPS[quant_key]

434 + )

435 +

436 + # Note: for fp8, fully_replaced=False because query quant ops remain in graph.

437 + # Only output quant ops are fused into attention.

438 + test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Quant)

439

440 # access the underlying `AttnFusionPass` on the `LazyInitPass`

441 assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)

tests/compile/test_fusions_e2e.py new

+305

@@ -0,0 +1,305 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from __future__ import annotations

5 +

6 +import itertools

7 +import logging

8 +from collections.abc import Iterable

9 +from typing import Any, NamedTuple

10 +

11 +import pytest

12 +import regex as re

13 +

14 +from tests.v1.attention.utils import _Backend

15 +from vllm import LLM, SamplingParams

16 +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig

17 +from vllm.platforms import current_platform

18 +from vllm.utils import is_torch_equal_or_newer

19 +from vllm.utils.flashinfer import has_flashinfer

20 +

21 +from ..utils import flat_product, multi_gpu_test

22 +

23 +

24 +class ModelBackendTestCase(NamedTuple):

25 + model_name: str

26 + model_kwargs: dict[str, Any]

27 + backend: _Backend

28 + attention_fusions: int

29 + allreduce_fusions: int | None = None

30 +

31 +

32 +MODELS_FP8: list[ModelBackendTestCase] = []

33 +MODELS_FP4: list[ModelBackendTestCase] = []

34 +MODELS: list[ModelBackendTestCase] = [] # tp-only

35 +

36 +if current_platform.is_cuda():

37 + MODELS_FP8 = [

38 + ModelBackendTestCase(

39 + # Use smaller model for L40s in CI

40 + model_name="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",

41 + model_kwargs=dict(max_model_len=1024),

42 + backend=_Backend.TRITON_ATTN,

43 + attention_fusions=32,

44 + allreduce_fusions=65,

45 + ),

46 + ModelBackendTestCase(

47 + model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP8",

48 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),

49 + backend=_Backend.FLASHINFER,

50 + attention_fusions=48,

51 + allreduce_fusions=96,

52 + ),

53 + ]

54 +

55 + MODELS_FP4 = [

56 + ModelBackendTestCase(

57 + model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-FP4",

58 + model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),

59 + backend=_Backend.FLASHINFER,

60 + attention_fusions=48,

61 + allreduce_fusions=96,

62 + ),

63 + ]

64 +

65 + # TP only

66 + MODELS = [

67 + ModelBackendTestCase(

68 + model_name="meta-llama/Llama-3.1-8B-Instruct",

69 + model_kwargs=dict(max_model_len=1024),

70 + backend=_Backend.TRITON_ATTN,

71 + attention_fusions=0,

72 + allreduce_fusions=65,

73 + ),

74 + ]

75 +

76 +elif current_platform.is_rocm():

77 + MODELS_FP8 = [

78 + ModelBackendTestCase(

79 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",

80 + model_kwargs=dict(max_model_len=1024),

81 + backend=_Backend.TRITON_ATTN,

82 + attention_fusions=32,

83 + ),

84 + ModelBackendTestCase(

85 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",

86 + model_kwargs=dict(max_model_len=1024),

87 + backend=_Backend.ROCM_ATTN,

88 + attention_fusions=32,

89 + ),

90 + ModelBackendTestCase(

91 + model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",

92 + model_kwargs=dict(max_model_len=1024),

93 + backend=_Backend.ROCM_AITER_UNIFIED_ATTN,

94 + attention_fusions=32,

95 + ),

96 + ]

97 +

98 +# TODO(luka) test both in nightly

99 +CUSTOM_OPS_FP8 = ["-quant_fp8"] # , "+quant_fp8"]

100 +

101 +

102 +@pytest.mark.parametrize(

103 + "model_name, model_kwargs, backend, "

104 + "attention_fusions, allreduce_fusions, custom_ops",

105 + # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8

106 + list(flat_product(MODELS_FP8, CUSTOM_OPS_FP8))

107 + # quant_fp4 only has the custom impl

108 + + list(flat_product(MODELS_FP4, [""])),

109 +)

110 +@pytest.mark.parametrize("inductor_graph_partition", [True, False])

111 +def test_attn_quant(

112 + model_name: str,

113 + model_kwargs: dict[str, Any],

114 + backend: _Backend,

115 + attention_fusions: int,

116 + allreduce_fusions: int,

117 + custom_ops: str,

118 + inductor_graph_partition: bool,

119 + caplog_mp_spawn,

120 + monkeypatch,

121 +):

122 + if backend == _Backend.FLASHINFER and (

123 + not current_platform.is_device_capability((10, 0)) or not has_flashinfer()

124 + ):

125 + pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")

126 + if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

127 + pytest.skip("Inductor graph partition requires torch>=2.9")

128 +

129 + custom_ops_list = custom_ops.split(",") if custom_ops else []

130 +

131 + if inductor_graph_partition:

132 + mode = CUDAGraphMode.FULL_AND_PIECEWISE

133 + splitting_ops: list[str] | None = None

134 + else:

135 + mode = CUDAGraphMode.FULL_DECODE_ONLY

136 + splitting_ops = []

137 +

138 + # Disable, compile cache to make sure custom passes run.

139 + # Otherwise, we can't verify fusion happened through the logs.

140 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

141 +

142 + # To capture subprocess logs, we need to know whether spawn or fork is used.

143 + # Force spawn as it is more general.

144 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

145 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)

146 +

147 + compilation_config = CompilationConfig(

148 + # Testing properties

149 + custom_ops=custom_ops_list,

150 + use_inductor_graph_partition=inductor_graph_partition,

151 + cudagraph_mode=mode,

152 + splitting_ops=splitting_ops,

153 + # Common

154 + level=CompilationMode.VLLM_COMPILE,

155 + pass_config=PassConfig(enable_attn_fusion=True, enable_noop=True),

156 + # Inductor caches custom passes by default as well via uuid

157 + inductor_compile_config={"force_disable_caches": True},

158 + )

159 +

160 + with caplog_mp_spawn(logging.DEBUG) as log_holder:

161 + run_model(compilation_config, model_name, **model_kwargs)

162 +

163 + matches = re.findall(

164 + r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",

165 + log_holder.text,

166 + )

167 + assert len(matches) == 1, log_holder.text

168 + assert int(matches[0]) == attention_fusions

169 +

170 +

171 +# TODO(luka) test both in nightly

172 +CUSTOM_OPS_RMS_NORM = ["-rms_norm"] # , "+rms_norm"]

173 +

174 +

175 +def custom_ops_product(*custom_ops_lists: list[str]) -> Iterable[str]:

176 + for op_list in itertools.product(*custom_ops_lists):

177 + yield ",".join(op_list)

178 +

179 +

180 +@multi_gpu_test(num_gpus=2)

181 +@pytest.mark.parametrize(

182 + "model_name, model_kwargs, backend, "

183 + "attention_fusions, allreduce_fusions, custom_ops",

184 + # Toggle RMSNorm and QuantFP8 for FP8 models

185 + list(

186 + flat_product(

187 + MODELS_FP8, custom_ops_product(CUSTOM_OPS_FP8, CUSTOM_OPS_RMS_NORM)

188 + )

189 + )

190 + # Toggle RMSNorm for FP4 models and unquant models

191 + + list(flat_product(MODELS_FP4 + MODELS, CUSTOM_OPS_RMS_NORM)),

192 +)

193 +@pytest.mark.parametrize("inductor_graph_partition", [True, False])

194 +@pytest.mark.skipif(

195 + not current_platform.is_cuda()

196 + or not has_flashinfer()

197 + or not current_platform.has_device_capability(90),

198 + reason="allreduce+rmsnorm fusion requires flashinfer",

199 +)

200 +def test_tp2_attn_quant_allreduce_rmsnorm(

201 + model_name: str,

202 + model_kwargs: dict,

203 + backend: _Backend,

204 + attention_fusions: int,

205 + allreduce_fusions: int,

206 + custom_ops: str,

207 + inductor_graph_partition: bool,

208 + caplog_mp_spawn,

209 + monkeypatch,

210 +):

211 + if inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):

212 + pytest.skip("Inductor graph partition requires torch>=2.9")

213 +

214 + custom_ops_list = custom_ops.split(",") if custom_ops else []

215 +

216 + if inductor_graph_partition:

217 + mode = CUDAGraphMode.FULL_AND_PIECEWISE

218 + splitting_ops: list[str] | None = None

219 + else:

220 + mode = CUDAGraphMode.FULL_DECODE_ONLY

221 + splitting_ops = []

222 +

223 + # Disable, compile cache to make sure custom passes run.

224 + # Otherwise, we can't verify fusion happened through the logs.

225 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

226 +

227 + # To capture subprocess logs, we need to know whether spawn or fork is used.

228 + # Force spawn as it is more general.

229 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

230 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend.name)

231 +

232 + compilation_config = CompilationConfig(

233 + # Testing properties

234 + use_inductor_graph_partition=inductor_graph_partition,

235 + cudagraph_mode=mode,

236 + custom_ops=custom_ops_list,

237 + splitting_ops=splitting_ops,

238 + # Common

239 + level=CompilationMode.VLLM_COMPILE,

240 + pass_config=PassConfig(

241 + enable_attn_fusion=True,

242 + enable_noop=True,

243 + enable_fi_allreduce_fusion=True,

244 + ),

245 + # Inductor caches custom passes by default as well via uuid

246 + inductor_compile_config={"force_disable_caches": True},

247 + )

248 +

249 + with caplog_mp_spawn(logging.DEBUG) as log_holder:

250 + run_model(

251 + compilation_config, model_name, tensor_parallel_size=2, **model_kwargs

252 + )

253 + matches = re.findall(

254 + r"fusion_attn.py:\d+] Fused quant onto (\d+) attention nodes",

255 + log_holder.text,

256 + )

257 + assert len(matches) == 2, log_holder.text

258 +

259 + assert int(matches[0]) == attention_fusions

260 + assert int(matches[1]) == attention_fusions

261 +

262 + matches = re.findall(

263 + r"collective_fusion.py:\d+] Replaced (\d+) patterns",

264 + log_holder.text,

265 + )

266 + assert len(matches) == 2, log_holder.text

267 +

268 + assert int(matches[0]) == allreduce_fusions

269 + assert int(matches[1]) == allreduce_fusions

270 +

271 +

272 +def run_model(compile_config: int | CompilationConfig, model: str, **model_kwargs):

273 + compilation_config = (

274 + compile_config

275 + if isinstance(compile_config, CompilationConfig)

276 + else CompilationConfig(level=compile_config)

277 + )

278 +

279 + prompts = [

280 + "Hello, my name is",

281 + "The president of the United States is",

282 + "The capital of France is",

283 + "The future of AI is",

284 + ]

285 + sampling_params = SamplingParams(temperature=0)

286 + # Allow override from model_kwargs

287 + model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}

288 + model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}

289 +

290 + # No cudagraphs by default

291 + if compilation_config.cudagraph_mode is None:

292 + compilation_config.cudagraph_mode = CUDAGraphMode.NONE

293 +

294 + llm = LLM(

295 + model=model,

296 + compilation_config=compilation_config,

297 + **model_kwargs,

298 + )

299 + outputs = llm.generate(prompts, sampling_params)

300 +

301 + # Print the outputs.

302 + for output in outputs:

303 + prompt = output.prompt

304 + generated_text = output.outputs[0].text

305 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

tests/compile/test_pass_manager.py

+3 -2

@@ -7,7 +7,7 @@ import torch

8 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass

9 from vllm.compilation.pass_manager import PostGradPassManager

10 -from vllm.config import VllmConfig

10 +from vllm.config import ModelConfig, VllmConfig

13 # dummy custom pass that doesn't inherit

@@ -42,7 +42,8 @@ class ProperPass(InductorPass):

42 ],

43 )

44 def test_pass_manager_uuid(callable):

45 - config = VllmConfig()

45 + # Some passes need dtype to be set

46 + config = VllmConfig(model_config=ModelConfig(dtype=torch.bfloat16))

48 pass_manager = PostGradPassManager()

49 pass_manager.configure(config)

tests/compile/test_sequence_parallelism.py

+60 -54

@@ -18,6 +18,8 @@ from vllm.config import (

18 ModelConfig,

19 PassConfig,

20 VllmConfig,

21 + get_current_vllm_config,

22 + set_current_vllm_config,

23 )

24 from vllm.distributed import tensor_model_parallel_all_reduce

25 from vllm.distributed.parallel_state import (

@@ -42,9 +44,7 @@ prompts = [

46 class TestModel(torch.nn.Module):

45 - def __init__(

46 - self, hidden_size=16, intermediate_size=32, vllm_config: VllmConfig = None

47 - ):

47 + def __init__(self, hidden_size=16, intermediate_size=32):

48 super().__init__()

49 self.hidden_size = hidden_size

50 self.intermediate_size = intermediate_size

@@ -95,13 +95,11 @@ class TestModel(torch.nn.Module):

97 class TestQuantModel(torch.nn.Module):

98 - def __init__(

99 - self, hidden_size=16, intermediate_size=32, vllm_config: VllmConfig = None

100 - ):

98 + def __init__(self, hidden_size=16, intermediate_size=32):

99 super().__init__()

100 self.hidden_size = hidden_size

101 self.intermediate_size = intermediate_size

104 - self.vllm_config = vllm_config

102 + self.vllm_config = get_current_vllm_config()

103 self.gate_proj = torch.nn.Parameter(

104 torch.empty((intermediate_size, hidden_size)), requires_grad=False

105 )

@@ -266,76 +264,84 @@ def sequence_parallelism_pass_on_test_model(

264 initialize_model_parallel(tensor_model_parallel_size=world_size)

265

266 # configure vllm config for SequenceParallelismPass

269 - vllm_config = VllmConfig()

270 - vllm_config.compilation_config = CompilationConfig(

267 + compilation_config = CompilationConfig(

268 pass_config=PassConfig(

269 enable_sequence_parallelism=True,

270 enable_fusion=enable_fusion,

271 enable_noop=True,

272 )

273 ) # NoOp needed for fusion

277 - vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))

274 + device_config = DeviceConfig(device=torch.device("cuda"))

275

276 # this is a fake model name to construct the model config

277 # in the vllm_config, it's not really used.

278 model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8"

282 - vllm_config.model_config = ModelConfig(

279 + model_config = ModelConfig(

280 model=model_name, trust_remote_code=True, dtype=dtype, seed=42

281 )

282

286 - noop_pass = NoOpEliminationPass(vllm_config)

287 - sequence_parallelism_pass = SequenceParallelismPass(vllm_config)

288 - assert (

289 - sequence_parallelism_pass.compilation_config.splitting_ops

290 - == vllm_config.compilation_config.splitting_ops

283 + vllm_config = VllmConfig(

284 + model_config=model_config,

285 + device_config=device_config,

286 + compilation_config=compilation_config,

287 )

292 - assert (

293 - sequence_parallelism_pass.compilation_config.use_inductor_graph_partition

294 - == vllm_config.compilation_config.use_inductor_graph_partition

295 - )

296 - func_pass = FixFunctionalizationPass(vllm_config)

297 - cleanup_pass = PostCleanupPass(vllm_config)

298 -

299 - passes_for_backend: list[VllmInductorPass] = [noop_pass, sequence_parallelism_pass]

288

301 - if enable_fusion:

302 - fusion_pass = RMSNormQuantFusionPass(vllm_config)

303 - passes_for_backend.append(fusion_pass)

289 + with set_current_vllm_config(vllm_config):

290 + noop_pass = NoOpEliminationPass(vllm_config)

291 + sequence_parallelism_pass = SequenceParallelismPass(vllm_config)

292 + func_pass = FixFunctionalizationPass(vllm_config)

293 + cleanup_pass = PostCleanupPass(vllm_config)

294 + assert (

295 + sequence_parallelism_pass.compilation_config.splitting_ops

296 + == vllm_config.compilation_config.splitting_ops

297 + )

298 + assert (

299 + sequence_parallelism_pass.compilation_config.use_inductor_graph_partition

300 + == vllm_config.compilation_config.use_inductor_graph_partition

301 + )

302 + passes_for_backend: list[VllmInductorPass] = [

303 + noop_pass,

304 + sequence_parallelism_pass,

305 + ]

306

305 - passes_for_backend.append(cleanup_pass)

307 + if enable_fusion:

308 + fusion_pass = RMSNormQuantFusionPass(vllm_config)

309 + passes_for_backend.append(fusion_pass)

310

307 - backend_no_func = TestBackend(*passes_for_backend)

308 - backend_func = TestBackend(*passes_for_backend, func_pass)

311 + passes_for_backend.append(cleanup_pass)

312

310 - model = test_model_cls(hidden_size, hidden_size * 2, vllm_config=vllm_config)

313 + backend_no_func = TestBackend(*passes_for_backend)

314 + backend_func = TestBackend(*passes_for_backend, func_pass)

315

312 - hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

313 - residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

316 + model = test_model_cls(hidden_size, hidden_size * 2)

317

315 - compiled_model_no_func = torch.compile(model, backend=backend_no_func)

316 - compiled_model_no_func(hidden_states, residual)

317 - compiled_model_func = torch.compile(model, backend=backend_func)

318 - compiled_model_func(hidden_states, residual)

318 + hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

319 + residual = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype)

320

320 - assert sequence_parallelism_pass.matched_count == 1

321 + compiled_model_no_func = torch.compile(model, backend=backend_no_func)

322 + compiled_model_no_func(hidden_states, residual)

323 + compiled_model_func = torch.compile(model, backend=backend_func)

324 + compiled_model_func(hidden_states, residual)

325

322 - # In pre-nodes, all reduce should be there,

323 - # reduce scatter and all gather should not

324 - backend_no_func.check_before_ops(model.ops_in_model_before())

326 + assert sequence_parallelism_pass.matched_count == 1

327

326 - # In post-nodes, reduce scatter and all gather should be there,

327 - # all reduce should not

328 - backend_no_func.check_after_ops(model.ops_in_model_after())

328 + # In pre-nodes, all reduce should be there,

329 + # reduce scatter and all gather should not

330 + backend_no_func.check_before_ops(model.ops_in_model_before())

331

330 - # check if the functionalization pass is applied

331 - for op in model.ops_in_model():

332 - find_auto_fn(backend_no_func.graph_post_pass.nodes, op)

333 - assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None

332 + # In post-nodes, reduce scatter and all gather should be there,

333 + # all reduce should not

334 + backend_no_func.check_after_ops(model.ops_in_model_after())

335

335 - # make sure the ops were all de-functionalized

336 - found = dict()

337 - for node in backend_func.graph_post_pass.nodes:

336 + # check if the functionalization pass is applied

337 for op in model.ops_in_model():

339 - if is_func(node, op):

340 - found[op] = True

341 - assert all(found[op] for op in model.ops_in_model())

338 + find_auto_fn(backend_no_func.graph_post_pass.nodes, op)

339 + assert find_auto_fn_maybe(backend_func.graph_post_pass.nodes, op) is None

340 +

341 + # make sure the ops were all de-functionalized

342 + found = dict()

343 + for node in backend_func.graph_post_pass.nodes:

344 + for op in model.ops_in_model():

345 + if is_func(node, op):

346 + found[op] = True

347 + assert all(found[op] for op in model.ops_in_model())

tests/conftest.py

+101 -3

@@ -1,10 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# ruff: noqa

3 +import contextlib

4 +import pathlib

5 +from copy import deepcopy

7 from tblib import pickling_support

9 +# ruff: noqa

10 +

11 # Install support for pickling exceptions so that we can nicely propagate

12 # failures from tests running in a subprocess.

13 # This should be run before any custom exception subclasses are defined.

@@ -40,7 +43,7 @@ from transformers import (

43 from transformers.models.auto.auto_factory import _BaseAutoModelClass

45 from tests.models.utils import TokensTextLogprobs, TokensTextLogprobsPromptLogprobs

43 -from vllm import LLM, SamplingParams

46 +from vllm import LLM, SamplingParams, envs

47 from vllm.assets.audio import AudioAsset

48 from vllm.assets.image import ImageAsset

49 from vllm.assets.video import VideoAsset

@@ -1070,6 +1073,101 @@ def caplog_vllm(temporary_enable_log_propagate, caplog):

1073 yield caplog

1074

1075

1076 +@pytest.fixture()

1077 +def caplog_mp_fork():

1078 + """

1079 + This fixture enables capturing logs from a forked MP subprocess.

1080 + It should be used in conjunction with caplog_vllm.

1081 +

1082 + By default, subprocess logs do not go through the parent process.

1083 + We instead create a queue listener in the parent process which

1084 + forwards logs to the logger's other handlers, and add a QueueHandler

1085 + to the root logger. Forked subprocesses will inherit the root logger

1086 + and pass their messages to the queue, which the listener will forward

1087 + to the root logger, which can be captured by caplog.

1088 +

1089 + Note that this workaround only works for fork; with spawn, the subprocess

1090 + reinitializes logging and does not automatically inherit the queue.

1091 + We'd have to manually pass the queue to the subprocess at the spawn point.

1092 + See caplog_mp_spawn below.

1093 + """

1094 +

1095 + @contextlib.contextmanager

1096 + def ctx():

1097 + import logging.handlers

1098 + import multiprocessing as mp

1099 +

1100 + logger_queue: mp.Queue[logging.LogRecord] = mp.Queue()

1101 + logger = logging.getLogger()

1102 + handlers = logger.handlers

1103 +

1104 + # The listener works on a background thread, not inherited by the child.

1105 + queue_listener = logging.handlers.QueueListener(logger_queue, *handlers)

1106 + queue_listener.start()

1107 +

1108 + # Add queue handler after creating the listener to avoid cycle

1109 + logger.addHandler(logging.handlers.QueueHandler(logger_queue))

1110 + yield

1111 + queue_listener.stop()

1112 +

1113 + return ctx

1114 +

1115 +

1116 +class LogHolder:

1117 + def __init__(self):

1118 + self.text = None

1119 +

1120 +

1121 +@pytest.fixture()

1122 +def caplog_mp_spawn(tmp_path, monkeypatch):

1123 + """

1124 + This fixture enables capturing logs from a forked MP subprocess.

1125 + It does not require caplog_vllm (but it only contains logs from the child).

1126 +

1127 + By default, subprocess logs do not go through the parent process.

1128 + We instead add a FileHandler to the config so the spawned child process

1129 + writes its logs to a temp file.

1130 + In the parent, we read the file and return the contents.

1131 +

1132 + Note: this method could be extended to fork by either reconfiguring logging

1133 + in the parent or using a SocketHandler:

1134 + https://docs.python.org/3/howto/logging-cookbook.html#sending-and-receiving-logging-events-across-a-network # noqa: E501

1135 + """

1136 +

1137 + @contextlib.contextmanager

1138 + def ctx(level: int | str):

1139 + from vllm.logger import DEFAULT_LOGGING_CONFIG

1140 +

1141 + config_path = tmp_path / "vllm_logging_config.json"

1142 + log_path = tmp_path / "vllm.log"

1143 + log_holder = LogHolder()

1144 +

1145 + config = deepcopy(DEFAULT_LOGGING_CONFIG)

1146 + if envs.VLLM_LOGGING_CONFIG_PATH:

1147 + path = pathlib.Path(envs.VLLM_LOGGING_CONFIG_PATH)

1148 + assert path.exists()

1149 + config = json.loads(path.read_text())

1150 +

1151 + config["loggers"]["vllm"]["handlers"] += ["vllm_file"]

1152 + config["handlers"]["vllm_file"] = {

1153 + "class": "logging.FileHandler",

1154 + "formatter": "vllm",

1155 + "level": level,

1156 + "filename": log_path.as_posix(),

1157 + }

1158 +

1159 + config_path.write_text(json.dumps(config))

1160 +

1161 + with monkeypatch.context() as monkeypatch_ctx:

1162 + monkeypatch_ctx.setenv("VLLM_LOGGING_CONFIG_PATH", config_path.as_posix())

1163 + monkeypatch_ctx.setenv("VLLM_CONFIGURE_LOGGING", "1")

1164 + yield log_holder

1165 +

1166 + log_holder.text = log_path.read_text()

1167 +

1168 + return ctx

1169 +

1170 +

1171 @pytest.fixture(scope="session")

1172 def num_gpus_available():

1173 """Get number of GPUs without initializing the CUDA context

tests/kernels/quant_utils.py

+1 -1

@@ -103,7 +103,7 @@ def ref_dynamic_per_tensor_fp8_quant(

103 .clamp(fp8_traits_min, fp8_traits_max)

104 .to(FP8_DTYPE)

105 )

106 - return ref_out, ref_scale.view((1,))

106 + return ref_out, ref_scale.view((1, 1))

107

108

109 def native_w8a8_block_matmul(

tests/test_logger.py

+46

@@ -501,3 +501,49 @@ def test_streaming_complete_logs_full_text_content():

501 assert call_args[1] == "test-streaming-full-text"

502 assert call_args[2] == " (streaming complete)"

503 assert call_args[5] == "streaming_complete"

504 +

505 +

506 +# Add vllm prefix to make sure logs go through the vllm logger

507 +test_logger = init_logger("vllm.test_logger")

508 +

509 +

510 +def mp_function(**kwargs):

511 + # This function runs in a subprocess

512 +

513 + test_logger.warning("This is a subprocess: %s", kwargs.get("a"))

514 + test_logger.error("This is a subprocess error.")

515 + test_logger.debug("This is a subprocess debug message: %s.", kwargs.get("b"))

516 +

517 +

518 +def test_caplog_mp_fork(caplog_vllm, caplog_mp_fork):

519 + with caplog_vllm.at_level(logging.DEBUG), caplog_mp_fork():

520 + import multiprocessing

521 +

522 + ctx = multiprocessing.get_context("fork")

523 + p = ctx.Process(

524 + target=mp_function,

525 + name=f"SubProcess{1}",

526 + kwargs={"a": "AAAA", "b": "BBBBB"},

527 + )

528 + p.start()

529 + p.join()

530 +

531 + assert "AAAA" in caplog_vllm.text

532 + assert "BBBBB" in caplog_vllm.text

533 +

534 +

535 +def test_caplog_mp_spawn(caplog_mp_spawn):

536 + with caplog_mp_spawn(logging.DEBUG) as log_holder:

537 + import multiprocessing

538 +

539 + ctx = multiprocessing.get_context("spawn")

540 + p = ctx.Process(

541 + target=mp_function,

542 + name=f"SubProcess{1}",

543 + kwargs={"a": "AAAA", "b": "BBBBB"},

544 + )

545 + p.start()

546 + p.join()

547 +

548 + assert "AAAA" in log_holder.text

549 + assert "BBBBB" in log_holder.text

tests/utils.py

+22 -1

@@ -6,6 +6,7 @@ import contextlib

6 import copy

7 import functools

8 import importlib

9 +import itertools

10 import json

11 import os

12 import random

@@ -15,7 +16,7 @@ import sys

16 import tempfile

17 import time

18 import warnings

18 -from collections.abc import Callable

19 +from collections.abc import Callable, Iterable

20 from contextlib import ExitStack, contextmanager, suppress

21 from multiprocessing import Process

22 from pathlib import Path

@@ -1261,3 +1262,23 @@ def check_answers(

1262 frac_ok = numok / len(answer)

1263 print(f"Num OK: {numok}/{len(answer)} {frac_ok}")

1264 assert frac_ok >= accept_rate

1265 +

1266 +

1267 +def flat_product(*iterables: Iterable[Any]):

1268 + """

1269 + Flatten lists of tuples of the cartesian product.

1270 + Useful when we want to avoid nested tuples to allow

1271 + test params to be unpacked directly from the decorator.

1272 +

1273 + Example:

1274 + flat_product([(1, 2), (3, 4)], ["a", "b"]) ->

1275 + [

1276 + (1, 2, "a"),

1277 + (1, 2, "b"),

1278 + (3, 4, "a"),

1279 + (3, 4, "b"),

1280 + ]

1281 + """

1282 + for element in itertools.product(*iterables):

1283 + normalized = (e if isinstance(e, tuple) else (e,) for e in element)

1284 + yield tuple(itertools.chain(*normalized))

tests/utils_/test_utils.py

+23 -1

@@ -40,7 +40,7 @@ from vllm.utils import (

40 unique_filepath,

41 )

43 -from ..utils import create_new_process_for_each_test

43 +from ..utils import create_new_process_for_each_test, flat_product

46 def test_get_open_port(monkeypatch: pytest.MonkeyPatch):

@@ -771,3 +771,25 @@ def test_unique_filepath():

771 paths.add(path)

772 assert len(paths) == 10

773 assert len(list(Path(temp_dir).glob("*.txt"))) == 10

774 +

775 +

776 +def test_flat_product():

777 + # Check regular itertools.product behavior

778 + result1 = list(flat_product([1, 2, 3], ["a", "b"]))

779 + assert result1 == [

780 + (1, "a"),

781 + (1, "b"),

782 + (2, "a"),

783 + (2, "b"),

784 + (3, "a"),

785 + (3, "b"),

786 + ]

787 +

788 + # check that the tuples get flattened

789 + result2 = list(flat_product([(1, 2), (3, 4)], ["a", "b"], [(5, 6)]))

790 + assert result2 == [

791 + (1, 2, "a", 5, 6),

792 + (1, 2, "b", 5, 6),

793 + (3, 4, "a", 5, 6),

794 + (3, 4, "b", 5, 6),

795 + ]

vllm/_custom_ops.py

+1 -1

@@ -1507,7 +1507,7 @@ def scaled_fp8_quant(

1507 output, input, scale, scale_ub

1508 )

1509 else:

1510 - scale = torch.empty(1, device=input.device, dtype=torch.float32)

1510 + scale = torch.empty((1, 1), device=input.device, dtype=torch.float32)

1511 torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)

1512 else:

1513 assert scale.numel() == 1, f"{scale.shape}"

vllm/compilation/collective_fusion.py

+70 -149

@@ -17,10 +17,14 @@ from vllm.distributed.parallel_state import (

17 get_tensor_model_parallel_world_size,

18 )

19 from vllm.logger import init_logger

20 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

21 + kFp8StaticTensorSym,

22 +)

23 from vllm.platforms import current_platform

24 from vllm.utils import direct_register_custom_op

26 from .inductor_pass import enable_fake_mode

27 +from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm

28 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

30 FP8_DTYPE = current_platform.fp8_dtype()

@@ -41,11 +45,8 @@ else:

46 logger = init_logger(__name__)

44 -ALLREDUCE_OP = torch.ops.vllm.all_reduce.default

45 -RMS_OP = torch.ops._C.rms_norm.default

46 -RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default

47 -STATIC_FP8_QUANT_OP = torch.ops._C.static_scaled_fp8_quant.default

48 -STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default

48 +if hasattr(torch.ops._C, "scaled_fp4_quant"):

49 + STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default

52 class BasePattern:

@@ -669,33 +670,24 @@ class AllReduceRMSNormPattern(BasePattern):

670 super().__init__(dtype, device)

671 self.epsilon = epsilon

672 self.allreduce_params = allreduce_params

673 + self.rmsnorm_matcher = MatcherRMSNorm(epsilon)

674

675 def get_inputs(self):

674 - input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)

675 - rms_result = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)

676 - weight = torch.empty([4], device=self.device, dtype=self.dtype)

676 + input, weight = self.rmsnorm_matcher.inputs()

677

678 - return [input, rms_result, weight]

678 + # input goes through allreduce first, always 16-bit

679 + return [input.to(self.dtype), weight]

680

681 def register(self, pm_pass: PatternMatcherPass):

681 - def pattern(

682 - input: torch.Tensor, rms_result: torch.Tensor, weight: torch.Tensor

683 - ):

682 + def pattern(input: torch.Tensor, weight: torch.Tensor):

683 allreduce_output = tensor_model_parallel_all_reduce(input)

685 - rms = auto_functionalized(

686 - RMS_OP,

687 - result=rms_result,

688 - input=allreduce_output,

689 - weight=weight,

690 - epsilon=self.epsilon,

691 - )

692 - # rms_result, allreduce_output

693 - return rms[1], allreduce_output

684 + rms = self.rmsnorm_matcher(allreduce_output, weight)

685

695 - def replacement(

696 - input: torch.Tensor, rms_result: torch.Tensor, weight: torch.Tensor

697 - ):

686 + return rms, allreduce_output

687 +

688 + def replacement(input: torch.Tensor, weight: torch.Tensor):

689 residual = torch.zeros_like(input)

690 + rms_result = torch.empty_like(input)

691 allreduce = auto_functionalized(

692 flashinfer_trtllm_fused_allreduce_norm,

693 allreduce_in=input,

@@ -733,29 +725,19 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):

725 super().__init__(dtype, device)

726 self.epsilon = epsilon

727 self.allreduce_params = allreduce_params

728 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)

729

730 def get_inputs(self):

738 - input = torch.empty([4, 4], device=self.device, dtype=self.dtype)

739 - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)

740 - weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)

741 - return [

742 - residual,

743 - input,

744 - weight,

745 - ]

731 + input, residual, weight = self.rmsnorm_matcher.inputs()

732 +

733 + # input goes through allreduce first, always 16-bit

734 + return [residual, input.to(self.dtype), weight]

735

736 def register(self, pm_pass: PatternMatcherPass):

737 def pattern(residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor):

738 allreduce_output = tensor_model_parallel_all_reduce(input)

750 - rms = auto_functionalized(

751 - RMS_ADD_OP,

752 - input=allreduce_output,

753 - residual=residual,

754 - weight=weight,

755 - epsilon=self.epsilon,

756 - )

757 - # input, residual

758 - return rms[1], rms[2]

739 + rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)

740 + return rms, residual

741

742 def replacement(

743 residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor

@@ -779,6 +761,18 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):

761 pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass

762 )

763

764 + # Same pattern, but only return the output and not residual

765 + # (helpful for end of graph where residual is not used again)

766 + first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0]

767 +

768 + pm.register_replacement(

769 + first_return_only(pattern),

770 + first_return_only(replacement),

771 + self.get_inputs(),

772 + pm.fwd_only,

773 + pm_pass,

774 + )

775 +

776

777 class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):

778 """

@@ -799,60 +793,37 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):

793 self.epsilon = epsilon

794 self.allreduce_params = allreduce_params

795 self.quant_dtype = torch.float8_e4m3fn

796 + self.rmsnorm_matcher = MatcherRMSNorm(epsilon)

797 + self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)

798

799 def register(self, pm_pass: PatternMatcherPass):

800 def get_inputs():

805 - input = torch.zeros([1, 8, 4], device=self.device, dtype=self.dtype)

806 - rmsnorm_result = torch.empty(

807 - [1, 8, 4], device=self.device, dtype=self.dtype

808 - )

809 - quant_result = torch.empty(

810 - [1, 8, 4], device=self.device, dtype=self.quant_dtype

811 - )

812 - weight = torch.empty([4], device=self.device, dtype=self.dtype)

813 - scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)

814 - return [input, rmsnorm_result, quant_result, weight, scale]

801 + input, weight = self.rmsnorm_matcher.inputs()

802 + _, scale = self.quant_matcher.inputs()

803 +

804 + # input goes through allreduce first, always 16-bit

805 + return [input.to(self.dtype), weight, scale]

806

807 def pattern(

808 input: torch.Tensor,

818 - rmsnorm_result: torch.Tensor,

819 - quant_result: torch.Tensor,

809 weight: torch.Tensor,

810 scale: torch.Tensor,

811 ):

812 all_reduce = tensor_model_parallel_all_reduce(input)

824 - rmsnorm_out_tuple = auto_functionalized(

825 - RMS_OP,

826 - result=rmsnorm_result,

827 - input=all_reduce,

828 - weight=weight,

829 - epsilon=self.epsilon,

830 - )

831 -

832 - quant_out_tuple = auto_functionalized(

833 - STATIC_FP8_QUANT_OP,

834 - result=quant_result,

835 - input=rmsnorm_out_tuple[1],

836 - scale=scale,

837 - )

838 -

839 - # quant_out, allreduce_output

840 - return quant_out_tuple[1], all_reduce

813 + rms = self.rmsnorm_matcher(all_reduce, weight)

814 + quant, _ = self.quant_matcher(rms, scale)

815 + return quant, all_reduce

816

842 - def replacement(

843 - input: torch.Tensor,

844 - result_rms: torch.Tensor,

845 - quant_result: torch.Tensor,

846 - weight: torch.Tensor,

847 - scale: torch.Tensor,

848 - ):

817 + def replacement(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):

818 residual = torch.zeros_like(input)

819 + result_rms = torch.empty_like(input)

820 + result_quant = torch.empty_like(input, dtype=self.quant_dtype)

821 allreduce = auto_functionalized(

822 flashinfer_trtllm_fused_allreduce_norm,

823 allreduce_in=input,

824 residual=residual,

825 norm_out=result_rms,

855 - quant_out=quant_result,

826 + quant_out=result_quant,

827 scale_out=None,

828 rms_gamma=weight,

829 rms_eps=self.epsilon,

@@ -892,64 +863,42 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):

863 self.allreduce_params = allreduce_params

864 self.quant_dtype = torch.float8_e4m3fn

865

866 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)

867 + self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)

868 +

869 def register(self, pm_pass: PatternMatcherPass):

870 def get_inputs():

897 - input = torch.empty([4, 4], device=self.device, dtype=self.dtype)

871 + input, residual, weight = self.rmsnorm_matcher.inputs()

872 + _, scale = self.quant_matcher.inputs()

873

899 - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype)

900 - weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)

901 - quant_result = torch.empty(

902 - [4, 4], device=self.device, dtype=self.quant_dtype

903 - )

904 - scale = torch.empty([1, 1], device=self.device, dtype=torch.float32)

905 -

906 - return [

907 - quant_result,

908 - residual,

909 - input,

910 - weight,

911 - scale,

912 - ]

874 + # input goes through allreduce first, always 16-bit

875 + return [residual, input.to(self.dtype), weight, scale]

876

877 def pattern(

915 - quant_result: torch.Tensor,

878 residual: torch.Tensor,

879 input: torch.Tensor,

880 weight: torch.Tensor,

881 scale: torch.Tensor,

882 ):

883 allreduce_output = tensor_model_parallel_all_reduce(input)

884 + rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual)

885 + quant, _ = self.quant_matcher(rms, scale)

886

923 - fused_add_rmsnorm_out_tuple = auto_functionalized(

924 - RMS_ADD_OP,

925 - input=allreduce_output,

926 - residual=residual,

927 - weight=weight,

928 - epsilon=self.epsilon,

929 - )

930 - quant_out_tuple = auto_functionalized(

931 - STATIC_FP8_QUANT_OP,

932 - result=quant_result,

933 - input=fused_add_rmsnorm_out_tuple[1],

934 - scale=scale,

935 - )

936 -

937 - # quant_out, allreduce_output

938 - return quant_out_tuple[1], fused_add_rmsnorm_out_tuple[2]

887 + return quant, res

888

889 def replacement(

941 - quant_result: torch.Tensor,

890 residual: torch.Tensor,

891 input: torch.Tensor,

892 weight: torch.Tensor,

893 scale: torch.Tensor,

894 ):

895 + result_quant = torch.empty_like(input, dtype=self.quant_dtype)

896 allreduce = auto_functionalized(

897 flashinfer_trtllm_fused_allreduce_norm,

898 allreduce_in=input,

899 residual=residual,

900 norm_out=None,

952 - quant_out=quant_result,

901 + quant_out=result_quant,

902 scale_out=None,

903 rms_gamma=weight,

904 rms_eps=self.epsilon,

@@ -986,14 +935,11 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):

935 super().__init__(dtype, device)

936 self.epsilon = epsilon

937 self.allreduce_params = allreduce_params

938 + self.rmsnorm_matcher = MatcherRMSNorm(epsilon)

939

940 def register(self, pm_pass: PatternMatcherPass):

941 def get_inputs():

942 input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype)

993 -

994 - rmsnorm_result = torch.empty(

995 - [1, 16, 16], device=self.device, dtype=self.dtype

996 - )

943 quant_result = torch.empty((16, 8), device=self.device, dtype=torch.uint8)

944 input_global_scale = torch.empty(

945 [1, 1], device=self.device, dtype=torch.float32

@@ -1001,36 +947,21 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):

947 weight = torch.empty([16], device=self.device, dtype=self.dtype)

948 output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32)

949

1004 - return [

1005 - input,

1006 - rmsnorm_result,

1007 - quant_result,

1008 - weight,

1009 - input_global_scale,

1010 - output_scale,

1011 - ]

950 + return [input, quant_result, weight, input_global_scale, output_scale]

951

952 def pattern(

953 input: torch.Tensor,

1015 - rmsnorm_result: torch.Tensor,

954 quant_result: torch.Tensor,

955 weight: torch.Tensor,

956 input_global_scale: torch.Tensor,

957 output_scale: torch.Tensor,

958 ):

959 all_reduce = tensor_model_parallel_all_reduce(input)

1022 - rmsnorm_out_tuple = auto_functionalized(

1023 - RMS_OP,

1024 - result=rmsnorm_result,

1025 - input=all_reduce,

1026 - weight=weight,

1027 - epsilon=self.epsilon,

1028 - )

1029 -

960 + rms = self.rmsnorm_matcher(all_reduce, weight)

961 quant_out_tuple = auto_functionalized(

962 STATIC_FP4_QUANT_OP,

963 output=quant_result,

1033 - input=rmsnorm_out_tuple[1],

964 + input=rms,

965 output_scale=output_scale,

966 input_scale=input_global_scale,

967 )

@@ -1040,13 +971,13 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):

971

972 def replacement(

973 input: torch.Tensor,

1043 - result_rms: torch.Tensor,

974 quant_result: torch.Tensor,

975 weight: torch.Tensor,

976 input_global_scale: torch.Tensor,

977 output_scale: torch.Tensor,

978 ):

979 residual = torch.zeros_like(input)

980 + result_rms = torch.empty_like(input)

981 allreduce = auto_functionalized(

982 flashinfer_trtllm_fused_allreduce_norm,

983 allreduce_in=input,

@@ -1090,6 +1021,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):

1021 super().__init__(dtype, device)

1022 self.epsilon = epsilon

1023 self.allreduce_params = allreduce_params

1024 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)

1025

1026 def register(self, pm_pass: PatternMatcherPass):

1027 def get_inputs():

@@ -1121,28 +1053,17 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):

1053 input_global_scale: torch.Tensor,

1054 ):

1055 allreduce_output = tensor_model_parallel_all_reduce(input)

1124 -

1125 - fused_add_rmsnorm_out_tuple = auto_functionalized(

1126 - RMS_ADD_OP,

1127 - input=allreduce_output,

1128 - residual=residual,

1129 - weight=weight,

1130 - epsilon=self.epsilon,

1131 - )

1056 + rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)

1057 quant_out_tuple = auto_functionalized(

1058 STATIC_FP4_QUANT_OP,

1059 output=quant_result,

1135 - input=fused_add_rmsnorm_out_tuple[1],

1060 + input=rms,

1061 output_scale=output_scale,

1062 input_scale=input_global_scale,

1063 )

1064

1065 # quant_out, allreduce_output, output_scale

1141 - return (

1142 - quant_out_tuple[1],

1143 - fused_add_rmsnorm_out_tuple[2],

1144 - quant_out_tuple[2],

1145 - )

1066 + return quant_out_tuple[1], residual, quant_out_tuple[2]

1067

1068 def replacement(

1069 quant_result: torch.Tensor,

vllm/compilation/fusion.py

+68 -129

@@ -9,7 +9,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized

9 from torch._inductor.pattern_matcher import PatternMatcherPass

10 from torch._ops import OpOverload

12 -from vllm.config import VllmConfig

12 +from vllm.config import VllmConfig, get_current_vllm_config

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.quantization.utils.quant_utils import (

15 GroupShape,

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

24 from vllm.platforms import current_platform

26 from .inductor_pass import enable_fake_mode

27 +from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm

28 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

30 logger = init_logger(__name__)

@@ -92,13 +93,19 @@ class RMSNormQuantPattern:

93 def __init__(self, epsilon: float, key: FusedRMSQuantKey):

94 self.epsilon = epsilon

95 self.quant_dtype = key.quant.dtype

95 -

96 - assert key.quant in QUANT_OPS, f"unsupported quantization scheme {key.quant}"

97 - self.QUANT_OP = QUANT_OPS[key.quant]

96 + config = get_current_vllm_config()

97 + self.model_dtype = config.model_config.dtype if config.model_config else None

99 assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"

100 self.FUSED_OP = FUSED_OPS[key]

101

102 + self.rmsnorm_matcher = (

103 + MatcherRMSNorm(epsilon)

104 + if not key.fused_add

105 + else MatcherFusedAddRMSNorm(epsilon)

106 + )

107 + self.quant_matcher = MatcherQuantFP8(key.quant)

108 +

109

110 class RMSNormStaticQuantPattern(RMSNormQuantPattern):

111 def __init__(self, epsilon: float, quant_dtype: torch.dtype, symmetric=True):

@@ -112,34 +119,18 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):

119

120 def register(self, pm_pass: PatternMatcherPass):

121 # Cannot use methods, as the self argument affects tracing

115 - def pattern(

116 - result: torch.Tensor,

117 - result_rms: torch.Tensor,

118 - input: torch.Tensor,

119 - weight: torch.Tensor,

120 - scale: torch.Tensor,

121 - ):

122 - at1 = auto_functionalized(

123 - RMS_OP,

124 - result=result_rms,

125 - input=input,

126 - weight=weight,

127 - epsilon=self.epsilon,

128 - )

129 - at2 = auto_functionalized(

130 - self.QUANT_OP, result=result, input=at1[1], scale=scale

131 - )

122 + def pattern(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):

123 + result_rms = self.rmsnorm_matcher(input, weight)

124 + return self.quant_matcher(result_rms, scale)[0]

125

133 - # result

134 - return at2[1]

126 + def replacement(input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor):

127 + # In case we're matching native rms-norm, conversions might be

128 + # optimized out. We convert here just to be safe.

129 + input = input.to(dtype=self.model_dtype)

130

136 - def replacement(

137 - result: torch.Tensor,

138 - result_rms: torch.Tensor,

139 - input: torch.Tensor,

140 - weight: torch.Tensor,

141 - scale: torch.Tensor,

142 - ):

131 + result = torch.empty(

132 + input.shape, device=input.device, dtype=self.quant_dtype

133 + )

134 at = auto_functionalized(

135 self.FUSED_OP,

136 result=result,

@@ -153,12 +144,11 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):

144 return at[1]

145

146 inputs = [

156 - torch.empty(5, 4, device="cuda", dtype=self.quant_dtype), # result

157 - empty_bf16(5, 4), # result_rms

158 - empty_bf16(5, 4), # input

159 - empty_bf16(1, 5), # weight

160 - empty_fp32(1, 1), # scale

147 + # input, weight

148 + *self.rmsnorm_matcher.inputs(),

149 + self.quant_matcher.inputs()[1], # scale

150 ]

151 + pattern(*inputs)

152

153 pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)

154

@@ -175,33 +165,27 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):

165

166 def register(self, pm_pass: PatternMatcherPass):

167 def pattern(

178 - result: torch.Tensor,

168 input: torch.Tensor,

180 - residual: torch.Tensor,

169 weight: torch.Tensor,

170 + residual: torch.Tensor,

171 scale: torch.Tensor,

172 ):

184 - at = auto_functionalized(

185 - RMS_ADD_OP,

186 - input=input,

187 - residual=residual,

188 - weight=weight,

189 - epsilon=self.epsilon,

190 - )

191 - at1 = auto_functionalized(

192 - self.QUANT_OP, result=result, input=at[1], scale=scale

193 - )

173 + result_rms, residual = self.rmsnorm_matcher(input, weight, residual)

174 + result, _ = self.quant_matcher(result_rms, scale)

175

195 - # result, residual

196 - return at1[1], at[2]

176 + return result, residual

177

178 def replacement(

199 - result: torch.Tensor,

179 input: torch.Tensor,

201 - residual: torch.Tensor,

180 weight: torch.Tensor,

181 + residual: torch.Tensor,

182 scale: torch.Tensor,

183 ):

184 + # In case we're matching native rms-norm, conversions might be

185 + # optimized out. We convert here just to be safe.

186 + input = input.to(dtype=self.model_dtype)

187 +

188 + result = torch.empty_like(input, dtype=self.quant_dtype)

189 at = auto_functionalized(

190 self.FUSED_OP,

191 result=result,

@@ -216,11 +200,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):

200 return at[1], at[2]

201

202 inputs = [

219 - torch.empty(5, 4, device="cuda", dtype=self.quant_dtype), # result

220 - empty_bf16(5, 4), # input

221 - empty_bf16(5, 4), # residual

222 - empty_bf16(1, 5), # weight

223 - empty_fp32(1, 1), # scale

203 + # input, weight, residual

204 + *self.rmsnorm_matcher.inputs(),

205 + self.quant_matcher.inputs()[1], # scale

206 ]

207

208 pm.register_replacement(

@@ -248,34 +230,18 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):

230 super().__init__(epsilon, key)

231

232 def register(self, pm_pass: PatternMatcherPass):

251 - def pattern(

252 - result: torch.Tensor,

253 - result_rms: torch.Tensor,

254 - input: torch.Tensor,

255 - weight: torch.Tensor,

256 - scale: torch.Tensor,

257 - ):

258 - at1 = auto_functionalized(

259 - RMS_OP,

260 - result=result_rms,

261 - input=input,

262 - weight=weight,

263 - epsilon=self.epsilon,

264 - )

265 - at2 = auto_functionalized(

266 - self.QUANT_OP, result=result, input=at1[1], scale=scale, scale_ub=None

267 - )

268 -

233 + def pattern(input: torch.Tensor, weight: torch.Tensor):

234 + result_rms = self.rmsnorm_matcher(input, weight)

235 # result, scale

270 - return at2[1], at2[2]

236 + return self.quant_matcher(result_rms)

237

272 - def replacement(

273 - result: torch.Tensor,

274 - result_rms: torch.Tensor,

275 - input: torch.Tensor,

276 - weight: torch.Tensor,

277 - scale: torch.Tensor,

278 - ):

238 + def replacement(input: torch.Tensor, weight: torch.Tensor):

239 + # In case we're matching native rms-norm, conversions might be

240 + # optimized out. We convert here just to be safe.

241 + input = input.to(dtype=self.model_dtype)

242 +

243 + result = torch.empty_like(input, dtype=self.quant_dtype)

244 + scale = self.quant_matcher.make_scale(input)

245 at = auto_functionalized(

246 self.FUSED_OP,

247 result=result,

@@ -290,18 +256,10 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):

256 # result, scale

257 return at[1], at[2]

258

293 - inputs = [

294 - torch.empty(5, 4, device="cuda", dtype=self.quant_dtype), # result

295 - empty_bf16(5, 4), # result_rms

296 - empty_bf16(5, 4), # input

297 - empty_bf16(1, 5), # weight

298 - empty_fp32(1, 1), # scale

299 - ]

300 -

259 pm.register_replacement(

260 pattern,

261 replacement,

304 - inputs,

262 + self.rmsnorm_matcher.inputs(),

263 pm.fwd_only,

264 pm_pass,

265 )

@@ -323,34 +281,21 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):

281 super().__init__(epsilon, key)

282

283 def register(self, pm_pass: PatternMatcherPass):

326 - def pattern(

327 - result: torch.Tensor,

328 - input: torch.Tensor,

329 - residual: torch.Tensor,

330 - weight: torch.Tensor,

331 - scale: torch.Tensor,

332 - ):

333 - at = auto_functionalized(

334 - RMS_ADD_OP,

335 - input=input,

336 - residual=residual,

337 - weight=weight,

338 - epsilon=self.epsilon,

339 - )

340 - at1 = auto_functionalized(

341 - self.QUANT_OP, result=result, input=at[1], scale=scale, scale_ub=None

342 - )

284 + def pattern(input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor):

285 + result_rms, residual = self.rmsnorm_matcher(input, weight, residual)

286 + result, scale = self.quant_matcher(result_rms)

287

344 - # result, residual, scale

345 - return at1[1], at[2], at1[2]

288 + return result, residual, scale

289

290 def replacement(

348 - result: torch.Tensor,

349 - input: torch.Tensor,

350 - residual: torch.Tensor,

351 - weight: torch.Tensor,

352 - scale: torch.Tensor,

291 + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor

292 ):

293 + # In case we're matching native rms-norm, conversions might be

294 + # optimized out. We convert here just to be safe.

295 + input = input.to(dtype=self.model_dtype)

296 +

297 + result = torch.empty_like(input, dtype=self.quant_dtype)

298 + scale = self.quant_matcher.make_scale(input)

299 at = auto_functionalized(

300 self.FUSED_OP,

301 result=result,

@@ -365,18 +310,10 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):

310 # result, residual, scale

311 return at[1], at[3], at[2]

312

368 - inputs = [

369 - torch.empty(5, 4, device="cuda", dtype=self.quant_dtype), # result

370 - empty_bf16(5, 4), # input

371 - empty_bf16(5, 4), # residual

372 - empty_bf16(1, 5), # weight

373 - empty_fp32(1, 1), # scale

374 - ]

375 -

313 pm.register_replacement(

314 pattern,

315 replacement,

379 - inputs,

316 + self.rmsnorm_matcher.inputs(),

317 pm.fwd_only,

318 pm_pass,

319 )

@@ -396,23 +333,25 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):

333 pass_name="rmsnorm_quant_fusion_pass"

334 )

335

336 + # Make sure fused add patterns are before simple rms norm,

337 + # as the latter is a subset of the former in torch ops

338 for epsilon in [1e-5, 1e-6]:

400 - # Fuse rms_norm + static fp8 quant

401 - RMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

402 -

339 # Fuse fused_add_rms_norm + static fp8 quant

340 FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(

341 self.patterns

342 )

343

408 - # Fuse rms_norm + dynamic per-token fp8 quant

409 - RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

344 + # Fuse rms_norm + static fp8 quant

345 + RMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

346

347 # Fuse fused_add_rms_norm + dynamic per-token fp8 quant

348 FusedAddRMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(

349 self.patterns

350 )

351

352 + # Fuse rms_norm + dynamic per-token fp8 quant

353 + RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

354 +

355 self.dump_patterns(config, self.patterns)

356

357 @VllmInductorPass.time_and_log

vllm/compilation/fusion_attn.py

+37 -18

@@ -2,9 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from abc import ABC, abstractmethod

5 +from collections.abc import Callable

7 import torch

8 import torch._inductor.pattern_matcher as pm

9 +from torch import fx

10 from torch._higher_order_ops.auto_functionalize import auto_functionalized

11 from torch._inductor.pattern_matcher import PatternMatcherPass

@@ -20,7 +22,9 @@ from vllm.platforms import current_platform

22 from vllm.utils import round_up

24 from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32

25 +from .fx_utils import is_func

26 from .inductor_pass import enable_fake_mode

27 +from .matcher_utils import MatcherQuantFP8

28 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

30 logger = init_logger(__name__)

@@ -66,9 +70,13 @@ class AttentionQuantPattern(ABC):

70 return torch.empty(*args, **kwargs)

72 @staticmethod

69 - def wrap_trace_fn(process_fx, trace_fn):

73 + def wrap_trace_fn(trace_fn, *process_fx_fns: Callable[[fx.GraphModule], None]):

74 def wrapped(*args, **kwargs):

71 - return process_fx(trace_fn(*args, **kwargs))

75 + gm = trace_fn(*args, **kwargs)

76 + for process_fx in process_fx_fns:

77 + process_fx(gm)

78 +

79 + return gm

81 return wrapped

@@ -77,7 +85,20 @@ class AttentionQuantPattern(ABC):

85 from torch._inductor.fx_passes.post_grad import view_to_reshape

87 view_to_reshape(gm)

80 - return gm

88 +

89 + @staticmethod

90 + def remove_noop_permutes(gm: torch.fx.GraphModule):

91 + for node in gm.graph.nodes:

92 + if not is_func(node, torch.ops.aten.permute.default):

93 + continue

94 +

95 + dims = node.args[1]

96 + if any(dim != i for i, dim in enumerate(dims)):

97 + continue

98 +

99 + # this is now an identity op, remove

100 + node.replace_all_uses_with(node.args[0])

101 + gm.graph.erase_node(node)

102

103 def register_if_supported(self, pm_pass: PatternMatcherPass):

104 if self.layer.impl.fused_output_quant_supported(self.quant_key):

@@ -108,6 +129,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

129 dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=symmetric

130 )

131 super().__init__(layer, quant_key, dtype)

132 + self.quant_matcher = MatcherQuantFP8(quant_key)

133

134 def _register(self, pm_pass: PatternMatcherPass):

135 def pattern(

@@ -115,7 +137,6 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

137 k: torch.Tensor,

138 v: torch.Tensor,

139 output_attn: torch.Tensor,

118 - output_quant: torch.Tensor,

140 scale: torch.Tensor,

141 ):

142 at1 = auto_functionalized(

@@ -131,17 +152,14 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

152 attn_out_view = RESHAPE_OP(

153 at1[1], [q.shape[0], self.num_heads * self.head_size]

154 )

134 - at2 = auto_functionalized(

135 - self.QUANT_OP, result=output_quant, input=attn_out_view, scale=scale

136 - )

137 - return at2[1]

155 +

156 + return self.quant_matcher(attn_out_view, scale)[0]

157

158 def replacement(

159 q: torch.Tensor,

160 k: torch.Tensor,

161 v: torch.Tensor,

162 output_attn: torch.Tensor,

144 - output_quant: torch.Tensor,

163 scale: torch.Tensor,

164 ):

165 # attn output in quant_dtype

@@ -164,13 +182,10 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

182 return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])

183

184 inputs = [

167 - self.empty(5, self.num_heads, self.head_size, dtype=self.dtype), # q

168 - self.empty(5, self.num_heads, self.head_size, dtype=self.dtype), # k

169 - self.empty(5, self.num_heads, self.head_size, dtype=self.dtype), # v

170 - self.empty(

171 - 5, self.num_heads, self.head_size, dtype=self.dtype

172 - ), # attn_output

173 - self.empty_quant(5, self.num_heads * self.head_size), # quant_output

185 + self.empty(5, self.num_heads, self.head_size), # q

186 + self.empty(5, self.num_heads, self.head_size), # k

187 + self.empty(5, self.num_heads, self.head_size), # v

188 + self.empty(5, self.num_heads, self.head_size), # attn_output

189 empty_fp32(1, 1), # scale

190 ]

191

@@ -179,7 +194,9 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):

194 replacement,

195 inputs,

196 AttentionQuantPattern.wrap_trace_fn(

182 - AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only

197 + pm.fwd_only,

198 + AttentionQuantPattern.fx_view_to_reshape,

199 + AttentionQuantPattern.remove_noop_permutes,

200 ),

201 pm_pass,

202 )

@@ -279,7 +296,9 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

296 replacement,

297 inputs,

298 AttentionQuantPattern.wrap_trace_fn(

282 - AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only

299 + pm.fwd_only,

300 + AttentionQuantPattern.fx_view_to_reshape,

301 + AttentionQuantPattern.remove_noop_permutes,

302 ),

303 pm_pass,

304 )

vllm/compilation/fx_utils.py

+12 -2

@@ -6,7 +6,7 @@ from collections.abc import Iterable, Iterator

7 from torch import fx

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

9 -from torch._ops import OpOverload

9 +from torch._ops import OpOverload, OpOverloadPacket

12 def is_func(node: fx.Node, target) -> bool:

@@ -64,7 +64,17 @@ def find_getitem(node: fx.Node, idx: int) -> fx.Node:

66 # An auto-functionalization-aware utility for finding nodes with a specific op

67 -def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]:

67 +# Also handles op overload packets and finds all overloads

68 +def find_op_nodes(

69 + op: OpOverload | OpOverloadPacket, graph: fx.Graph

70 +) -> Iterator[fx.Node]:

71 + if isinstance(op, OpOverloadPacket):

72 + for overload in op.overloads():

73 + overload_op = getattr(op, overload)

74 + yield from find_op_nodes(overload_op, graph)

75 + return

76 +

77 + assert isinstance(op, OpOverload)

78 if not op._schema.is_mutable:

79 yield from graph.find_nodes(op="call_function", target=op)

vllm/compilation/matcher_utils.py new

+208

@@ -0,0 +1,208 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from abc import ABC, abstractmethod

4 +

5 +import torch

6 +from torch._higher_order_ops import auto_functionalized

7 +from torch._ops import OpOverload

8 +

9 +from vllm.config import get_current_vllm_config

10 +from vllm.model_executor.layers.layernorm import RMSNorm

11 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8

12 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

13 + QuantKey,

14 + _normalize_quant_group_shape,

15 + kFp8DynamicTensorSym,

16 + kFp8DynamicTokenSym,

17 + kFp8StaticTensorSym,

18 + kNvfp4Quant,

19 +)

20 +from vllm.platforms import current_platform

21 +

22 +RMS_OP = torch.ops._C.rms_norm.default

23 +RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default

24 +

25 +QUANT_OPS: dict[QuantKey, OpOverload] = {

26 + kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501

27 + kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501

28 + kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501

29 +}

30 +

31 +if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

32 + QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default # noqa: E501

33 +

34 +

35 +class MatcherCustomOp(ABC):

36 + def __init__(self, enabled: bool):

37 + config = get_current_vllm_config()

38 + self.model_dtype = config.model_config.dtype if config.model_config else None

39 + self.device = config.device_config.device if config.device_config else None

40 +

41 + self.enabled = enabled

42 + self.forward = self.forward_custom if enabled else self.forward_native

43 +

44 + @abstractmethod

45 + def forward_custom(self, *args, **kws):

46 + pass

47 +

48 + @abstractmethod

49 + def forward_native(self, *args, **kws):

50 + pass

51 +

52 + def __call__(self, *args, **kws):

53 + return self.forward(*args, **kws)

54 +

55 + def empty(self, *args, **kws):

56 + return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kws)

57 +

58 + def empty_f32(self, *args, **kws):

59 + return torch.empty(*args, dtype=torch.float32, device=self.device, **kws)

60 +

61 + def inputs(self) -> list[torch.Tensor]:

62 + """Utility for inputs to the pattern"""

63 + raise NotImplementedError

64 +

65 +

66 +class MatcherRMSNorm(MatcherCustomOp):

67 + def __init__(self, epsilon: float, enabled: bool | None = None):

68 + if enabled is None:

69 + enabled = RMSNorm.enabled()

70 +

71 + super().__init__(enabled)

72 + self.epsilon = epsilon

73 +

74 + def inputs(self):

75 + input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)

76 + weight = self.empty(16)

77 + return [input, weight]

78 +

79 + def forward_custom(

80 + self,

81 + input: torch.Tensor,

82 + weight: torch.Tensor,

83 + ) -> torch.Tensor:

84 + result = torch.empty_like(input)

85 + _, result = auto_functionalized(

86 + RMS_OP,

87 + result=result,

88 + input=input,

89 + weight=weight,

90 + epsilon=self.epsilon,

91 + )

92 +

93 + return result

94 +

95 + def forward_native(

96 + self,

97 + input: torch.Tensor,

98 + weight: torch.Tensor,

99 + ) -> torch.Tensor:

100 + return RMSNorm.forward_static(

101 + input, self.epsilon, input.size(-1), self.model_dtype, weight

102 + )

103 +

104 +

105 +class MatcherFusedAddRMSNorm(MatcherCustomOp):

106 + def __init__(self, epsilon: float, enabled: bool | None = None):

107 + if enabled is None:

108 + enabled = RMSNorm.enabled()

109 +

110 + super().__init__(enabled)

111 + self.epsilon = epsilon

112 +

113 + def inputs(self):

114 + input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)

115 + weight = self.empty(16)

116 + residual = self.empty(5, 16)

117 + return [input, weight, residual]

118 +

119 + def forward_custom(

120 + self,

121 + input: torch.Tensor,

122 + weight: torch.Tensor,

123 + residual: torch.Tensor,

124 + ) -> tuple[torch.Tensor, torch.Tensor]:

125 + _, result, residual = auto_functionalized(

126 + RMS_ADD_OP,

127 + input=input,

128 + residual=residual,

129 + weight=weight,

130 + epsilon=self.epsilon,

131 + )

132 +

133 + return result, residual

134 +

135 + def forward_native(

136 + self,

137 + input: torch.Tensor,

138 + weight: torch.Tensor,

139 + residual: torch.Tensor,

140 + ) -> tuple[torch.Tensor, torch.Tensor]:

141 + return RMSNorm.forward_static(

142 + input, self.epsilon, input.size(-1), self.model_dtype, weight, residual

143 + )

144 +

145 +

146 +class MatcherQuantFP8(MatcherCustomOp):

147 + def __init__(self, quant_key: QuantKey, enabled: bool | None = None):

148 + if enabled is None:

149 + enabled = QuantFP8.enabled()

150 +

151 + super().__init__(enabled)

152 + self.quant_key = quant_key

153 + assert quant_key in QUANT_OPS, f"unsupported quantization scheme {quant_key}"

154 + self.QUANT_OP = QUANT_OPS[quant_key]

155 +

156 + assert quant_key.dtype == current_platform.fp8_dtype(), (

157 + "Only QuantFP8 supported by"

158 + )

159 + assert quant_key.scale2 is None

160 + self.quant_fp8 = QuantFP8(quant_key.scale.static, quant_key.scale.group_shape)

161 +

162 + def forward_custom(

163 + self,

164 + input: torch.Tensor,

165 + scale: torch.Tensor | None = None,

166 + ) -> tuple[torch.Tensor, torch.Tensor]:

167 + result = torch.empty(

168 + input.shape, device=input.device, dtype=self.quant_key.dtype

169 + )

170 +

171 + if self.quant_key.scale.static:

172 + assert scale is not None

173 + _, result = auto_functionalized(

174 + self.QUANT_OP, result=result, input=input, scale=scale

175 + )

176 + return result, scale

177 + else:

178 + assert scale is None

179 + scale = self.make_scale(input)

180 + _, result, scale = auto_functionalized(

181 + self.QUANT_OP, result=result, input=input, scale=scale, scale_ub=None

182 + )

183 + return result, scale

184 +

185 + def forward_native(

186 + self,

187 + input: torch.Tensor,

188 + scale: torch.Tensor | None = None,

189 + ) -> tuple[torch.Tensor, torch.Tensor]:

190 + return self.quant_fp8(input, scale)

191 +

192 + def make_scale(self, input: torch.Tensor):

193 + normalized_group_shape = _normalize_quant_group_shape(

194 + input, self.quant_key.scale.group_shape

195 + )

196 + scale_shape = (

197 + input.shape[0] // normalized_group_shape[0],

198 + input.shape[1] // normalized_group_shape[1],

199 + )

200 +

201 + return torch.empty(scale_shape, device=input.device, dtype=torch.float32)

202 +

203 + def inputs(self) -> list[torch.Tensor]:

204 + input = self.empty(5, 16)

205 + if self.quant_key.scale.static:

206 + return [input, self.empty_f32(1, 1)]

207 +

208 + return [input]

vllm/compilation/monitor.py

@@ -22,6 +22,7 @@ def start_monitoring_torch_compile(vllm_config: VllmConfig):

22 import depyf

24 path.mkdir(parents=True, exist_ok=True)

25 + logger.debug("Dumping depyf output to %s", path)

26 global context_manager

27 context_manager = depyf.prepare_debug(path.as_posix())

28 context_manager.__enter__()

vllm/compilation/pass_manager.py

+20 -17

@@ -5,7 +5,7 @@ import functools

5 from torch import fx as fx

7 from vllm import envs

8 -from vllm.config import VllmConfig

8 +from vllm.config import VllmConfig, set_current_vllm_config

9 from vllm.logger import init_logger

10 from vllm.platforms import current_platform

11 from vllm.utils import set_env_var

@@ -88,27 +88,30 @@ class PostGradPassManager(CustomGraphPass):

89 def configure(self, config: VllmConfig):

90 self.pass_config = config.compilation_config.pass_config

91 - if self.pass_config.enable_noop:

92 - self.passes += [NoOpEliminationPass(config)]

94 - if self.pass_config.enable_sequence_parallelism:

95 - self.passes += [SequenceParallelismPass(config)]

96 - if self.pass_config.enable_async_tp:

97 - self.passes += [AsyncTPPass(config)]

92 + # Set the current vllm config to allow tracing CustomOp instances

93 + with set_current_vllm_config(config, check_compile=False):

94 + if self.pass_config.enable_noop:

95 + self.passes += [NoOpEliminationPass(config)]

99 - if self.pass_config.enable_fi_allreduce_fusion:

100 - self.passes += [AllReduceFusionPass(config)]

97 + if self.pass_config.enable_sequence_parallelism:

98 + self.passes += [SequenceParallelismPass(config)]

99 + if self.pass_config.enable_async_tp:

100 + self.passes += [AsyncTPPass(config)]

101

102 - if self.pass_config.enable_fusion:

103 - self.passes += [RMSNormQuantFusionPass(config)]

104 - self.passes += [ActivationQuantFusionPass(config)]

102 + if self.pass_config.enable_fi_allreduce_fusion:

103 + self.passes += [AllReduceFusionPass(config)]

104

106 - if self.pass_config.enable_attn_fusion:

107 - self.passes += [AttnFusionPass(config)]

105 + if self.pass_config.enable_fusion:

106 + self.passes += [RMSNormQuantFusionPass(config)]

107 + self.passes += [ActivationQuantFusionPass(config)]

108

109 - # needs a functional graph

110 - self.post_cleanup = PostCleanupPass(config)

111 - self.fix_functionalization = FixFunctionalizationPass(config)

109 + if self.pass_config.enable_attn_fusion:

110 + self.passes += [AttnFusionPass(config)]

111 +

112 + # needs a functional graph

113 + self.post_cleanup = PostCleanupPass(config)

114 + self.fix_functionalization = FixFunctionalizationPass(config)

115

116 # [HACK: Bug with Inductor graph partition and torch.compile cache]

117 # In PyTorch 2.9, torch.compile has a bug where the graph

vllm/compilation/vllm_inductor_pass.py

+2 -1

@@ -128,7 +128,8 @@ class VllmPatternMatcherPass(VllmInductorPass):

128 f" please add to dump_patterns if there are any errors.\n\n"

129 f"from torch._higher_order_ops.auto_functionalize import "

130 f"auto_functionalized as auto_functionalized\n"

131 - f"from torch._inductor.pattern_matcher import *",

131 + f"from torch._inductor.pattern_matcher import *\n"

132 + f"vllm = torch.ops.vllm",

133 file=f,

134 )

135

vllm/model_executor/layers/layernorm.py

+39 -20

@@ -178,14 +178,11 @@ class RMSNorm(CustomOp):

178 self.variance_size_override = (

179 None if var_hidden_size == hidden_size else var_hidden_size

180 )

181 + weight_dtype = dtype or torch.get_default_dtype()

182 self.has_weight = has_weight

182 - if dtype is not None:

183 - self.weight = torch.ones(hidden_size, dtype=dtype)

184 - else:

185 - self.weight = torch.ones(hidden_size)

183 + self.weight = torch.ones(hidden_size, dtype=weight_dtype)

184 if self.has_weight:

185 self.weight = nn.Parameter(self.weight)

188 - weight_dtype = self.weight.data.dtype

186

187 if current_platform.is_rocm():

188 self.rocm_norm_func = dispatch_rocm_rmsnorm_func(

@@ -195,47 +192,69 @@ class RMSNorm(CustomOp):

192 with_fused_add=True, dtype=weight_dtype

193 )

194

198 - def forward_native(

199 - self,

195 + @staticmethod

196 + def forward_static(

197 x: torch.Tensor,

198 + variance_epsilon: float,

199 + hidden_size: int,

200 + orig_dtype: torch.dtype,

201 + weight: torch.Tensor | None = None,

202 residual: torch.Tensor | None = None,

203 + variance_size_override: int | None = None,

204 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

205 """PyTorch-native implementation equivalent to forward()."""

204 - orig_dtype = x.dtype

206 x = x.to(torch.float32)

207 if residual is not None:

207 - x = x + residual.to(torch.float32)

208 + # residual promoted f16->f32 automatically,

209 + # otherwise Inductor eliminates the casts to and from f16,

210 + # increasing memory usage (and complicating pattern matching)

211 + x = x + residual

212 residual = x.to(orig_dtype)

213

210 - hidden_size = x.shape[-1]

211 - if hidden_size != self.hidden_size:

214 + if x.shape[-1] != hidden_size:

215 raise ValueError(

213 - "Expected hidden_size to be "

214 - f"{self.hidden_size}, but found: {hidden_size}"

216 + f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}"

217 )

218

217 - if self.variance_size_override is None:

219 + if variance_size_override is None:

220 x_var = x

221 else:

220 - if hidden_size < self.variance_size_override:

222 + if hidden_size < variance_size_override:

223 raise ValueError(

224 "Expected hidden_size to be at least "

223 - f"{self.variance_size_override}, but found: {hidden_size}"

225 + f"{variance_size_override}, but found: {hidden_size}"

226 )

227

226 - x_var = x[:, :, : self.variance_size_override]

228 + x_var = x[:, :, :variance_size_override]

229

230 variance = x_var.pow(2).mean(dim=-1, keepdim=True)

231

230 - x = x * torch.rsqrt(variance + self.variance_epsilon)

232 + x = x * torch.rsqrt(variance + variance_epsilon)

233 x = x.to(orig_dtype)

232 - if self.has_weight:

233 - x = x * self.weight

234 + if weight is not None:

235 + x = x * weight

236 if residual is None:

237 return x

238 else:

239 return x, residual

240

241 + def forward_native(

242 + self,

243 + x: torch.Tensor,

244 + residual: torch.Tensor | None = None,

245 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

246 + """PyTorch-native implementation equivalent to forward()."""

247 +

248 + return self.forward_static(

249 + x,

250 + self.variance_epsilon,

251 + self.hidden_size,

252 + x.dtype,

253 + self.weight.data if self.has_weight else None,

254 + residual,

255 + self.variance_size_override,

256 + )

257 +

258 def forward_cuda(

259 self,

260 x: torch.Tensor,