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, 12:15 UTC
[torch.compile] Add torch inductor pass for fusing silu_and_mul with subsequent scaled_fp8_quant operations (#10867)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
May 1, 2025 at 07:59 UTC
460a2b11004229ccba1f5d5200c541496c8cf513
11 files changed
+406
-9
CMakeLists.txt
+1
@@ -241,6 +241,7 @@ set(VLLM_EXT_SRC
241
"csrc/quantization/fp8/common.cu"
242
"csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
243
"csrc/quantization/gguf/gguf_kernel.cu"
244
+ "csrc/quantization/activation_kernels.cu"
245
"csrc/cuda_utils_kernels.cu"
246
"csrc/prepare_inputs/advance_step.cu"
247
"csrc/custom_all_reduce.cu"
csrc/core/math.hpp
+19
@@ -7,3 +7,22 @@ inline constexpr uint32_t next_pow_2(uint32_t const num) {
7
if (num <= 1) return num;
8
return 1 << (CHAR_BIT * sizeof(num) - __builtin_clz(num - 1));
9
}
10
+
11
+template <typename A, typename B>
12
+static inline constexpr auto div_ceil(A a, B b) {
13
+ return (a + b - 1) / b;
14
+}
15
+
16
+// Round a down to the next multiple of b. The caller is responsible for making
17
+// sure that b is non-zero
18
+template <typename T>
19
+inline constexpr T round_to_previous_multiple_of(T a, T b) {
20
+ return a % b == 0 ? a : (a / b) * b;
21
+}
22
+
23
+// Round a up to the next multiple of b. The caller is responsible for making
24
+// sure that b is non-zero
25
+template <typename T>
26
+inline constexpr T round_to_next_multiple_of(T a, T b) {
27
+ return a % b == 0 ? a : ((a / b) + 1) * b;
28
+}
csrc/ops.h
+3
@@ -97,6 +97,9 @@ void batched_rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
97
98
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
99
100
+void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
101
+ torch::Tensor& scale);
102
+
103
void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
104
105
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
csrc/quantization/activation_kernels.cu
new
+120
@@ -0,0 +1,120 @@
1
+#include <ATen/cuda/CUDAContext.h>
2
+#include <torch/all.h>
3
+#include <c10/cuda/CUDAGuard.h>
4
+
5
+#include <cmath>
6
+#include "core/math.hpp"
7
+#include "cuda_compat.h"
8
+#include "dispatch_utils.h"
9
+
10
+#include "quantization/fp8/common.cuh"
11
+
12
+namespace vllm {
13
+
14
+template <typename T>
15
+__device__ __forceinline__ T silu_kernel(const T& x) {
16
+ // x * sigmoid(x)
17
+ return (T)(((float)x) / (1.0f + expf((float)-x)));
18
+}
19
+
20
+// Activation and gating kernel template.
21
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
22
+ typename fp8_type>
23
+__global__ void act_and_mul_quant_kernel(
24
+ fp8_type* __restrict__ out, // [..., d]
25
+ const scalar_t* __restrict__ input, // [..., 2, d]
26
+ const float* scale, const int d) {
27
+ const int32_t blocks_per_token = gridDim.y;
28
+
29
+ const int32_t elems_per_128bit_load = (128 / 8) / sizeof(scalar_t);
30
+
31
+ // We don't expect the hidden dimension to exceed 32 bits so int32 should
32
+ // be safe here.
33
+ const int32_t tgt_elems_per_block = div_ceil(d, blocks_per_token);
34
+ const int32_t elems_per_block =
35
+ round_to_next_multiple_of(tgt_elems_per_block, elems_per_128bit_load);
36
+ const int32_t block_start = blockIdx.y * elems_per_block;
37
+ int32_t block_end = block_start + elems_per_block;
38
+ block_end = block_end > d ? d : block_end;
39
+
40
+ // token_idx is 64 bit to prevent 32 bit overflow when the number of tokens
41
+ // is very large
42
+ const int64_t token_idx = blockIdx.x;
43
+ const scalar_t* __restrict__ x_ptr = input + token_idx * 2 * d;
44
+ const scalar_t* __restrict__ y_ptr = input + token_idx * 2 * d + d;
45
+ fp8_type* __restrict__ out_ptr = out + token_idx * d;
46
+
47
+ // 128-bit vectorized code
48
+ const int32_t vec_loop_end =
49
+ round_to_previous_multiple_of(elems_per_128bit_load, block_end);
50
+ const int32_t vec_end_idx = vec_loop_end / elems_per_128bit_load;
51
+ const int32_t vec_start_idx = block_start / elems_per_128bit_load;
52
+
53
+ const int4* __restrict__ x_128bit_ptr = reinterpret_cast<const int4*>(x_ptr);
54
+ const int4* __restrict__ y_128bit_ptr = reinterpret_cast<const int4*>(y_ptr);
55
+ int2* __restrict__ out_128bit_ptr = reinterpret_cast<int2*>(out_ptr);
56
+
57
+ float inverted_scale = 1 / *scale;
58
+#pragma unroll
59
+ for (int32_t vec_idx = vec_start_idx + threadIdx.x; vec_idx < vec_end_idx;
60
+ vec_idx += blockDim.x) {
61
+ const int4 x_128bit = VLLM_LDG(&x_128bit_ptr[vec_idx]);
62
+ const int4 y_128bit = VLLM_LDG(&y_128bit_ptr[vec_idx]);
63
+ using scalar_128bit_vec_t = std::array<scalar_t, elems_per_128bit_load>;
64
+ using scalar_64bit_vec_t = std::array<fp8_type, elems_per_128bit_load>;
65
+
66
+ scalar_64bit_vec_t out_vec;
67
+ const auto x_vec = reinterpret_cast<scalar_128bit_vec_t const&>(x_128bit);
68
+ const auto y_vec = reinterpret_cast<scalar_128bit_vec_t const&>(y_128bit);
69
+
70
+#pragma unroll
71
+ for (int i = 0; i < elems_per_128bit_load; i++) {
72
+ out_vec[i] = scaled_fp8_conversion<true, fp8_type>(
73
+ ACT_FN(x_vec[i]) * y_vec[i], inverted_scale);
74
+ }
75
+
76
+ out_128bit_ptr[vec_idx] = reinterpret_cast<const int2&>(out_vec);
77
+ }
78
+
79
+ // Scalar cleanup code
80
+ if (block_end > vec_loop_end) {
81
+ for (int64_t idx = vec_loop_end + threadIdx.x; idx < block_end;
82
+ idx += blockDim.x) {
83
+ const scalar_t x = VLLM_LDG(&x_ptr[idx]);
84
+ const scalar_t y = VLLM_LDG(&y_ptr[idx]);
85
+ out_ptr[idx] =
86
+ scaled_fp8_conversion<true, fp8_type>(ACT_FN(x) * y, inverted_scale);
87
+ }
88
+ }
89
+}
90
+} // namespace vllm
91
+
92
+// Launch activation, gating, and quantize kernel.
93
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
94
+ int d = input.size(-1) / 2; \
95
+ int64_t num_tokens = input.numel() / input.size(-1); \
96
+ dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
97
+ dim3 block(std::min(d, 512)); \
98
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
99
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
100
+ VLLM_DISPATCH_FLOATING_TYPES( \
101
+ input.scalar_type(), "act_and_mul_kernel", [&] { \
102
+ VLLM_DISPATCH_FP8_TYPES( \
103
+ out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
104
+ vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
105
+ fp8_t> \
106
+ <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
107
+ input.data_ptr<scalar_t>(), \
108
+ scale.data_ptr<float>(), d); \
109
+ }); \
110
+ });
111
+
112
+void silu_and_mul_quant(torch::Tensor& out, // [..., d]
113
+ torch::Tensor& input, // [..., 2 * d]
114
+ torch::Tensor& scale) {
115
+ TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn);
116
+ TORCH_CHECK(input.dtype() == torch::kFloat16 ||
117
+ input.dtype() == torch::kBFloat16);
118
+ TORCH_CHECK(input.size(-1) % 2 == 0);
119
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
120
+}
csrc/torch_bindings.cpp
+5
-1
@@ -81,9 +81,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
81
82
// Activation ops
83
// Activation function used in SwiGLU.
84
- ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
84
+ ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
85
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
86
87
+ ops.def(
88
+ "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
89
+ ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
90
+
91
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
92
ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
93
tests/compile/test_functionalization.py
+14
-3
@@ -5,6 +5,7 @@ import torch
5
6
import vllm.envs as envs
7
from vllm import LLM, SamplingParams
8
+from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
9
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
10
from vllm.compilation.fusion import (FUSED_OPS, FusionPass, QuantKey,
11
kFp8DynamicTokenSym, kFp8StaticTensorSym)
@@ -17,7 +18,6 @@ from .backend import TestBackend
18
OPS_IN_MODEL = [
19
torch.ops._C.rotary_embedding.default,
20
torch.ops._C.fused_add_rms_norm.default,
20
- torch.ops._C.silu_and_mul.default,
21
]
22
23
RMS_OP = torch.ops._C.rms_norm.default
@@ -29,6 +29,9 @@ RMS_QUANT_OPS = {
29
],
30
}
31
32
+SILU_MUL_OP = torch.ops._C.silu_and_mul.default
33
+
34
+SILU_MUL_QUANT_OP = torch.ops._C.silu_and_mul_quant.default
35
prompts = [
36
"Hello, my name is",
37
"The president of the United States is",
@@ -55,8 +58,10 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
58
enable_noop=True))
59
noop_pass = NoOpEliminationPass(vllm_config)
60
fusion_pass = FusionPass.instance(vllm_config)
61
+ act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)
62
59
- passes = [noop_pass, fusion_pass] if do_fusion else [noop_pass]
63
+ passes = [noop_pass, fusion_pass, act_quant_fusion_pass
64
+ ] if do_fusion else [noop_pass]
65
func_pass = FixFunctionalizationPass(vllm_config)
66
backend_func = TestBackend(*passes, func_pass)
67
backend_no_func = TestBackend(*passes)
@@ -79,6 +84,7 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
84
model_runner.model = torch.compile(orig_model,
85
fullgraph=True,
86
backend=backend_no_func)
87
+
88
gen_no_func = llm.generate(prompts, sampling_params)
89
90
for output_func, output_no_func in zip(gen_func, gen_no_func):
@@ -88,7 +94,12 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
94
# and replaced by fused quantized ops in RMS_QUANT_OPS.
95
rms_ops = [FUSED_OPS[(quant_key, True)], FUSED_OPS[(quant_key, False)]
96
] if do_fusion else [RMS_OP]
91
- ops = OPS_IN_MODEL + rms_ops
97
+ silu_mul_ops = [SILU_MUL_QUANT_OP] if do_fusion and \
98
+ quant_key == kFp8StaticTensorSym else [
99
+ SILU_MUL_OP
100
+ ]
101
+
102
+ ops = OPS_IN_MODEL + rms_ops + silu_mul_ops
103
104
for op in ops:
105
find_auto_fn(backend_no_func.graph_post_pass.nodes, op)
tests/compile/test_silu_mul_quant_fusion.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+import torch
4
+
5
+import vllm.envs as envs
6
+from vllm._custom_ops import scaled_fp8_quant
7
+from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
8
+from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe
9
+from vllm.config import CompilationConfig, VllmConfig
10
+from vllm.model_executor.layers.activation import SiluAndMul
11
+
12
+from .backend import TestBackend
13
+
14
+
15
+class TestModel(torch.nn.Module):
16
+
17
+ def __init__(self, *args, **kwargs):
18
+ super().__init__(*args, **kwargs)
19
+ self.silu_and_mul = SiluAndMul()
20
+ self.scale = torch.rand(1, dtype=torch.float32)
21
+
22
+ def forward(self, x):
23
+ y = self.silu_and_mul(x)
24
+ x2 = scaled_fp8_quant(y, self.scale)
25
+ return x2
26
+
27
+
28
+@pytest.mark.parametrize("num_tokens", [256])
29
+@pytest.mark.parametrize("hidden_size", [64])
30
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda",
31
+ reason="Only test on CUDA")
32
+def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):
33
+ torch.set_default_device("cuda")
34
+ torch.set_default_dtype(torch.float16)
35
+
36
+ # Reshape pass is needed for the fusion pass to work
37
+ config = VllmConfig()
38
+ config.compilation_config = CompilationConfig(
39
+ pass_config=CompilationConfig.PassConfig(enable_fusion=True,
40
+ enable_reshape=True))
41
+ fusion_pass = ActivationQuantFusionPass(config)
42
+
43
+ backend = TestBackend(fusion_pass)
44
+ model = TestModel()
45
+
46
+ # First dimension dynamic
47
+ x = torch.rand(num_tokens, hidden_size)
48
+ torch._dynamo.mark_dynamic(x, 0)
49
+
50
+ result = model(x)
51
+
52
+ model2 = torch.compile(model, backend=backend)
53
+ result2 = model2(x)
54
+
55
+ # Check that it gives the same answer
56
+ torch.testing.assert_close(result[0].to(dtype=torch.float16),
57
+ result2[0].to(dtype=torch.float16),
58
+ atol=1e-3,
59
+ rtol=1e-3)
60
+
61
+ # Check substitution worked
62
+ pre_nodes = backend.graph_pre_pass.nodes
63
+ post_nodes = backend.graph_post_pass.nodes
64
+
65
+ silu_and_mul_quant = torch.ops._C.silu_and_mul_quant.default
66
+ fp8_quant = torch.ops._C.static_scaled_fp8_quant.default
67
+
68
+ # In pre-nodes, fp8 quant should be present and fused kernels should not
69
+ assert find_auto_fn_maybe(pre_nodes, silu_and_mul_quant) is None
70
+ find_auto_fn(pre_nodes, fp8_quant)
71
+
72
+ # In post-nodes, fused kernels should be present and fp8 quant should not
73
+ find_auto_fn(post_nodes, silu_and_mul_quant)
74
+ assert find_auto_fn_maybe(post_nodes, fp8_quant) is None
tests/kernels/test_fused_quant_activation.py
new
+69
@@ -0,0 +1,69 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+import torch
4
+
5
+import vllm._custom_ops as ops
6
+from tests.kernels.utils import opcheck
7
+from vllm.model_executor.layers.activation import SiluAndMul
8
+
9
+DTYPES = [torch.bfloat16, torch.float16]
10
+QUANT_DTYPES = [torch.float8_e4m3fn]
11
+NUM_TOKENS = [1, 17, 86, 1234, 3045] # Arbitrary values for testing
12
+HIDDEN_SIZES = [16, 48, 128, 1562, 4096] # Arbitrary values for testing
13
+SEEDS = [0]
14
+CUDA_DEVICES = [
15
+ f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
16
+]
17
+
18
+
19
+def ref_impl(silu_and_mul: SiluAndMul, x: torch.Tensor,
20
+ scale: torch.Tensor) -> torch.Tensor:
21
+ silu_and_mul_out = silu_and_mul.forward_native(x)
22
+ out, scales = ops.scaled_fp8_quant(silu_and_mul_out, scale)
23
+ return out
24
+
25
+
26
+def ops_impl(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
27
+ out_shape = (x.shape[0], x.shape[1] // 2)
28
+ out = torch.empty(out_shape,
29
+ dtype=torch.torch.float8_e4m3fn,
30
+ device=x.device)
31
+ torch.ops._C.silu_and_mul_quant(out, x, scale)
32
+ return out
33
+
34
+
35
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
36
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
37
+@pytest.mark.parametrize("dtype", DTYPES)
38
+@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
39
+@pytest.mark.parametrize("seed", SEEDS)
40
+@pytest.mark.parametrize("device", CUDA_DEVICES)
41
+@torch.inference_mode()
42
+def test_silu_and_mul(
43
+ num_tokens: int,
44
+ hidden_size: int,
45
+ dtype: torch.dtype,
46
+ quant_dtype: torch.dtype,
47
+ seed: int,
48
+ device: str,
49
+) -> None:
50
+ torch.random.manual_seed(seed)
51
+ if torch.cuda.is_available():
52
+ torch.cuda.manual_seed(seed)
53
+ torch.set_default_device(device)
54
+
55
+ layer = SiluAndMul()
56
+
57
+ # Make inputs
58
+ scale = (torch.randn((1), device=device, dtype=torch.float32))
59
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype)
60
+
61
+ ref_out = ref_impl(layer, x, scale)
62
+ ops_out = ops_impl(x, scale)
63
+
64
+ assert ref_out.dtype == quant_dtype
65
+ assert ops_out.dtype == quant_dtype
66
+ assert ref_out.shape == ops_out.shape
67
+ assert torch.allclose(ref_out.to(dtype=torch.float32),
68
+ ops_out.to(dtype=torch.float32))
69
+ opcheck(torch.ops._C.silu_and_mul_quant, (ops_out, x, scale))
vllm/compilation/activation_quant_fusion.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import torch
4
+from torch._higher_order_ops.auto_functionalize import auto_functionalized
5
+from torch._inductor.pattern_matcher import (PatternMatcherPass, fwd_only,
6
+ register_replacement)
7
+
8
+from vllm.config import VllmConfig
9
+from vllm.logger import init_logger
10
+
11
+from .vllm_inductor_pass import VllmInductorPass
12
+
13
+logger = init_logger(__name__)
14
+
15
+
16
+def silu_mul_pattern_static(result: torch.Tensor,
17
+ result_silu_mul: torch.Tensor, input: torch.Tensor,
18
+ scale: torch.Tensor):
19
+ at1 = auto_functionalized(torch.ops._C.silu_and_mul.default,
20
+ result=result_silu_mul,
21
+ input=input)
22
+ at2 = auto_functionalized(torch.ops._C.static_scaled_fp8_quant.default,
23
+ result=result,
24
+ input=at1[1],
25
+ scale=scale)
26
+ return at2[1]
27
+
28
+
29
+def silu_mul_replacement_static(result: torch.Tensor,
30
+ result_silu_mul: torch.Tensor,
31
+ input: torch.Tensor, scale: torch.Tensor):
32
+ at = auto_functionalized(torch.ops._C.silu_and_mul_quant.default,
33
+ result=result,
34
+ input=input,
35
+ scale=scale)
36
+ return at[1]
37
+
38
+
39
+def empty_bf16(*args, **kwargs):
40
+ return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
41
+
42
+
43
+def empty_fp8(*args, **kwargs):
44
+ fp8 = torch.float8_e4m3fn
45
+ return torch.empty(*args, **kwargs, dtype=fp8, device="cuda")
46
+
47
+
48
+def empty_fp32(*args, **kwargs):
49
+ return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
50
+
51
+
52
+class ActivationQuantFusionPass(VllmInductorPass):
53
+ """
54
+ This pass fuses a pre-defined set of custom ops into fused ops.
55
+ It uses the torch pattern matcher to find the patterns and replace them.
56
+
57
+ Because patterns can only be registered once, the pass is a singleton.
58
+ This will be addressed in a future version of PyTorch:
59
+ https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
60
+ """
61
+
62
+ def __init__(self, config: VllmConfig):
63
+ super().__init__(config)
64
+
65
+ self.patterns: PatternMatcherPass = PatternMatcherPass(
66
+ pass_name="activation_quant_fusion_pass")
67
+
68
+ inputs = [
69
+ empty_fp8(5, 4), # Quant output
70
+ empty_bf16(5, 4), # Silu_and_mul output
71
+ empty_bf16(5, 4), # Input
72
+ empty_fp32(1, 1) # Scale
73
+ ]
74
+ register_replacement(silu_mul_pattern_static,
75
+ silu_mul_replacement_static, inputs, fwd_only,
76
+ self.patterns)
77
+
78
+ def __call__(self, graph: torch.fx.Graph):
79
+ self.begin()
80
+ self.dump_graph(graph, "before_act_quant_fusion")
81
+
82
+ count = self.patterns.apply(graph)
83
+ logger.debug("Replaced %s patterns in ActivationQuantFusionPass",
84
+ count)
85
+
86
+ self.dump_graph(graph, "after_act_quant_fusion")
87
+ self.end_and_log()
vllm/compilation/fix_functionalization.py
+12
-5
@@ -68,18 +68,25 @@ class FixFunctionalizationPass(VllmInductorPass):
68
self.defunctionalize(graph, node, mutated_args)
69
elif at_target in [
70
torch.ops._C.rms_norm.default,
71
- torch.ops._C.rms_norm_static_fp8_quant.default
71
+ torch.ops._C.rms_norm_static_fp8_quant.default,
72
]:
73
mutated_args = {1: 'result'}
74
self.defunctionalize(graph, node, mutated_args)
75
-
75
+ # For some reason we need to specify the args for both
76
+ # silu_and_mul and silu_and_mul_quant. The kwargs
77
+ # pathway gets the wrong answer.
78
elif at_target == torch.ops._C.silu_and_mul.default:
77
- mutated_args = {1: 'out'}
78
- # Because we have an 'out', need to specify args directly
79
+ mutated_args = {1: 'result'}
80
+ self.defunctionalize(graph,
81
+ node,
82
+ mutated_args,
83
+ args=('result', 'input'))
84
+ elif at_target == torch.ops._C.silu_and_mul_quant.default:
85
+ mutated_args = {1: 'result'}
86
self.defunctionalize(graph,
87
node,
88
mutated_args,
82
- args=('out', 'input'))
89
+ args=('result', 'input', 'scale'))
90
else:
91
continue # skip the count
92
vllm/compilation/pass_manager.py
+2
@@ -7,6 +7,7 @@ from torch import fx as fx
7
from vllm.config import VllmConfig
8
from vllm.logger import init_logger
9
10
+from .activation_quant_fusion import ActivationQuantFusionPass
11
from .fix_functionalization import FixFunctionalizationPass
12
from .fusion import FusionPass
13
from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
@@ -51,6 +52,7 @@ class PostGradPassManager(CustomGraphPass):
52
53
if self.pass_config.enable_fusion:
54
self.passes += [FusionPass.instance(config)]
55
+ self.passes += [ActivationQuantFusionPass(config)]
56
57
if self.pass_config.enable_sequence_parallelism:
58
self.passes += [SequenceParallelismPass(config)]