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, 11:00 UTC
[AMD][torch.compile] Enable silu+fp8_quant fusion for rocm (#18082)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
May 14, 2025 at 00:13 UTC
7b2f28deba3ce0ad773611f1612f9fc092b0e923
6 files changed
+14
-9
.buildkite/test-pipeline.yaml
+1
@@ -309,6 +309,7 @@ steps:
309
commands:
310
- pytest -v -s compile/test_pass_manager.py
311
- pytest -v -s compile/test_fusion.py
312
+ - pytest -v -s compile/test_silu_mul_quant_fusion.py
313
- pytest -v -s compile/test_sequence_parallelism.py
314
315
- label: PyTorch Fullgraph Smoke Test # 9min
csrc/quantization/activation_kernels.cu
+2
-1
@@ -112,7 +112,8 @@ __global__ void act_and_mul_quant_kernel(
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);
115
+ TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
116
+ out.dtype() == torch::kFloat8_e4m3fnuz);
117
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
118
input.dtype() == torch::kBFloat16);
119
TORCH_CHECK(input.size(-1) % 2 == 0);
tests/compile/test_silu_mul_quant_fusion.py
+3
-3
@@ -27,8 +27,8 @@ class TestModel(torch.nn.Module):
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")
30
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
31
+ reason="Only test on CUDA and ROCm")
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)
@@ -36,7 +36,7 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):
36
# Reshape pass is needed for the fusion pass to work
37
config = VllmConfig()
38
config.compilation_config = CompilationConfig(
39
- pass_config=PassConfig(enable_fusion=True, enable_reshape=True))
39
+ pass_config=PassConfig(enable_fusion=True, enable_noop=True))
40
fusion_pass = ActivationQuantFusionPass(config)
41
42
backend = TestBackend(fusion_pass)
tests/kernels/quantization/test_rocm_skinny_gemms.py
+3
-2
@@ -58,8 +58,9 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
58
@pytest.mark.parametrize("m", M + [28672]) # m >= 16
59
@pytest.mark.parametrize("dtype", DTYPES)
60
@pytest.mark.parametrize("seed", SEEDS)
61
-@pytest.mark.skipif(not current_platform.is_rocm(),
62
- reason="only test for rocm")
61
+@pytest.mark.skipif(
62
+ not (current_platform.is_rocm() and current_platform.supports_fp8()),
63
+ reason="only test for rocm fp8")
64
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
65
torch.manual_seed(seed)
66
tests/kernels/test_fused_quant_activation.py
+3
-2
@@ -5,9 +5,10 @@ import torch
5
import vllm._custom_ops as ops
6
from tests.kernels.utils import opcheck
7
from vllm.model_executor.layers.activation import SiluAndMul
8
+from vllm.platforms import current_platform
9
10
DTYPES = [torch.bfloat16, torch.float16]
10
-QUANT_DTYPES = [torch.float8_e4m3fn]
11
+QUANT_DTYPES = [current_platform.fp8_dtype()]
12
NUM_TOKENS = [1, 17, 86, 1234, 3045] # Arbitrary values for testing
13
HIDDEN_SIZES = [16, 48, 128, 1562, 4096] # Arbitrary values for testing
14
SEEDS = [0]
@@ -26,7 +27,7 @@ def ref_impl(silu_and_mul: SiluAndMul, x: torch.Tensor,
27
def ops_impl(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
28
out_shape = (x.shape[0], x.shape[1] // 2)
29
out = torch.empty(out_shape,
29
- dtype=torch.torch.float8_e4m3fn,
30
+ dtype=current_platform.fp8_dtype(),
31
device=x.device)
32
torch.ops._C.silu_and_mul_quant(out, x, scale)
33
return out
vllm/compilation/activation_quant_fusion.py
+2
-1
@@ -7,6 +7,7 @@ from torch._inductor.pattern_matcher import (PatternMatcherPass, fwd_only,
7
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
+from vllm.platforms import current_platform
11
12
from .vllm_inductor_pass import VllmInductorPass
13
@@ -41,7 +42,7 @@ def empty_bf16(*args, **kwargs):
42
43
44
def empty_fp8(*args, **kwargs):
44
- fp8 = torch.float8_e4m3fn
45
+ fp8 = current_platform.fp8_dtype()
46
return torch.empty(*args, **kwargs, dtype=fp8, device="cuda")
47
48