@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bd56c983d6
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 15, 21:00 UTC

[torch.compile] Fix RMSNorm + quant fusion in the non-cutlass-fp8 case, rename RedundantReshapesPass to NoopEliminationPass (#10902)

Signed-off-by: luka <luka@neuralmagic.com>

Luka Govedič committed Feb 28, 2025 at 18:20 UTC bd56c983d6fe8ff93bddd5faaf8d96e01c90fd83
9 files changed +249 -170
tests/compile/backend.py
+9 -4
@@ -13,21 +13,26 @@ class TestBackend:
13 This class provides a simple Inductor backend that can be used for testing.
14 It takes a list of custom passes and runs them after Inductor's passes.
15 It also saves the graph before and after the custom passes for inspection.
16 +
17 + Inductor config can be modified directly by editing the inductor_config
18 + property. This can be helpful for adding passes like the
19 + 'pre_grad_custom_pass' and the 'post_grad_custom_pre_pass'.
20 """
21
22 def __init__(self, *passes: Union[InductorPass, Callable[[fx.Graph],
23 None]]):
24 self.custom_passes = list(passes)
25 from torch._inductor import config
22 - self.current_config = config.shallow_copy_dict()
23 - self.current_config['force_disable_caches'] = True
24 - self.current_config['post_grad_custom_post_pass'] = self.post_pass
26 + self.inductor_config = config.shallow_copy_dict()
27 + self.inductor_config['force_disable_caches'] = True
28 + self.inductor_config['post_grad_custom_post_pass'] = self.post_pass
29
30 def __call__(self, graph: fx.GraphModule, example_inputs):
31 + self.graph_pre_compile = deepcopy(graph)
32 from torch._inductor.compile_fx import compile_fx
33 return compile_fx(graph,
34 example_inputs,
30 - config_patches=self.current_config)
35 + config_patches=self.inductor_config)
36
37 def post_pass(self, graph: fx.Graph):
38 self.graph_pre_pass = deepcopy(graph)
tests/compile/test_functionalization.py
+4 -4
@@ -9,7 +9,7 @@ from vllm.compilation.fix_functionalization import FixFunctionalizationPass
9 from vllm.compilation.fusion import (FUSED_OPS, FusionPass, QuantKey,
10 kFp8DynamicTokenSym, kFp8StaticTensorSym)
11 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
12 -from vllm.compilation.reshapes import RedundantReshapesPass
12 +from vllm.compilation.noop_elimination import NoOpEliminationPass
13 from vllm.config import CompilationConfig
14
15 from .backend import TestBackend
@@ -50,11 +50,11 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
50 torch.set_default_device("cuda")
51
52 config = CompilationConfig.PassConfig(enable_fusion=do_fusion,
53 - enable_reshape=True)
54 - reshape_pass = RedundantReshapesPass(config)
53 + enable_noop=True)
54 + noop_pass = NoOpEliminationPass(config)
55 fusion_pass = FusionPass.instance(config)
56
57 - passes = [reshape_pass, fusion_pass] if do_fusion else [reshape_pass]
57 + passes = [noop_pass, fusion_pass] if do_fusion else [noop_pass]
58 func_pass = FixFunctionalizationPass(config)
59 backend_func = TestBackend(*passes, func_pass)
60 backend_no_func = TestBackend(*passes)
tests/compile/test_fusion.py
+69 -58
@@ -5,23 +5,25 @@ import torch
5 from compressed_tensors.quantization import FP8_DTYPE
6
7 import vllm.envs as envs
8 +import vllm.plugins
9 from vllm.compilation.fusion import (FUSED_OPS, QUANT_OPS, FusedRMSQuantKey,
10 FusionPass, QuantKey)
11 from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe
11 -from vllm.compilation.reshapes import RedundantReshapesPass
12 -from vllm.config import CompilationConfig
12 +from vllm.compilation.noop_elimination import NoOpEliminationPass
13 +from vllm.config import CompilationConfig, CompilationLevel, VllmConfig
14 from vllm.model_executor.layers.layernorm import RMSNorm
15 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15 - apply_fp8_linear)
16 + CUTLASS_FP8_SUPPORTED, apply_fp8_linear, maybe_create_device_identity)
17
18 from .backend import TestBackend
19
20
21 class TestModel(torch.nn.Module):
22
22 - def __init__(self, hidden_size: int, eps: float, static: bool, *args,
23 - **kwargs):
23 + def __init__(self, hidden_size: int, eps: float, static: bool,
24 + cutlass_fp8_enabled: bool, *args, **kwargs):
25 super().__init__(*args, **kwargs)
26 + self.cutlass_fp8_enabled = cutlass_fp8_enabled
27 self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]
28 self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
29 if static:
@@ -41,7 +43,8 @@ class TestModel(torch.nn.Module):
43 self.w[0],
44 self.wscale[0],
45 self.scale[0],
44 - use_per_token_if_dynamic=True)
46 + use_per_token_if_dynamic=True,
47 + cutlass_fp8_supported=self.cutlass_fp8_enabled)
48 # make sure resid is used for replacement to work
49 y2, resid = self.norm[1](x2, resid)
50
@@ -49,7 +52,8 @@ class TestModel(torch.nn.Module):
52 self.w[1],
53 self.wscale[1],
54 self.scale[1],
52 - use_per_token_if_dynamic=True)
55 + use_per_token_if_dynamic=True,
56 + cutlass_fp8_supported=self.cutlass_fp8_enabled)
57 y3, resid = self.norm[2](x3, resid) # use resid here
58 return y3
59
@@ -59,60 +63,67 @@ class TestModel(torch.nn.Module):
63 @pytest.mark.parametrize("num_tokens", [7, 256, 533, 2048, 2049])
64 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
65 @pytest.mark.parametrize("static", [True, False])
66 +@pytest.mark.parametrize("cutlass_fp8_enabled",
67 + [True, False] if CUTLASS_FP8_SUPPORTED else [False])
68 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda",
69 reason="Only test on CUDA")
64 -def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static):
70 +def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
71 + cutlass_fp8_enabled):
72 torch.set_default_device("cuda")
73 torch.set_default_dtype(dtype)
74 torch.manual_seed(1)
75 + maybe_create_device_identity() # needed for certain non-cutlass fp8 paths
76
69 - # Reshape pass is needed for the fusion pass to work
70 - config = CompilationConfig.PassConfig(enable_fusion=True,
71 - enable_reshape=True)
72 - reshape_pass = RedundantReshapesPass(config)
73 - fusion_pass = FusionPass.instance(config)
74 -
75 - backend = TestBackend(reshape_pass, fusion_pass)
76 - model = TestModel(hidden_size, eps, static)
77 -
78 - # First dimension dynamic
79 - x = torch.rand(num_tokens, hidden_size)
80 - torch._dynamo.mark_dynamic(x, 0)
81 -
82 - result = model(x)
83 -
84 - model2 = torch.compile(model, backend=backend)
85 - result2 = model2(x)
86 -
87 - # Higher tol for dynamic, even higher for bfloat16
88 - if static:
89 - ATOL, RTOL = (1e-3, 1e-3)
90 - elif dtype == torch.float16:
91 - ATOL, RTOL = (2e-3, 2e-3)
92 - else:
93 - ATOL, RTOL = (1e-2, 1e-2)
94 -
95 - torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
96 -
97 - # Check substitution worked
98 - pre_nodes = backend.graph_pre_pass.nodes
99 - post_nodes = backend.graph_post_pass.nodes
100 -
101 - # static is per-tensor, dynamic is per-token
102 - key = QuantKey(dtype=FP8_DTYPE,
103 - static=static,
104 - per_tensor=static,
105 - symmetric=True)
106 - rms_quant = FUSED_OPS[FusedRMSQuantKey(key, False)]
107 - add_rms_quant = FUSED_OPS[FusedRMSQuantKey(key, True)]
108 - fp8_quant = QUANT_OPS[key]
109 -
110 - # In pre-nodes, fp8 quant should be present and fused kernels should not
111 - assert find_auto_fn_maybe(pre_nodes, rms_quant) is None
112 - assert find_auto_fn_maybe(pre_nodes, add_rms_quant) is None
113 - find_auto_fn(pre_nodes, fp8_quant)
114 -
115 - # In post-nodes, fused kernels should be present and fp8 quant should not
116 - find_auto_fn(post_nodes, rms_quant)
117 - find_auto_fn(post_nodes, add_rms_quant)
118 - assert find_auto_fn_maybe(post_nodes, fp8_quant) is None
77 + vllm_config = VllmConfig(compilation_config=CompilationConfig(
78 + level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm"]))
79 + with vllm.config.set_current_vllm_config(vllm_config):
80 + # Reshape pass is needed for the fusion pass to work
81 + config = CompilationConfig.PassConfig(enable_fusion=True,
82 + enable_noop=True)
83 + noop_pass = NoOpEliminationPass(config)
84 + fusion_pass = FusionPass.instance(config)
85 +
86 + backend = TestBackend(noop_pass, fusion_pass)
87 + model = TestModel(hidden_size, eps, static, cutlass_fp8_enabled)
88 +
89 + # First dimension dynamic
90 + x = torch.rand(num_tokens, hidden_size)
91 + torch._dynamo.mark_dynamic(x, 0)
92 +
93 + result = model(x)
94 +
95 + model2 = torch.compile(model, backend=backend)
96 + result2 = model2(x)
97 +
98 + # Higher tol for dynamic, even higher for bfloat16
99 + if static:
100 + ATOL, RTOL = (1e-3, 1e-3)
101 + elif dtype == torch.float16:
102 + ATOL, RTOL = (2e-3, 2e-3)
103 + else:
104 + ATOL, RTOL = (1e-2, 1e-2)
105 +
106 + torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
107 +
108 + # Check substitution worked
109 + pre_nodes = backend.graph_pre_pass.nodes
110 + post_nodes = backend.graph_post_pass.nodes
111 +
112 + # static is per-tensor, dynamic is per-token
113 + key = QuantKey(dtype=FP8_DTYPE,
114 + static=static,
115 + per_tensor=static,
116 + symmetric=True)
117 + rms_quant = FUSED_OPS[FusedRMSQuantKey(key, False)]
118 + add_rms_quant = FUSED_OPS[FusedRMSQuantKey(key, True)]
119 + fp8_quant = QUANT_OPS[key]
120 +
121 + # In pre-nodes, fp8 quant should be there and fused kernels should not
122 + assert find_auto_fn_maybe(pre_nodes, rms_quant) is None
123 + assert find_auto_fn_maybe(pre_nodes, add_rms_quant) is None
124 + find_auto_fn(pre_nodes, fp8_quant)
125 +
126 + # In post-nodes, fused kernels should be there and fp8 quant should not
127 + find_auto_fn(post_nodes, rms_quant)
128 + find_auto_fn(post_nodes, add_rms_quant)
129 + assert find_auto_fn_maybe(post_nodes, fp8_quant) is None
vllm/compilation/noop_elimination.py new
+135
@@ -0,0 +1,135 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Iterable, Union
4 +
5 +import torch.fx
6 +from torch import SymInt
7 +
8 +from vllm.logger import init_logger
9 +
10 +from .fx_utils import is_func
11 +from .vllm_inductor_pass import VllmInductorPass
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +class NoOpEliminationPass(VllmInductorPass):
17 + """
18 + This is an inductor pass that removes redundant reshape/slice operations.
19 + It is required for RMSNorm-quant fusion to work properly.
20 + That's because apply_fp8_linear adds a reshape, which is redundant
21 + in the 2D-case. Additionally, torch internal no-op elimination pass does
22 + not handle certain slice variants.
23 +
24 + Example graph 1:
25 + getitem_1: "f16[s0, 4096]" = ...
26 + view_1: "f16[s0, 4096]" = torch.reshape(getitem_1, [-1, 4096])
27 + at = auto_functionalized(static_scaled_fp8_quant, input = view_1, ...)
28 + out: "f8e4m3fn[s0, 4096]" = at[1]
29 +
30 + Can be replaced with:
31 + getitem_1: "f16[s0, 4096]" = ...
32 + at = auto_functionalized(static_scaled_fp8_quant, input = getitem_1, ...)
33 + out: "f8e4m3fn[s0, 4096]" = at[1]
34 +
35 + Example graph 2:
36 + arg0: "s0" = SymInt(s0)
37 + scaled_mm: "f16[s0, 4096]" = ...
38 + slice_1: "f16[s0, 4096]" = torch.slice(scaled_mm, -1, 0, arg0)
39 + at = auto_functionalized(fused_add_rms_norm, input = slice_1, ...)
40 + out: "f16[s0, 4096]" = torch.slice_scatter(scaled_mm, at[1], 0, 0, arg0)
41 +
42 + Can be replaced with:
43 + arg0: "s0" = SymInt(s0)
44 + scaled_mm: "f16[s0, 4096]" = ...
45 + at = auto_functionalized(fused_add_rms_norm, input = scaled_mm, ...)
46 + out: "f16[s0, 4096]" = at[1]
47 +
48 + TODO(luka): This is currently tested in test_fusion,
49 + but separate tests could be good.
50 + """
51 +
52 + def __call__(self, graph: torch.fx.Graph):
53 + self.begin()
54 + self.dump_graph(graph, "before_noop_elimination")
55 + count = 0
56 + # Remove no-op reshapes/views:
57 + for node in graph.nodes:
58 + if is_func(node, torch.ops.aten.reshape.default):
59 + input, shape = node.args[:2]
60 + input_shape = input.meta["val"].shape
61 + if len(shape) != len(input_shape):
62 + # Reshape changing rank, skip
63 + continue
64 +
65 + if shape.count(-1) > 1:
66 + # Invalid reshape args, skip
67 + continue
68 +
69 + if self.all_dims_equivalent(shape, input_shape):
70 + node.replace_all_uses_with(input)
71 + graph.erase_node(node)
72 + count += 1
73 +
74 + elif is_func(node, torch.ops.aten.slice.Tensor):
75 + input, dim_index, start, end = node.args[:4]
76 + input_shape = input.meta["val"].shape
77 + i_dim = input_shape[dim_index]
78 +
79 + if start == 0 and self.dims_equivalent(end, i_dim):
80 + node.replace_all_uses_with(input)
81 + graph.erase_node(node)
82 + count += 1
83 +
84 + elif is_func(node, torch.ops.aten.slice_scatter.default):
85 + base, view, dim_index, start, end = node.args[:5]
86 + base_shape = base.meta["val"].shape
87 + view_shape = view.meta["val"].shape
88 +
89 + view_dim = view_shape[dim_index]
90 +
91 + # Check that view fully covers base and the full view is used
92 + # (if the view fully covered the base after slicing but was not
93 + # fully used, we could replace slice_scatter with a simple slice
94 + # but that's a niche case).
95 + if (base_shape == view_shape and start == 0
96 + and self.dims_equivalent(end, view_dim)):
97 + node.replace_all_uses_with(view)
98 + graph.erase_node(node)
99 + count += 1
100 +
101 + logger.debug("Removed %s no-op reshapes and slices", count)
102 + self.dump_graph(graph, "after_noop_elimination")
103 + self.end_and_log()
104 +
105 + def all_dims_equivalent(self, dims: Iterable[Union[int, torch.fx.Node]],
106 + i_dims: Iterable[Union[int, SymInt]]):
107 + return all(
108 + self.dims_equivalent(s, i_s) for s, i_s in zip(dims, i_dims))
109 +
110 + def dims_equivalent(self, dim: Union[int, torch.fx.Node],
111 + i_dim: Union[int, SymInt]) -> bool:
112 + """
113 + This function checks if two dimensions are equivalent.
114 + :param dim: The dimension arg to reshape/slice
115 + :param i_dim: The corresponding dimension in the input tensor
116 + :return: Are the dimensions equivalent?
117 +
118 + There are three cases in which the dimensions are equivalent:
119 + 1. The dimensions are equal (both integers)
120 + 2. The reshape dimension is -1 (i.e. inferred)
121 + 3. The dimensions both correspond to the same SymInt
122 +
123 + While case 2 does not guarantee the dimensions are equal,
124 + they are equal if all other dimensions are equal.
125 +
126 + In case 3, the reshape dimension is a torch.fx.Node,
127 + and its value is a SymInt. That value is equal to the
128 + input dimension.
129 +
130 + """
131 + # Case 1 and 2
132 + if dim == i_dim or dim == -1:
133 + return True
134 + # Case 3
135 + return isinstance(dim, torch.fx.Node) and dim.meta["val"] == i_dim
vllm/compilation/pass_manager.py
+4 -4
@@ -11,7 +11,7 @@ from vllm.logger import init_logger
11 from .fix_functionalization import FixFunctionalizationPass
12 from .fusion import FusionPass
13 from .inductor_pass import InductorPass
14 -from .reshapes import RedundantReshapesPass
14 +from .noop_elimination import NoOpEliminationPass
15
16 logger = init_logger(__name__)
17
@@ -36,7 +36,7 @@ class PostGradPassManager(Parent):
36
37 The order of the post-grad post-passes is:
38 1. passes (constructor parameter)
39 - 2. default passes (RedundantReshapesPass, FusionPass)
39 + 2. default passes (NoopEliminationPass, FusionPass)
40 3. config["post_grad_custom_post_pass"] (if it exists)
41 4. fix_functionalization
42 This way, all passes operate on a functionalized graph.
@@ -54,8 +54,8 @@ class PostGradPassManager(Parent):
54
55 def configure(self, pass_config: CompilationConfig.PassConfig):
56 self.pass_config = pass_config
57 - if pass_config.enable_reshape:
58 - self.passes += [RedundantReshapesPass(pass_config)]
57 + if pass_config.enable_noop:
58 + self.passes += [NoOpEliminationPass(pass_config)]
59
60 if pass_config.enable_fusion:
61 self.passes += [FusionPass.instance(pass_config)]
vllm/compilation/reshapes.py deleted
-90
@@ -1,90 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -from typing import Union
4 -
5 -import torch.fx
6 -from torch import SymInt
7 -
8 -from vllm.logger import init_logger
9 -
10 -from .fx_utils import is_func
11 -from .vllm_inductor_pass import VllmInductorPass
12 -
13 -logger = init_logger(__name__)
14 -
15 -
16 -class RedundantReshapesPass(VllmInductorPass):
17 - """
18 - This is an inductor pass that removes redundant reshape operations.
19 - It is required for RMSNorm-quant fusion to work properly.
20 - That's because apply_fp8_linear adds a reshape, which is redundant
21 - in the 2D-case.
22 -
23 - Example graph:
24 -
25 - getitem_1: "f16[s0, 4096]" = ...
26 - view_1: "f16[s0, 4096]" = torch.reshape(getitem_1, [-1, 4096])
27 - at = auto_functionalized(static_scaled_fp8_quant, input = view_1, ...)
28 - out: "f8e4m3fn[s0, 4096]" = at[1]
29 -
30 - Can be replaced with:
31 - getitem_1: "f16[s0, 4096]" = ...
32 - at = auto_functionalized(static_scaled_fp8_quant, input = getitem_1, ...)
33 - out: "f8e4m3fn[s0, 4096]" = at[1]
34 - """
35 -
36 - def __call__(self, graph: torch.fx.Graph):
37 - self.begin()
38 - self.dump_graph(graph, "before_reshapes")
39 - count = 0
40 - # Remove no-op reshapes/views:
41 - for node in graph.nodes:
42 - if is_func(node, torch.ops.aten.reshape.default):
43 - input, shape = node.args[:2]
44 - input_shape = input.meta["val"].shape
45 - if len(shape) != len(input_shape):
46 - # Reshape changing rank, skip
47 - continue
48 -
49 - if shape.count(-1) > 1:
50 - # Invalid reshape args, skip
51 - continue
52 -
53 - if all(
54 - self.dims_equivalent(s, i_s)
55 - for s, i_s in zip(shape, input_shape)):
56 - node.replace_all_uses_with(input)
57 - graph.erase_node(node)
58 - count += 1
59 -
60 - logger.debug("Removed %s no-op reshapes", count)
61 -
62 - self.dump_graph(graph, "after_reshapes")
63 - self.end_and_log()
64 -
65 - def dims_equivalent(self, dim: Union[int, torch.fx.Node],
66 - i_dim: Union[int, SymInt]) -> bool:
67 - """
68 - This function checks if two dimensions are equivalent.
69 - :param dim: The dimension arg to reshape
70 - :param i_dim: The corresponding dimension in the input tensor
71 - :return: Are the dimensions equivalent?
72 -
73 - There are three cases in which the dimensions are equivalent:
74 - 1. The dimensions are equal (both integers)
75 - 2. The reshape dimension is -1 (i.e. inferred)
76 - 3. The dimensions both correspond to the same SymInt
77 -
78 - While case 2 does not guarantee the dimensions are equal,
79 - they are equal if all other dimensions are equal.
80 -
81 - In case 3, the reshape dimension is a torch.fx.Node,
82 - and its value is a SymInt. That value is equal to the
83 - input dimension.
84 -
85 - """
86 - # Case 1 and 2
87 - if dim == i_dim or dim == -1:
88 - return True
89 - # Case 3
90 - return isinstance(dim, torch.fx.Node) and dim.meta["val"] == i_dim
vllm/compilation/vllm_inductor_pass.py
+16 -2
@@ -28,8 +28,8 @@ class VllmInductorPass(InductorPass):
28 self.config = config
29 self.pass_name = self.__class__.__name__
30
31 - def dump_graph(self, graph: torch.fx.Graph, stage: str):
32 - if stage in self.config.dump_graph_stages:
31 + def dump_graph(self, graph: torch.fx.Graph, stage: str, always=False):
32 + if stage in self.config.dump_graph_stages or always:
33 # Make sure filename includes rank in the distributed setting
34 parallel = p_is_init() and get_tp_world_size() > 1
35 rank = f"-{get_tp_rank()}" if parallel else ""
@@ -49,3 +49,17 @@ class VllmInductorPass(InductorPass):
49 self._end_time = time.perf_counter_ns()
50 duration_ms = float(self._end_time - self._start_time) / 1.0e6
51 logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)
52 +
53 +
54 +class PrinterInductorPass(VllmInductorPass):
55 +
56 + def __init__(self,
57 + name: str,
58 + config: CompilationConfig.PassConfig,
59 + always=False):
60 + super().__init__(config)
61 + self.name = name
62 + self.always = always
63 +
64 + def __call__(self, graph: torch.fx.Graph):
65 + self.dump_graph(graph, self.name, always=self.always)
vllm/config.py
+6 -7
@@ -2993,13 +2993,13 @@ class CompilationConfig(BaseModel):
2993 Each pass defines its own stages (before, after, maybe in-between).
2994 - dump_graph_dir: directory to dump the graphs. Default is .
2995 - enable_fusion: whether to enable the custom fusion pass.
2996 - - enable_reshape: whether to enable the custom reshape elimination pass.
2997 - TODO better pass enabling system.
2996 + - enable_noop: whether to enable the custom no-op elimination pass.
2997 + TODO(luka) better pass enabling system.
2998 """
2999 dump_graph_stages: List[str] = Field(default_factory=list)
3000 dump_graph_dir: Path = Field(default=Path("."))
3001 enable_fusion: bool = True
3002 - enable_reshape: bool = True
3002 + enable_noop: bool = True
3003
3004 def uuid(self):
3005 """
@@ -3008,13 +3008,12 @@ class CompilationConfig(BaseModel):
3008 Do not include dump_graph_* in the hash - they don't affect
3009 compilation.
3010 """
3011 - dict_ = self.model_dump(
3012 - include={"enable_fusion", "enable_reshape"})
3011 + dict_ = self.model_dump(include={"enable_fusion", "enable_noop"})
3012 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
3013 return hashlib.sha256(encoded).digest()
3014
3015 def model_post_init(self, __context: Any) -> None:
3017 - if not self.enable_reshape and self.enable_fusion:
3016 + if not self.enable_noop and self.enable_fusion:
3017 logger.warning_once(
3018 "Fusion enabled but reshape elimination disabled. "
3019 "RMSNorm + quant (fp8) fusion might not work")
@@ -3411,7 +3410,7 @@ class VllmConfig:
3410 self.compilation_config.use_inductor = True
3411 self.compilation_config.cudagraph_num_of_warmups = 1
3412 self.compilation_config.pass_config.enable_fusion = False
3414 - self.compilation_config.pass_config.enable_reshape = False
3413 + self.compilation_config.pass_config.enable_noop = False
3414 self.compilation_config.level = CompilationLevel.PIECEWISE
3415
3416 self._set_cudagraph_sizes()
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+6 -1
@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple, Union
5 import torch
6
7 from vllm import _custom_ops as ops
8 +from vllm.config import CompilationLevel, get_current_vllm_config
9 from vllm.platforms import current_platform
10
11 # Input scaling factors are no longer optional in _scaled_mm starting
@@ -161,10 +162,14 @@ def apply_fp8_linear(
162 # Note: we pad the input because torch._scaled_mm is more performant
163 # for matrices with batch dimension > 16.
164 # This could change in the future.
165 + # We also don't pad when using torch.compile,
166 + # as it breaks with dynamic shapes.
167 + config = get_current_vllm_config().compilation_config
168 + do_pad = config.level < CompilationLevel.PIECEWISE
169 qinput, x_scale = ops.scaled_fp8_quant(
170 input_2d,
171 input_scale,
167 - num_token_padding=17,
172 + num_token_padding=17 if do_pad else None,
173 use_per_token_if_dynamic=use_per_token_if_dynamic)
174
175 per_tensor_weights = (weight_scale.numel() == 1)