[Fusion] normalize fusion naming and enable e2e test (#4693)
### What this PR does / why we need it?
This PR standardizes the fusion naming, changing
`enable_quantization_fusion` to `fuse_norm_quant`, and enables e2e
testing.
### Does this PR introduce _any_ user-facing change?
N/A
### How was this patch tested?
CI passed with new added/existing test.
- vLLM version: v0.12.0
- vLLM main:
ad32e3e19c
---------
Signed-off-by: wxsIcey <1790571317@qq.com>
This commit is contained in:
130
tests/e2e/singlecard/compile/backend.py
Normal file
130
tests/e2e/singlecard/compile/backend.py
Normal file
@@ -0,0 +1,130 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
||||
# Copyright 2023 The vLLM team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from copy import deepcopy
|
||||
from typing import Any, Callable, List, Optional, Sequence
|
||||
|
||||
import torch.fx as fx
|
||||
from torch._inductor.decomposition import select_decomp_table
|
||||
from vllm.compilation.fx_utils import OpOverload
|
||||
from vllm.config import get_current_vllm_config
|
||||
|
||||
from vllm_ascend.compilation.compiler_interface import compile_fx
|
||||
|
||||
|
||||
class TestBackend:
|
||||
"""
|
||||
A custom compilation backend for testing operator fusion passes.
|
||||
It applies the AddRMSNormQuantFusionPass during graph compilation and
|
||||
records the FX graph before and after the transformation.
|
||||
"""
|
||||
|
||||
def __init__(self, custom_passes: Optional[List[Any]] = None):
|
||||
vllm_config = get_current_vllm_config()
|
||||
compile_config = vllm_config.compilation_config
|
||||
self.inductor_config = compile_config.inductor_compile_config
|
||||
self.inductor_config["graph_fusion_manager"] = self.post_pass
|
||||
self.custom_passes = custom_passes
|
||||
|
||||
# Placeholders to store FX graphs for verification
|
||||
self.graph_pre_pass = None
|
||||
self.graph_post_pass = None
|
||||
|
||||
def post_pass(self,
|
||||
graph: fx.Graph,
|
||||
runtime_shape: int | None = None) -> fx.Graph:
|
||||
"""
|
||||
Apply custom graph transformation passes.
|
||||
"""
|
||||
self.graph_pre_pass = deepcopy(graph)
|
||||
if self.custom_passes is not None:
|
||||
for pass_ in self.custom_passes:
|
||||
pass_(graph)
|
||||
self.graph_post_pass = deepcopy(graph)
|
||||
return graph
|
||||
|
||||
def compile(
|
||||
self,
|
||||
graph: fx.GraphModule,
|
||||
example_inputs: list[Any],
|
||||
compiler_config: dict[str, Any],
|
||||
runtime_shape: Optional[int] = None,
|
||||
key: Optional[str] = None
|
||||
) -> tuple[Optional[Callable], Optional[Any]]:
|
||||
"""
|
||||
Compile the FX graph using vLLM's Ascend compiler interface.
|
||||
Wraps the post-pass logic into the inner_compile callback.
|
||||
"""
|
||||
|
||||
def compile_inner(graph, example_inputs):
|
||||
current_pass_manager = compiler_config["graph_fusion_manager"]
|
||||
return current_pass_manager(graph, runtime_shape)
|
||||
|
||||
decompositions = select_decomp_table()
|
||||
compiled_fn = compile_fx(
|
||||
graph=graph,
|
||||
example_inputs=example_inputs,
|
||||
inner_compile=compile_inner,
|
||||
decompositions=decompositions,
|
||||
)
|
||||
return compiled_fn, None
|
||||
|
||||
def __call__(self, gm: fx.GraphModule,
|
||||
example_inputs: Optional[List[Any]]):
|
||||
"""
|
||||
Make the backend callable by torch.compile().
|
||||
Returns a compiled executable function.
|
||||
"""
|
||||
assert example_inputs is not None
|
||||
compiled_fn, _ = self.compile(
|
||||
gm,
|
||||
example_inputs,
|
||||
compiler_config={"graph_fusion_manager": self.post_pass},
|
||||
runtime_shape=None,
|
||||
key=None,
|
||||
)
|
||||
return compiled_fn
|
||||
|
||||
def find_nodes_by_target(self, graph: fx.GraphModule,
|
||||
target: OpOverload) -> List[fx.Node]:
|
||||
"""Helper to find all FX nodes that call a specific operator."""
|
||||
return [
|
||||
node for node in graph.graph.nodes
|
||||
if hasattr(node, 'target') and node.target == target
|
||||
]
|
||||
|
||||
def check_before_ops(self,
|
||||
ops: Sequence[OpOverload],
|
||||
fully_replaced: bool = True):
|
||||
"""
|
||||
Verify that the original (unfused) operators exist before the pass
|
||||
and are fully removed afterward (if fully_replaced=True).
|
||||
"""
|
||||
for op in ops:
|
||||
num_pre = len(self.find_nodes_by_target(self.graph_pre_pass, op))
|
||||
num_post = len(self.find_nodes_by_target(self.graph_post_pass, op))
|
||||
print(f"Op {op}: pre={num_pre}, post={num_post}")
|
||||
|
||||
assert num_pre > 0, f"Op {op} not found in pre-pass graph"
|
||||
if fully_replaced:
|
||||
assert num_post == 0, f"Unexpected op {op} in post-pass graph: {num_post} nodes remain"
|
||||
|
||||
def check_after_ops(self, ops: Sequence[OpOverload]):
|
||||
"""Verify that the fused operator appears in the transformed graph."""
|
||||
for op in ops:
|
||||
num_post = len(self.find_nodes_by_target(self.graph_post_pass, op))
|
||||
print(f"Op {op}: post={num_post}")
|
||||
assert num_post > 0, f"Op {op} not found in post-pass graph"
|
||||
Reference in New Issue
Block a user