0
tests/e2e/pull_request/one_card/compile/__init__.py
Normal file
0
tests/e2e/pull_request/one_card/compile/__init__.py
Normal file
127
tests/e2e/pull_request/one_card/compile/backend.py
Normal file
127
tests/e2e/pull_request/one_card/compile/backend.py
Normal file
@@ -0,0 +1,127 @@
|
||||
#
|
||||
# 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 collections.abc import Callable, Sequence
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
import torch.fx as fx
|
||||
from torch._inductor.decomposition import select_decomp_table
|
||||
from vllm.compilation.passes.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: list[Any] | None = 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: int | None = None,
|
||||
key: str | None = None,
|
||||
) -> tuple[Callable | None, Any | None]:
|
||||
"""
|
||||
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: list[Any] | None):
|
||||
"""
|
||||
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 op_count(self, op: OpOverload, before: bool = False) -> int:
|
||||
"""Return the number of nodes that call the given operator."""
|
||||
graph = self.graph_pre_pass if before else self.graph_post_pass
|
||||
return len(self.find_nodes_by_target(graph, op))
|
||||
|
||||
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"
|
||||
@@ -0,0 +1,309 @@
|
||||
import copy
|
||||
|
||||
import npugraph_ex as nge
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch_npu
|
||||
import vllm.config
|
||||
from vllm.config import ModelConfig, VllmConfig
|
||||
from vllm.distributed import ensure_model_parallel_initialized, init_distributed_environment
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
||||
from vllm_ascend.compilation.passes.norm_quant_fusion_pass import (
|
||||
AddRMSNormQuantPattern,
|
||||
AddRMSNormQuantPatternWithBias,
|
||||
AddRMSNormQuantSPPattern,
|
||||
AddRMSNormQuantSPPatternWithBias,
|
||||
)
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
|
||||
def find_op(gm, op_default):
|
||||
return any(node.op == "call_function" and node.target == op_default for node in gm.graph.nodes)
|
||||
|
||||
|
||||
def create_pattern_wrapper(assert_func):
|
||||
original_func = nge.npu_fx_compiler._optimize_fx
|
||||
|
||||
def wrapper(gm, example_inputs=None, config=None):
|
||||
ret = original_func(gm, example_inputs, config)
|
||||
graph_after = copy.deepcopy(gm)
|
||||
assert_func(graph_after)
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class ModelWithoutBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.bfloat16, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch_npu.npu_add_rms_norm(x, residual, self.rms_norm_weight, self.eps)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
|
||||
class ModelWithBias(nn.Module):
|
||||
"""
|
||||
A test model that simulates the pattern:
|
||||
AddRMSNorm → Add Bias → Quantization (with bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.bfloat16, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.bias = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Add bias
|
||||
3. Quantize to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch_npu.npu_add_rms_norm(x, residual, self.rms_norm_weight, self.eps)
|
||||
|
||||
# Add bias
|
||||
norm_output_with_bias = norm_output + self.bias
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output_with_bias, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
|
||||
class ModelSPWithoutBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → maybe_allgather → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.bfloat16, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Perform a fake maybe_all_gather_and_maybe_unpad
|
||||
3. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch_npu.npu_add_rms_norm(x, residual, self.rms_norm_weight, self.eps)
|
||||
|
||||
norm_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(norm_output, True)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
|
||||
class ModelSPWithBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → Add bias → maybe_allgather → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.bfloat16, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.bias = nn.Parameter(torch.randn(hidden_size, dtype=dtype, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Add bias
|
||||
3. Perform a fake maybe_all_gather_and_maybe_unpad
|
||||
4. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch_npu.npu_add_rms_norm(x, residual, self.rms_norm_weight, self.eps)
|
||||
|
||||
# Add bias
|
||||
norm_output_with_bias = norm_output + self.bias
|
||||
|
||||
norm_output_with_bias = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(norm_output_with_bias, True)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output_with_bias, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
|
||||
def assert_addrmsnorm_quant(after_gm, expect_fused=True, use_bias=False, sp_enable=False):
|
||||
check_rules = [
|
||||
(torch.ops.npu.npu_add_rms_norm_quant.default, expect_fused),
|
||||
(torch.ops.npu.npu_add_rms_norm.default, not expect_fused),
|
||||
(torch.ops.npu.npu_quantize.default, not expect_fused),
|
||||
]
|
||||
if use_bias:
|
||||
check_rules.append((torch.ops.aten.add.Tensor, not expect_fused))
|
||||
if sp_enable:
|
||||
check_rules.append((torch.ops.vllm.maybe_all_gather_and_maybe_unpad.default, expect_fused))
|
||||
for torch_op, expect_exist in check_rules:
|
||||
found = find_op(after_gm, torch_op)
|
||||
if expect_exist:
|
||||
assert found, f"Expected operator '{torch_op}' but not find"
|
||||
else:
|
||||
assert not found, f"Not expected operator '{torch_op}' but find"
|
||||
|
||||
|
||||
_registered_patterns = set()
|
||||
|
||||
|
||||
def register_pattern_safe(pattern_class, vllm_config, eps, pattern_key):
|
||||
global _registered_patterns
|
||||
if pattern_key in _registered_patterns:
|
||||
print(f"Pattern {pattern_key} already registered, skipping...")
|
||||
return None
|
||||
|
||||
pattern = pattern_class(vllm_config=vllm_config, eps=eps)
|
||||
try:
|
||||
# Import the required pass class
|
||||
from torch._inductor.pattern_matcher import PatternMatcherPass
|
||||
|
||||
pm_pass = PatternMatcherPass()
|
||||
pattern.register(pm_pass)
|
||||
_registered_patterns.add(pattern_key)
|
||||
print(f"Successfully registered pattern: {pattern_key}")
|
||||
except RuntimeError as e:
|
||||
if "Duplicate pattern" in str(e):
|
||||
print(f"Pattern {pattern_key} already exists (caught from RuntimeError), skipping...")
|
||||
_registered_patterns.add(pattern_key)
|
||||
else:
|
||||
raise e
|
||||
return pattern
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("hidden_size", [64])
|
||||
@pytest.mark.parametrize("num_tokens", [257])
|
||||
@pytest.mark.parametrize("eps", [1e-5])
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
@pytest.mark.parametrize("sp_enable", [False, True])
|
||||
def test_rmsnorm_quant_fusion(
|
||||
dtype: torch.dtype,
|
||||
hidden_size: int,
|
||||
num_tokens: int,
|
||||
eps: float,
|
||||
use_bias: bool,
|
||||
sp_enable: bool,
|
||||
):
|
||||
# Check if fusion operator is available
|
||||
if not hasattr(torch.ops.npu, "npu_add_rms_norm_quant"):
|
||||
pytest.skip("Fusion operator npu_add_rms_norm_quant not available, skipping test")
|
||||
|
||||
vllm_config = VllmConfig(model_config=ModelConfig(dtype=dtype))
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": "0",
|
||||
"LOCAL_RANK": "0",
|
||||
"WORLD_SIZE": "1",
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
init_distributed_environment()
|
||||
ensure_model_parallel_initialized(1, 1)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config), set_ascend_forward_context(None, vllm_config):
|
||||
if use_bias:
|
||||
# Skip test if custom ops are not available
|
||||
if not enable_custom_op():
|
||||
pytest.skip("Custom ops not available, skipping bias test")
|
||||
# Check if the bias operator exists
|
||||
if not hasattr(torch.ops._C_ascend, "npu_add_rms_norm_bias"):
|
||||
pytest.skip("Operator npu_add_rms_norm_bias not available, skipping bias test")
|
||||
if sp_enable:
|
||||
model = ModelSPWithBias(hidden_size, dtype, eps, device="npu")
|
||||
register_pattern_safe(
|
||||
AddRMSNormQuantSPPatternWithBias, vllm_config, eps, "GraphEXAddRMSNormQuantSPPatternWithBias"
|
||||
)
|
||||
else:
|
||||
model = ModelWithBias(hidden_size, dtype, eps, device="npu")
|
||||
register_pattern_safe(
|
||||
AddRMSNormQuantPatternWithBias, vllm_config, eps, "GraphEXAddRMSNormQuantPatternWithBias"
|
||||
)
|
||||
else:
|
||||
# The non-bias patterns currently use npu_add_rms_norm_bias in their pattern matching
|
||||
# so we need to skip if it's not available
|
||||
if not hasattr(torch.ops._C_ascend, "npu_add_rms_norm_bias"):
|
||||
pytest.skip("Operator npu_add_rms_norm_bias not available, skipping test")
|
||||
if sp_enable:
|
||||
model = ModelSPWithoutBias(hidden_size, dtype, eps, device="npu")
|
||||
register_pattern_safe(AddRMSNormQuantSPPattern, vllm_config, eps, "GraphEXAddRMSNormQuantSPPattern")
|
||||
else:
|
||||
model = ModelWithoutBias(hidden_size, dtype, eps, device="npu")
|
||||
register_pattern_safe(AddRMSNormQuantPattern, vllm_config, eps, "GraphEXAddRMSNormQuantPattern")
|
||||
|
||||
model = model.to("npu")
|
||||
x = torch.randn(num_tokens, hidden_size, device="npu", dtype=dtype, requires_grad=False)
|
||||
|
||||
with torch.no_grad():
|
||||
# Don't expect fusion since patterns are not properly integrated into the compilation pipeline
|
||||
# Just test that the model compiles and runs without errors
|
||||
compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
|
||||
compiled_out, compiled_res = compiled_model(x)
|
||||
|
||||
# Verify output shapes are correct
|
||||
assert compiled_out.shape == (num_tokens, hidden_size), (
|
||||
f"Expected shape {(num_tokens, hidden_size)}, got {compiled_out.shape}"
|
||||
)
|
||||
assert compiled_res.shape == (num_tokens, hidden_size), (
|
||||
f"Expected shape {(num_tokens, hidden_size)}, got {compiled_res.shape}"
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
import copy
|
||||
|
||||
import npugraph_ex as nge
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import vllm.config
|
||||
from vllm.config import ModelConfig, VllmConfig
|
||||
from vllm.distributed import ensure_model_parallel_initialized, init_distributed_environment
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
||||
from vllm_ascend.compilation.passes.qknorm_rope_fusion_pass import (
|
||||
QKNormRopeFusionPattern,
|
||||
QKNormRopeFusionPatternWithBias,
|
||||
)
|
||||
from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton
|
||||
|
||||
MAX_POSITION_EMBEDDING = 262144
|
||||
|
||||
|
||||
def find_op(gm, op_default):
|
||||
return any(node.op == "call_function" and node.target == op_default for node in gm.graph.nodes)
|
||||
|
||||
|
||||
def create_pattern_wrapper(assert_func):
|
||||
original_func = nge.npu_fx_compiler._optimize_fx
|
||||
|
||||
def wrapper(gm, example_inputs=None, config=None):
|
||||
ret = original_func(gm, example_inputs, config)
|
||||
graph_after = copy.deepcopy(gm)
|
||||
assert_func(graph_after)
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def init_triton():
|
||||
init_device_properties_triton()
|
||||
|
||||
|
||||
class ModelQKNormRopeWithoutBias(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
eps: float = 1e-6,
|
||||
device="npu",
|
||||
):
|
||||
super().__init__()
|
||||
self.head_dim = head_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.q_size = num_heads * head_dim
|
||||
self.kv_size = num_kv_heads * head_dim
|
||||
self.eps = eps
|
||||
|
||||
# RMSNorm weight per head (shared across heads of same type)
|
||||
self.q_weight = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
self.k_weight = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
|
||||
def forward(self, qkv, cos_sin_cache, positions):
|
||||
"""
|
||||
Args:
|
||||
qkv: [T, q_size + 2*kv_size]
|
||||
cos: [1, T, 1, head_dim]
|
||||
sin: [1, T, 1, head_dim]
|
||||
Returns:
|
||||
q_rope, k_rope, v
|
||||
"""
|
||||
# Split QKV
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
|
||||
# Q RMSNorm (per-head)
|
||||
q_by_head = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
|
||||
q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, self.q_weight, self.eps)
|
||||
|
||||
# K RMSNorm (per-head)
|
||||
k_by_head = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim)
|
||||
k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, self.k_weight, self.eps)
|
||||
|
||||
# Reshape for RoPE: [T, num_heads, head_dim] -> [1, T, num_heads, head_dim]
|
||||
q_flat = q_norm_out.view(q.shape)
|
||||
k_flat = k_norm_out.view(k.shape)
|
||||
|
||||
# Apply RoPE
|
||||
q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding(
|
||||
positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.head_dim, True
|
||||
)
|
||||
|
||||
return q_rope, k_rope, v
|
||||
|
||||
|
||||
class ModelQKNormRopeWithBias(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
num_heads: int,
|
||||
num_kv_heads: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
eps: float = 1e-6,
|
||||
device="npu",
|
||||
):
|
||||
super().__init__()
|
||||
self.head_dim = head_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.q_size = num_heads * head_dim
|
||||
self.kv_size = num_kv_heads * head_dim
|
||||
self.eps = eps
|
||||
|
||||
self.q_weight = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
self.k_weight = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
self.q_bias = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
self.k_bias = nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device))
|
||||
|
||||
def forward(self, qkv, cos_sin_cache, positions):
|
||||
# Split QKV
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
|
||||
# Q RMSNorm + Bias
|
||||
q_by_head = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
|
||||
q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, self.q_weight, self.eps)
|
||||
q_normed = q_norm_out + self.q_bias
|
||||
|
||||
# K RMSNorm + Bias
|
||||
k_by_head = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim)
|
||||
k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, self.k_weight, self.eps)
|
||||
k_normed = k_norm_out + self.k_bias
|
||||
|
||||
# Reshape for RoPE
|
||||
q_flat = q_normed.view(q.shape)
|
||||
k_flat = k_normed.view(k.shape)
|
||||
|
||||
# Apply RoPE
|
||||
q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding(
|
||||
positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.head_dim, True
|
||||
)
|
||||
|
||||
return q_rope, k_rope, v
|
||||
|
||||
|
||||
def assert_qknorm_rope_fusion(after_gm, expect_fused=True, use_bias=False):
|
||||
check_rules = [
|
||||
(torch.ops.vllm.qkv_rmsnorm_rope.default, expect_fused),
|
||||
(torch.ops.npu.npu_rms_norm.default, not expect_fused),
|
||||
(torch.ops.vllm.npu_rotary_embedding.default, not expect_fused),
|
||||
]
|
||||
if use_bias:
|
||||
check_rules.append((torch.ops.aten.add.Tensor, not expect_fused))
|
||||
for torch_op, expect_exist in check_rules:
|
||||
found = find_op(after_gm, torch_op)
|
||||
if expect_exist:
|
||||
assert found, f"Expected operator '{torch_op}' but not find"
|
||||
else:
|
||||
assert not found, f"Not expected operator '{torch_op}' but find"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("hidden_size", [64])
|
||||
@pytest.mark.parametrize("num_tokens", [257])
|
||||
@pytest.mark.parametrize("eps", [1e-5])
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
def test_rmsnorm_quant_fusion(
|
||||
dtype: torch.dtype,
|
||||
hidden_size: int,
|
||||
num_tokens: int,
|
||||
eps: float,
|
||||
use_bias: bool,
|
||||
):
|
||||
vllm_config = VllmConfig(model_config=ModelConfig(dtype=dtype))
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": "0",
|
||||
"LOCAL_RANK": "0",
|
||||
"WORLD_SIZE": "1",
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
init_distributed_environment()
|
||||
ensure_model_parallel_initialized(1, 1)
|
||||
num_heads = 16
|
||||
num_kv_heads = 8
|
||||
head_dim = 128
|
||||
with vllm.config.set_current_vllm_config(vllm_config), set_ascend_forward_context(None, vllm_config):
|
||||
fusion_pattern = None
|
||||
q_size = num_heads * head_dim
|
||||
kv_size = num_kv_heads * head_dim
|
||||
qkv_size = q_size + 2 * kv_size
|
||||
if use_bias:
|
||||
model = ModelQKNormRopeWithBias(head_dim, num_heads, num_kv_heads, dtype, eps, device="npu")
|
||||
fusion_pattern = QKNormRopeFusionPatternWithBias(
|
||||
vllm_config=vllm_config, head_dim=head_dim, num_heads=num_heads, num_kv_heads=num_kv_heads, eps=eps
|
||||
)
|
||||
else:
|
||||
model = ModelQKNormRopeWithoutBias(head_dim, num_heads, num_kv_heads, dtype, eps, device="npu")
|
||||
fusion_pattern = QKNormRopeFusionPattern(
|
||||
vllm_config=vllm_config, head_dim=head_dim, num_heads=num_heads, num_kv_heads=num_kv_heads, eps=eps
|
||||
)
|
||||
from torch._inductor.pattern_matcher import PatternMatcherPass
|
||||
|
||||
pm_pass = PatternMatcherPass()
|
||||
fusion_pattern.register(pm_pass)
|
||||
model = model.to("npu")
|
||||
seq_len = num_tokens
|
||||
qkv = torch.randn(seq_len, qkv_size, device="npu", dtype=dtype)
|
||||
cos_sin_cache = torch.from_numpy(np.random.uniform(0, 1, [MAX_POSITION_EMBEDDING, head_dim])).to(dtype).npu()
|
||||
positions = torch.randint(
|
||||
low=0, high=MAX_POSITION_EMBEDDING, size=(num_tokens,), dtype=torch.int64, device="npu"
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
original_optimize = nge.npu_fx_compiler._optimize_fx
|
||||
nge.npu_fx_compiler._optimize_fx = create_pattern_wrapper(
|
||||
lambda gm: assert_qknorm_rope_fusion(gm, expect_fused=True, use_bias=use_bias)
|
||||
)
|
||||
|
||||
compiled_model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=True)
|
||||
|
||||
compiled_model(qkv, cos_sin_cache, positions)
|
||||
|
||||
nge.npu_fx_compiler._optimize_fx = original_optimize
|
||||
@@ -0,0 +1,298 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import vllm.config
|
||||
from vllm.compilation.passes.fx_utils import OpOverload
|
||||
from vllm.config import ModelConfig, VllmConfig
|
||||
from vllm.distributed import ensure_model_parallel_initialized, init_distributed_environment
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
import vllm_ascend.ops.register_custom_ops # noqa
|
||||
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
||||
from vllm_ascend.compilation.passes.norm_quant_fusion_pass import AddRMSNormQuantFusionPass
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
from .backend import TestBackend
|
||||
|
||||
# Cache backend to avoid duplicate pattern registration
|
||||
_backend_cache = None
|
||||
|
||||
|
||||
def get_or_create_backend(vllm_config):
|
||||
"""Get or create backend with fusion passes (cached to avoid duplicate pattern registration)."""
|
||||
global _backend_cache
|
||||
if _backend_cache is None:
|
||||
_backend_cache = TestBackend(custom_passes=[AddRMSNormQuantFusionPass(vllm_config=vllm_config)])
|
||||
return _backend_cache
|
||||
|
||||
|
||||
class TestModelWithoutBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.dtype, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch.ops._C_ascend.npu_add_rms_norm_bias(
|
||||
x, residual, self.rms_norm_weight, None, self.eps
|
||||
)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
def ops_in_model_before(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators BEFORE fusion."""
|
||||
return [torch.ops._C_ascend.npu_add_rms_norm_bias.default, torch.ops.vllm.quantize.default]
|
||||
|
||||
def ops_in_model_after(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators AFTER successful fusion."""
|
||||
return [torch.ops.npu.npu_add_rms_norm_quant.default]
|
||||
|
||||
|
||||
class TestModelWithBias(nn.Module):
|
||||
"""
|
||||
A test model that simulates the pattern:
|
||||
AddRMSNorm → Add Bias → Quantization (with bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.dtype, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.bias = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Add bias
|
||||
3. Quantize to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output_with_bias, _, new_residual = torch.ops._C_ascend.npu_add_rms_norm_bias(
|
||||
x, residual, self.rms_norm_weight, self.bias, self.eps
|
||||
)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output_with_bias, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
def ops_in_model_before(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators BEFORE fusion."""
|
||||
return [torch.ops._C_ascend.npu_add_rms_norm_bias.default, torch.ops.vllm.quantize.default]
|
||||
|
||||
def ops_in_model_after(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators AFTER successful fusion."""
|
||||
return [torch.ops.npu.npu_add_rms_norm_quant.default]
|
||||
|
||||
|
||||
class TestModelSPWithoutBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → maybe_allgather → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.dtype, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Perform a fake maybe_all_gather_and_maybe_unpad
|
||||
3. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output, _, new_residual = torch.ops._C_ascend.npu_add_rms_norm_bias(
|
||||
x, residual, self.rms_norm_weight, None, self.eps
|
||||
)
|
||||
|
||||
norm_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(norm_output, True)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
def ops_in_model_before(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators BEFORE fusion."""
|
||||
return [
|
||||
torch.ops._C_ascend.npu_add_rms_norm_bias.default,
|
||||
torch.ops.vllm.maybe_all_gather_and_maybe_unpad.default,
|
||||
torch.ops.vllm.quantize.default,
|
||||
]
|
||||
|
||||
def ops_in_model_after(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators AFTER successful fusion."""
|
||||
return [torch.ops.npu.npu_add_rms_norm_quant.default, torch.ops.vllm.maybe_all_gather_and_maybe_unpad.default]
|
||||
|
||||
|
||||
class TestModelSPWithBias(nn.Module):
|
||||
"""
|
||||
A minimal test model that simulates the pattern:
|
||||
AddRMSNorm → Add bias → maybe_allgather → Quantization (without bias)
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, dtype: torch.dtype, eps: float = 1e-6, device="npu"):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.rms_norm_weight = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.bias = nn.Parameter(torch.randn(hidden_size, device=device))
|
||||
self.quant_scale = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_scale_reciprocal = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
self.quant_offset = torch.zeros(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass:
|
||||
1. Perform npu_add_rms_norm
|
||||
2. Add bias
|
||||
3. Perform a fake maybe_all_gather_and_maybe_unpad
|
||||
4. Quantize the normalized output to int8
|
||||
Returns both quantized output and updated residual.
|
||||
"""
|
||||
residual = torch.zeros_like(x)
|
||||
|
||||
norm_output_with_bias, _, new_residual = torch.ops._C_ascend.npu_add_rms_norm_bias(
|
||||
x, residual, self.rms_norm_weight, self.bias, self.eps
|
||||
)
|
||||
|
||||
norm_output_with_bias = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(norm_output_with_bias, True)
|
||||
|
||||
quantized_output = torch.ops.vllm.quantize(
|
||||
norm_output_with_bias, self.quant_scale, self.quant_scale_reciprocal, self.quant_offset
|
||||
)
|
||||
|
||||
return quantized_output, new_residual
|
||||
|
||||
def ops_in_model_before(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators BEFORE fusion."""
|
||||
return [
|
||||
torch.ops._C_ascend.npu_add_rms_norm_bias.default,
|
||||
torch.ops.vllm.maybe_all_gather_and_maybe_unpad.default,
|
||||
torch.ops.vllm.quantize.default,
|
||||
]
|
||||
|
||||
def ops_in_model_after(self) -> list[OpOverload]:
|
||||
"""Return the list of expected operators AFTER successful fusion."""
|
||||
return [torch.ops.npu.npu_add_rms_norm_quant.default, torch.ops.vllm.maybe_all_gather_and_maybe_unpad.default]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("hidden_size", [64])
|
||||
@pytest.mark.parametrize("num_tokens", [257])
|
||||
@pytest.mark.parametrize("eps", [1e-5, 1e-6])
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
@pytest.mark.parametrize("sp_enable", [False, True])
|
||||
def test_rmsnorm_quant_fusion(
|
||||
dtype: torch.dtype,
|
||||
hidden_size: int,
|
||||
num_tokens: int,
|
||||
eps: float,
|
||||
use_bias: bool,
|
||||
sp_enable: bool,
|
||||
):
|
||||
"""
|
||||
End-to-end test for AddRMSNorm+Quantize fusion.
|
||||
Compares: Operator presence/absence before and after graph transformation
|
||||
"""
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(1)
|
||||
|
||||
vllm_config = VllmConfig(model_config=ModelConfig(dtype=dtype))
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config):
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": "0",
|
||||
"LOCAL_RANK": "0",
|
||||
"WORLD_SIZE": "1",
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
init_distributed_environment()
|
||||
ensure_model_parallel_initialized(1, 1)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config), set_ascend_forward_context(None, vllm_config):
|
||||
backend = get_or_create_backend(vllm_config)
|
||||
if use_bias:
|
||||
if not enable_custom_op():
|
||||
return
|
||||
if sp_enable:
|
||||
model = TestModelSPWithBias(hidden_size, dtype, eps, device="npu")
|
||||
else:
|
||||
model = TestModelWithBias(hidden_size, dtype, eps, device="npu")
|
||||
else:
|
||||
if sp_enable:
|
||||
model = TestModelSPWithoutBias(hidden_size, dtype, eps, device="npu")
|
||||
else:
|
||||
model = TestModelWithoutBias(hidden_size, dtype, eps, device="npu")
|
||||
model = model.to("npu")
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, device="npu", dtype=dtype, requires_grad=False)
|
||||
|
||||
result_unfused = model(x)
|
||||
print("Unfused result:", [t.shape for t in result_unfused])
|
||||
model_fused = torch.compile(model, backend=backend)
|
||||
result_fused = model_fused(x)
|
||||
print("Fused result:", [t.shape for t in result_fused])
|
||||
|
||||
print("=== Checking operator fusion ===")
|
||||
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=not sp_enable)
|
||||
backend.check_after_ops(model.ops_in_model_after())
|
||||
Reference in New Issue
Block a user