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

[torch.compile][ROCm] Fuse quantization onto attention using a torch.compile pass (#16756)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Sage Moore <sage@neuralmagic.com>

Luka Govedič committed Jun 12, 2025 at 11:31 UTC f98548b9da39bce8aaff64974e1e68208aeada5e
33 files changed +622 -79
.buildkite/test-pipeline.yaml
+1
@@ -305,6 +305,7 @@ steps:
305 commands:
306 - pytest -v -s compile/test_pass_manager.py
307 - pytest -v -s compile/test_fusion.py
308 + - pytest -v -s compile/test_fusion_attn.py
309 - pytest -v -s compile/test_silu_mul_quant_fusion.py
310 - pytest -v -s compile/test_sequence_parallelism.py
311 - pytest -v -s compile/test_async_tp.py
tests/compile/backend.py
+16 -14
@@ -1,13 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Sequence
5 from copy import deepcopy
6 from typing import Callable, Union
7
8 from torch import fx
9 +from torch._ops import OpOverload
10
9 -from vllm.compilation.fx_utils import (find_specified_fn,
10 - find_specified_fn_maybe)
11 +from vllm.compilation.fx_utils import find_op_nodes
12 from vllm.compilation.inductor_pass import InductorPass
13 from vllm.config import get_current_vllm_config
14
@@ -48,18 +49,19 @@ class TestBackend:
49 # assign by reference, will reflect the final state of the graph
50 self.final_graph = graph
51
51 - def check_before_ops(self, ops,
52 - find_fn=find_specified_fn, \
53 - find_fn_maybe=find_specified_fn_maybe, \
54 - ops_fully_replaced=True):
52 + def check_before_ops(self, ops: Sequence[OpOverload], fully_replaced=True):
53 for op in ops:
56 - find_fn(self.graph_pre_pass.nodes, op)
57 - if ops_fully_replaced:
58 - assert find_fn_maybe(self.graph_post_pass.nodes, op) is None
54 + num_pre = len(list(find_op_nodes(op, self.graph_pre_pass)))
55 + num_post = len(list(find_op_nodes(op, self.graph_post_pass)))
56 + assert num_pre > 0, f"Op {op.name()} not found in pre-pass graph"
57 + assert num_pre > num_post, f"All nodes remain for op {op.name()}"
58 + if fully_replaced:
59 + assert num_post == 0, \
60 + f"Unexpected op {op.name()} in post-pass graph"
61
60 - def check_after_ops(self, ops,
61 - find_fn=find_specified_fn, \
62 - find_fn_maybe=find_specified_fn_maybe):
62 + def check_after_ops(self, ops: Sequence[OpOverload]):
63 for op in ops:
64 - find_fn(self.graph_post_pass.nodes, op)
65 - assert find_fn_maybe(self.graph_pre_pass.nodes, op) is None
64 + num_pre = len(list(find_op_nodes(op, self.graph_pre_pass)))
65 + num_post = len(list(find_op_nodes(op, self.graph_post_pass)))
66 + assert num_pre == 0, f"Unexpected op {op.name()} in pre-pass graph"
67 + assert num_post > 0, f"Op {op.name()} not found in post-pass graph"
\ No newline at end of file
tests/compile/test_async_tp.py
+1 -2
@@ -169,8 +169,7 @@ def async_tp_pass_on_test_model(local_rank: int, world_size: int,
169
170 # In pre-nodes, all gather or reduce scatter should exist,
171 # fused_matmul_reduce_scatter or fused_all_gather_matmul should not
172 - backend.check_before_ops(model.ops_in_model_before(),
173 - ops_fully_replaced=False)
172 + backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
173
174 # In post-nodes, fused_matmul_reduce_scatter or \
175 # fused_all_gather_matmul should exist
tests/compile/test_fusion.py
+5 -7
@@ -7,8 +7,7 @@ import torch
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
10 + FusionPass, GroupShape, QuantKey)
11 from vllm.compilation.noop_elimination import NoOpEliminationPass
12 from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,
13 VllmConfig)
@@ -30,9 +29,10 @@ class TestModel(torch.nn.Module):
29 self.cutlass_fp8_enabled = cutlass_fp8_enabled
30 self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]
31 self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
32 + group_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN
33 self.key = QuantKey(dtype=FP8_DTYPE,
34 static=static,
35 - per_tensor=static,
35 + group_shape=group_shape,
36 symmetric=True)
37 if static:
38 self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]
@@ -122,9 +122,7 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
122 torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
123
124 # In pre-nodes, fp8 quant should be there and fused kernels should not
125 - backend.check_before_ops(model.ops_in_model_before(), find_auto_fn,
126 - find_auto_fn_maybe)
125 + backend.check_before_ops(model.ops_in_model_before())
126
127 # In post-nodes, fused kernels should be there and fp8 quant should not
129 - backend.check_after_ops(model.ops_in_model_after(), find_auto_fn,
130 - find_auto_fn_maybe)
128 + backend.check_after_ops(model.ops_in_model_after())
tests/compile/test_fusion_attn.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3 +
4 +import pytest
5 +import torch._dynamo
6 +
7 +from tests.compile.backend import TestBackend
8 +from tests.models.utils import check_outputs_equal
9 +from vllm import LLM, SamplingParams
10 +from vllm.compilation.fusion import QUANT_OPS, QuantKey, kFp8StaticTensorSym
11 +from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
12 +from vllm.compilation.fx_utils import find_op_nodes
13 +from vllm.compilation.noop_elimination import NoOpEliminationPass
14 +from vllm.config import CompilationConfig, CompilationLevel, VllmConfig
15 +from vllm.platforms import current_platform
16 +
17 +# globals needed for string-import custom Dynamo backend field
18 +backend: Optional[TestBackend] = None
19 +backend_unfused: Optional[TestBackend] = None
20 +
21 +
22 +@pytest.mark.parametrize(
23 + "model, quant_key",
24 + [("amd/Llama-3.1-8B-Instruct-FP8-KV", kFp8StaticTensorSym)])
25 +@pytest.mark.parametrize(
26 + "use_triton_fa", [True, False] if current_platform.is_rocm() else [False])
27 +@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
28 +@pytest.mark.skipif(not current_platform.is_cuda_alike(),
29 + reason="Only test CUDA and ROCm")
30 +def test_attention_fusion(example_prompts, monkeypatch, model: str,
31 + quant_key: QuantKey, use_triton_fa: bool):
32 + # Clean Dynamo cache to avoid reusing other test cases
33 + # (for some reason the reset at the end is not enough)
34 + torch._dynamo.reset()
35 +
36 + # Use global backends
37 + global backend, backend_unfused
38 +
39 + use_v1 = False # can be made a param once V1 support added
40 + monkeypatch.setenv("VLLM_USE_V1", str(int(use_v1)))
41 + monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", str(int(use_triton_fa)))
42 +
43 + # Prompt 4 seems too open-ended, differs between fused and unfused
44 + # (both outputs look reasonable though)
45 + prompts = example_prompts[:4] + example_prompts[5:]
46 +
47 + compile_config = CompilationConfig(
48 + # DYNAMO_AS_IS triggers custom backend & does full Dynamo compilation
49 + # DYNAMO_ONCE does not properly propagate shapes.
50 + level=CompilationLevel.DYNAMO_AS_IS,
51 + backend="tests.compile.test_fusion_attn.backend_unfused",
52 + )
53 + vllm_config = VllmConfig(compilation_config=compile_config)
54 + backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
55 +
56 + llm = LLM(model,
57 + enforce_eager=True,
58 + compilation_config=compile_config,
59 + gpu_memory_utilization=0.9,
60 + max_model_len=2048)
61 +
62 + sampling_params = SamplingParams(temperature=0.0,
63 + max_tokens=10,
64 + top_p=0.95)
65 +
66 + unfused_output = llm.generate(prompts, sampling_params)
67 + backend_unfused = None # Reset backend to make sure llm gets released
68 + del llm
69 +
70 + compile_config = CompilationConfig(
71 + # DYNAMO_AS_IS triggers custom backend & does full Dynamo compilation
72 + # DYNAMO_ONCE does not properly propagate shapes.
73 + level=CompilationLevel.DYNAMO_AS_IS,
74 + backend="tests.compile.test_fusion_attn.backend",
75 + )
76 + vllm_config = VllmConfig(compilation_config=compile_config)
77 +
78 + # AttnFusionPass needs attention layers to be registered in config upon init
79 + # so we initialize it during compilation.
80 + attn_pass = lambda *args, **kw: AttnFusionPass(vllm_config)(*args, **kw)
81 + backend = TestBackend(NoOpEliminationPass(vllm_config), attn_pass)
82 + llm2 = LLM(model,
83 + enforce_eager=True,
84 + compilation_config=compile_config,
85 + gpu_memory_utilization=0.9,
86 + max_model_len=2048)
87 +
88 + # check support
89 + attn_fusion_supported = [
90 + layer.impl.fused_output_quant_supported(quant_key.dtype,
91 + quant_key.static,
92 + quant_key.group_shape)
93 + for key, layer in compile_config.static_forward_context.items()
94 + ]
95 +
96 + print(f"{attn_fusion_supported=}")
97 + if any(attn_fusion_supported):
98 + # Check quant ops
99 + backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=False)
100 +
101 + # attention ops present in both, just output_scale param changes
102 + attn_nodes_pre = list(find_op_nodes(ATTN_OP, backend.graph_pre_pass))
103 + attn_nodes_post = list(find_op_nodes(ATTN_OP, backend.graph_post_pass))
104 + assert len(attn_nodes_pre) == len(attn_nodes_post)
105 +
106 + for i in range(len(attn_nodes_pre)):
107 + assert attn_nodes_pre[i].kwargs["output_scale"] is None
108 + fused = attn_nodes_post[i].kwargs["output_scale"] is not None
109 + assert fused == attn_fusion_supported[i], \
110 + f"Node {i} {'' if fused else 'not '} expected " \
111 + f"to have fused output quant"
112 +
113 + # check outputs
114 + fused_output = llm2.generate(prompts, sampling_params)
115 +
116 + # transform outputs to format expected by check_outputs_equal
117 + sample_outs = lambda s: (list(s.token_ids), s.text)
118 + outs_lst = lambda ros: [sample_outs(ro.outputs[0]) for ro in ros]
119 +
120 + check_outputs_equal(
121 + outputs_0_lst=outs_lst(unfused_output),
122 + outputs_1_lst=outs_lst(fused_output),
123 + name_0="unfused",
124 + name_1="fused",
125 + )
126 +
127 + # Clean Dynamo cache to avoid polluting other case(s)
128 + torch._dynamo.reset()
129 +
130 + # Reset backend to make sure llm2 gets released
131 + backend = None
vllm/_custom_ops.py
+7 -1
@@ -1225,6 +1225,7 @@ def scaled_fp8_quant(
1225 num_token_padding: Optional[int] = None,
1226 scale_ub: Optional[torch.Tensor] = None,
1227 use_per_token_if_dynamic: bool = False,
1228 + output: Optional[torch.Tensor] = None,
1229 ) -> tuple[torch.Tensor, torch.Tensor]:
1230 """
1231 Quantize input tensor to FP8 and return quantized tensor and scale.
@@ -1256,7 +1257,12 @@ def scaled_fp8_quant(
1257 out_dtype: torch.dtype = current_platform.fp8_dtype()
1258 if num_token_padding:
1259 shape = (max(num_token_padding, input.shape[0]), shape[1])
1259 - output = torch.empty(shape, device=input.device, dtype=out_dtype)
1260 + if output is None:
1261 + output = torch.empty(shape, device=input.device, dtype=out_dtype)
1262 + else:
1263 + assert num_token_padding is None, \
1264 + "padding not supported if output passed in"
1265 + assert output.dtype == out_dtype
1266
1267 if scale is None:
1268 if use_per_token_if_dynamic:
vllm/attention/backends/abstract.py
+17
@@ -284,9 +284,25 @@ class AttentionImpl(ABC, Generic[T]):
284 kv_cache: torch.Tensor,
285 attn_metadata: T,
286 output: Optional[torch.Tensor] = None,
287 + output_scale: Optional[torch.Tensor] = None,
288 ) -> torch.Tensor:
289 raise NotImplementedError
290
291 + def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
292 + group_shape: tuple[int, int]):
293 + """
294 + Does this attention implementation support fused output quantization.
295 + This is used by the AttnFusionPass to only fuse output quantization
296 + onto implementations that support it.
297 +
298 + TODO(luka) merge parameters into QuantDescriptor
299 + :param dtype: quantized dtype
300 + :param static: static or dynamic quantization
301 + :param group_shape: quant group shape. (-1, -1) for per-tensor.
302 + :return: is fusion supported for this type of quantization
303 + """
304 + return False
305 +
306
307 class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
308
@@ -300,6 +316,7 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
316 kv_cache: torch.Tensor,
317 attn_metadata: T,
318 output: Optional[torch.Tensor] = None,
319 + output_scale: Optional[torch.Tensor] = None,
320 ) -> torch.Tensor:
321 raise NotImplementedError
322
vllm/attention/backends/blocksparse_attn.py
+6
@@ -374,6 +374,7 @@ class BlocksparseFlashAttentionImpl(AttentionImpl):
374 kv_cache: torch.Tensor,
375 attn_metadata: BlocksparseFlashAttentionMetadata,
376 output: Optional[torch.Tensor] = None,
377 + output_scale: Optional[torch.Tensor] = None,
378 ) -> torch.Tensor:
379 """Forward pass with FlashAttention and PagedAttention.
380
@@ -388,6 +389,11 @@ class BlocksparseFlashAttentionImpl(AttentionImpl):
389 Returns:
390 shape = [num_tokens, num_heads * head_size]
391 """
392 + if output_scale is not None:
393 + raise NotImplementedError(
394 + "fused output quantization is not yet supported"
395 + " for BlocksparseFlashAttentionImpl")
396 +
397 num_tokens, hidden_size = query.shape
398 # Reshape the query, key, and value tensors.
399 query = query.view(-1, self.num_heads, self.head_size)
vllm/attention/backends/dual_chunk_flash_attn.py
+9
@@ -370,6 +370,8 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
370 value: torch.Tensor,
371 kv_cache: torch.Tensor,
372 attn_metadata: DualChunkFlashAttentionMetadata,
373 + output: Optional[torch.Tensor] = None,
374 + output_scale: Optional[torch.Tensor] = None,
375 ) -> torch.Tensor:
376 """Forward pass with DualChunkFlashAttention.
377 Args:
@@ -383,6 +385,13 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
385 Returns:
386 shape = [num_tokens, num_heads * head_size]
387 """
388 + assert output is None, "Output tensor not supported for DualChunk"
389 +
390 + if output_scale is not None:
391 + raise NotImplementedError(
392 + "fused output quantization is not yet supported"
393 + " for FlashAttentionImpl")
394 +
395 (
396 query,
397 query_succ,
vllm/attention/backends/flash_attn.py
+6
@@ -673,6 +673,7 @@ class FlashAttentionImpl(AttentionImpl):
673 kv_cache: torch.Tensor,
674 attn_metadata: FlashAttentionMetadata,
675 output: Optional[torch.Tensor] = None,
676 + output_scale: Optional[torch.Tensor] = None,
677 ) -> torch.Tensor:
678 """Forward pass with FlashAttention.
679
@@ -692,6 +693,11 @@ class FlashAttentionImpl(AttentionImpl):
693 """
694 assert output is not None, "Output tensor must be provided."
695
696 + if output_scale is not None:
697 + raise NotImplementedError(
698 + "fused output quantization is not yet supported"
699 + " for FlashAttentionImpl")
700 +
701 # NOTE(woosuk): FlashAttention2 does not support FP8 KV cache.
702 if not flash_attn_supports_fp8() or output.dtype != torch.bfloat16:
703 assert (
vllm/attention/backends/flashinfer.py
+6
@@ -975,8 +975,14 @@ class FlashInferImpl(AttentionImpl):
975 kv_cache: torch.Tensor,
976 attn_metadata: FlashInferMetadata,
977 output: Optional[torch.Tensor] = None,
978 + output_scale: Optional[torch.Tensor] = None,
979 ) -> torch.Tensor:
980
981 + if output_scale is not None:
982 + raise NotImplementedError(
983 + "fused output quantization is not yet supported"
984 + " for FlashInferImpl")
985 +
986 # TODO: directly write to output tensor
987 num_heads: int = self.num_heads
988 head_size: int = self.head_size
vllm/attention/backends/hpu_attn.py
+6
@@ -181,6 +181,7 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
181 kv_cache: torch.Tensor,
182 attn_metadata: HPUAttentionMetadata,
183 output: Optional[torch.Tensor] = None,
184 + output_scale: Optional[torch.Tensor] = None,
185 ) -> torch.Tensor:
186 """Forward pass with xFormers and PagedAttention.
187
@@ -193,6 +194,11 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
194 Returns:
195 shape = [num_tokens, num_heads * head_size]
196 """
197 + if output_scale is not None:
198 + raise NotImplementedError(
199 + "fused output quantization is not yet supported"
200 + " for HPUAttentionImpl")
201 +
202 batch_size, seq_len, hidden_size = query.shape
203 _, seq_len_kv, _ = key.shape
204
vllm/attention/backends/ipex_attn.py
+6
@@ -192,6 +192,7 @@ class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):
192 kv_cache: torch.Tensor,
193 attn_metadata: IpexAttnMetadata, # type: ignore
194 output: Optional[torch.Tensor] = None,
195 + output_scale: Optional[torch.Tensor] = None,
196 ) -> torch.Tensor:
197 """Forward pass with IPEX varlen_attention and PagedAttention.
198
@@ -206,6 +207,11 @@ class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):
207 Returns:
208 shape = [num_tokens, num_heads * head_size]
209 """
210 + if output_scale is not None:
211 + raise NotImplementedError(
212 + "fused output quantization is not yet supported"
213 + " for IpexAttentionImpl")
214 +
215 assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
216 num_tokens, hidden_size = query.shape
217 # Reshape the query, key, and value tensors.
vllm/attention/backends/mla/common.py
+6
@@ -1319,11 +1319,17 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1319 kv_cache: torch.Tensor,
1320 attn_metadata: T,
1321 output: Optional[torch.Tensor] = None,
1322 + output_scale: Optional[torch.Tensor] = None,
1323 ) -> torch.Tensor:
1324 if output is not None:
1325 raise NotImplementedError(
1326 "output is not yet supported for MLAImplBase")
1327
1328 + if output_scale is not None:
1329 + raise NotImplementedError(
1330 + "fused output quantization is not yet supported"
1331 + " for MLAImplBase")
1332 +
1333 if attn_metadata.is_profile_run and \
1334 attn_metadata.context_chunk_workspace is not None:
1335 # During the profile run try to simulate to worse case output size
vllm/attention/backends/pallas.py
+6
@@ -172,6 +172,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
172 kv_cache: Tuple[torch.Tensor, torch.Tensor],
173 attn_metadata: PallasMetadata,
174 output: Optional[torch.Tensor] = None,
175 + output_scale: Optional[torch.Tensor] = None,
176 ) -> torch.Tensor:
177 """Forward pass with Pallas attention.
178
@@ -187,6 +188,11 @@ class PallasAttentionBackendImpl(AttentionImpl):
188 Returns:
189 shape = [batch_size, seq_len, num_heads * head_size]
190 """
191 + if output_scale is not None:
192 + raise NotImplementedError(
193 + "fused output quantization is not yet supported"
194 + " for PallasAttentionImpl")
195 +
196 assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0
197 batch_size, seq_len, hidden_size = query.shape
198 query = query.view(batch_size, seq_len, self.num_heads, self.head_size)
vllm/attention/backends/rocm_flash_attn.py
+37 -4
@@ -38,11 +38,11 @@ def is_rocm_aiter_paged_attn_enabled() -> bool:
38 @cache
39 def _get_paged_attn_module() -> PagedAttention:
40 """
41 - Initializes the appropriate PagedAttention module from `attention/ops`,
41 + Initializes the appropriate PagedAttention module from `attention/ops`,
42 which is used as helper function
43 by `ROCmFlashAttentionImpl` and `ROCmFlashAttentionBackend`.
44
45 - The choice of attention module depends on whether
45 + The choice of attention module depends on whether
46 AITER paged attention is enabled:
47 - If enabled, `ROCmFlashAttentionImpl` uses `AITERPagedAttention`.
48 - Otherwise, it defaults to using the original `PagedAttention`.
@@ -598,6 +598,15 @@ class ROCmFlashAttentionImpl(AttentionImpl):
598 head_dim).reshape(tokens, n_kv_heads * n_rep,
599 head_dim))
600
601 + def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
602 + group_shape: tuple[int, int]):
603 + if self.use_triton_flash_attn:
604 + return dtype == current_platform.fp8_dtype(
605 + ) and static and group_shape == (-1, -1) # per-tensor
606 +
607 + # Only supported in the Triton backend
608 + return False
609 +
610 def forward(
611 self,
612 layer: AttentionLayer,
@@ -607,6 +616,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
616 kv_cache: torch.Tensor,
617 attn_metadata: ROCmFlashAttentionMetadata,
618 output: Optional[torch.Tensor] = None,
619 + output_scale: Optional[torch.Tensor] = None,
620 ) -> torch.Tensor:
621 """Forward pass with FlashAttention and PagedAttention.
622
@@ -660,6 +670,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
670 """
671 assert output is not None, "Output tensor must be provided."
672
673 + if output_scale is not None and not self.use_triton_flash_attn:
674 + raise NotImplementedError(
675 + "fused output quantization only supported for Triton"
676 + " implementation in ROCMFlashAttentionImpl for now")
677 +
678 query = query.view(-1, self.num_heads, self.head_size)
679 if key is not None:
680 assert value is not None
@@ -799,6 +814,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
814 attn_masks[0][None]
815 if attn_masks is not None else None,
816 full_scales,
817 + output_scale,
818 )
819 elif self.use_naive_attn:
820 if self.num_kv_heads != self.num_heads:
@@ -876,6 +892,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
892 decode_query.dtype, head_size, block_size, gqa_ratio,
893 decode_meta.max_decode_seq_len, self.sliding_window,
894 self.kv_cache_dtype, self.alibi_slopes)
895 +
896 if use_custom:
897 max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type
898 != AttentionType.ENCODER_DECODER else
@@ -887,7 +904,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
904 assert _PARTITION_SIZE_ROCM % block_size == 0
905 tmp_output = torch.empty(
906 size=(num_seqs, num_heads, max_num_partitions, head_size),
890 - dtype=output.dtype,
907 + dtype=query.dtype,
908 device=output.device,
909 )
910 exp_sums = torch.empty(
@@ -921,9 +938,17 @@ class ROCmFlashAttentionImpl(AttentionImpl):
938 self.kv_cache_dtype,
939 layer._k_scale,
940 layer._v_scale,
941 + output_scale,
942 )
943 else:
926 - output[num_prefill_tokens:] = paged_attn.forward_decode(
944 + # PagedAttention does not support fused quant, manually quantize
945 + if output_scale is None:
946 + out_pa = output[num_prefill_tokens:]
947 + else:
948 + out_pa = torch.empty_like(output[num_prefill_tokens:],
949 + dtype=query.dtype)
950 +
951 + out_pa[:] = paged_attn.forward_decode(
952 decode_query,
953 key_cache,
954 value_cache,
@@ -944,6 +969,14 @@ class ROCmFlashAttentionImpl(AttentionImpl):
969 layer._v_scale,
970 )
971
972 + # Manually perform quantization
973 + if output_scale is not None:
974 + out_uq = out_pa.view(-1, self.num_heads * self.head_size)
975 + out_q = output.view(-1, self.num_heads * self.head_size)
976 + ops.scaled_fp8_quant(out_uq,
977 + output_scale,
978 + output=out_q[num_prefill_tokens:])
979 +
980 # Reshape the output tensor.
981 return output.view(-1, self.num_heads * self.head_size)
982
vllm/attention/backends/torch_sdpa.py
+5
@@ -459,6 +459,7 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
459 kv_cache: torch.Tensor,
460 attn_metadata: TorchSDPAMetadata, # type: ignore
461 output: Optional[torch.Tensor] = None,
462 + output_scale: Optional[torch.Tensor] = None,
463 ) -> torch.Tensor:
464 """Forward pass with torch SDPA and PagedAttention.
465
@@ -473,6 +474,10 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
474 Returns:
475 shape = [num_tokens, num_heads * head_size]
476 """
477 + if output_scale is not None:
478 + raise NotImplementedError(
479 + "fused output quantization is not yet supported"
480 + " for TorchSDPABackendImpl")
481
482 # For warming-up
483 if attn_metadata is None:
vllm/attention/backends/xformers.py
+6
@@ -435,6 +435,7 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
435 kv_cache: torch.Tensor,
436 attn_metadata: "XFormersMetadata",
437 output: Optional[torch.Tensor] = None,
438 + output_scale: Optional[torch.Tensor] = None,
439 ) -> torch.Tensor:
440 """Forward pass with xFormers and PagedAttention.
441
@@ -487,6 +488,11 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
488 Returns:
489 shape = [num_tokens, num_heads * head_size]
490 """
491 + if output_scale is not None:
492 + raise NotImplementedError(
493 + "fused output quantization is not yet supported"
494 + " for XFormersImpl")
495 +
496 attn_type = self.attn_type
497 # Check that appropriate attention metadata attributes are
498 # selected for the desired attention type
vllm/attention/layer.py
+4 -1
@@ -430,6 +430,7 @@ def unified_attention_with_output(
430 value: torch.Tensor,
431 output: torch.Tensor,
432 layer_name: str,
433 + output_scale: Optional[torch.Tensor] = None,
434 ) -> None:
435 wait_for_kv_layer_from_connector(layer_name)
436 forward_context: ForwardContext = get_forward_context()
@@ -444,7 +445,8 @@ def unified_attention_with_output(
445 value,
446 kv_cache,
447 attn_metadata,
447 - output=output)
448 + output=output,
449 + output_scale=output_scale)
450
451 maybe_save_kv_layer_to_connector(layer_name, kv_cache)
452
@@ -455,6 +457,7 @@ def unified_attention_with_output_fake(
457 value: torch.Tensor,
458 output: torch.Tensor,
459 layer_name: str,
460 + output_scale: Optional[torch.Tensor] = None,
461 ) -> None:
462 return
463
vllm/compilation/fusion.py
+63 -36
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Callable, NamedTuple, Optional
3 +from typing import Callable, ClassVar, NamedTuple, Optional
4
5 import torch
6 import torch._inductor.pattern_matcher as pm
@@ -34,36 +33,66 @@ RMS_OP = torch.ops._C.rms_norm.default
33 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
34
35
36 +# Use proxy as NamedTuple direct subclasses cannot have static members
37 +class _GroupShape(NamedTuple):
38 + row: int
39 + col: int
40 +
41 +
42 +class GroupShape(_GroupShape):
43 + """
44 + This class describes the quantization group shape.
45 + It includes static members for common shapes (per-tensor, per-token).
46 + """
47 +
48 + # Aliases for common quantization group shapes
49 + PER_TENSOR: ClassVar['GroupShape']
50 + PER_TOKEN: ClassVar['GroupShape']
51 +
52 +
53 +GroupShape.PER_TENSOR = GroupShape(-1, -1)
54 +GroupShape.PER_TOKEN = GroupShape(1, -1)
55 +
56 +
57 class QuantKey(NamedTuple):
58 """
59 Named tuple for identifying the type of quantization.
60 dtype: quantized data type
61 static: static quantization if True, dynamic if False
42 - per_tensor: per-tensor quantization if True, per-token if False
62 + group_shape: quantization group shape
63 symmetric: symmetric if True, asymmetric if False
64 +
65 + TODO(luka) use QuantDescriptor once standardized:
66 + https://github.com/vllm-project/vllm/issues/8913
67 +
68 """
69 dtype: torch.dtype
70 static: bool
47 - per_tensor: bool = True
71 + group_shape: GroupShape
72 symmetric: bool = True
73
74 def __str__(self):
75 + group_shape = ('per_tensor'
76 + if self.group_shape == GroupShape.PER_TENSOR else
77 + ('per_token' if self.group_shape == GroupShape.PER_TOKEN
78 + else str(self.group_shape)))
79 +
80 return (f"QuantKey({'static' if self.static else 'dynamic'},"
52 - f"{fx.graph.dtype_abbrs[self.dtype]},"
53 - f"{'per_tensor' if self.per_tensor else 'per_token'},"
81 + f"{fx.graph.dtype_abbrs[self.dtype]},{group_shape},"
82 f"{'a' if not self.symmetric else ''}symmetric)")
83
84
57 -kFp8StaticTensorSym = QuantKey(FP8_DTYPE, True, True, True)
58 -kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, True, True)
59 -kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, False, True)
85 +kFp8StaticTensorSym = QuantKey(FP8_DTYPE, True, GroupShape.PER_TENSOR, True)
86 +kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TENSOR, True)
87 +kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TOKEN, True)
88
89 QUANT_OPS: dict[QuantKey, OpOverload] = {
62 - kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa
90 + kFp8StaticTensorSym:
91 + torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501
92 kFp8DynamicTensorSym:
64 - torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa
93 + torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501
94 kFp8DynamicTokenSym:
66 - torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa
95 + torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
96 }
97
98
@@ -83,13 +112,13 @@ class FusedRMSQuantKey(NamedTuple):
112
113 FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
114 FusedRMSQuantKey(kFp8StaticTensorSym, False):
86 - torch.ops._C.rms_norm_static_fp8_quant.default, # noqa
115 + torch.ops._C.rms_norm_static_fp8_quant.default, # noqa: E501
116 FusedRMSQuantKey(kFp8StaticTensorSym, True):
88 - torch.ops._C.fused_add_rms_norm_static_fp8_quant.default, # noqa
117 + torch.ops._C.fused_add_rms_norm_static_fp8_quant.default, # noqa: E501
118 FusedRMSQuantKey(kFp8DynamicTokenSym, False):
90 - torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa
119 + torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
120 FusedRMSQuantKey(kFp8DynamicTokenSym, True):
92 - torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa
121 + torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
122 }
123
124
@@ -177,10 +206,11 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
206 quant_dtype: torch.dtype,
207 symmetric=True):
208 fused_key = FusedRMSQuantKey(fused_add=False,
180 - quant=QuantKey(dtype=quant_dtype,
181 - static=True,
182 - per_tensor=True,
183 - symmetric=symmetric))
209 + quant=QuantKey(
210 + dtype=quant_dtype,
211 + static=True,
212 + group_shape=GroupShape.PER_TENSOR,
213 + symmetric=symmetric))
214 super().__init__(epsilon, fused_key)
215
216 def register(self, pm_pass: PatternMatcherPass):
@@ -233,10 +263,11 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
263 quant_dtype: torch.dtype,
264 symmetric=True):
265 key = FusedRMSQuantKey(fused_add=True,
236 - quant=QuantKey(dtype=quant_dtype,
237 - static=True,
238 - per_tensor=True,
239 - symmetric=symmetric))
266 + quant=QuantKey(
267 + dtype=quant_dtype,
268 + static=True,
269 + group_shape=GroupShape.PER_TENSOR,
270 + symmetric=symmetric))
271 super().__init__(epsilon, key)
272
273 def register(self, pm_pass: PatternMatcherPass,
@@ -323,12 +354,12 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
354 def __init__(self,
355 epsilon: float,
356 quant_dtype: torch.dtype,
326 - per_tensor: bool,
357 + group_shape: GroupShape = GroupShape.PER_TOKEN,
358 symmetric=True):
359 key = FusedRMSQuantKey(fused_add=False,
360 quant=QuantKey(dtype=quant_dtype,
361 static=False,
331 - per_tensor=per_tensor,
362 + group_shape=group_shape,
363 symmetric=symmetric))
364 super().__init__(epsilon, key)
365
@@ -421,12 +452,12 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
452 def __init__(self,
453 epsilon: float,
454 quant_dtype: torch.dtype,
424 - per_tensor: bool = True,
455 + group_shape: GroupShape = GroupShape.PER_TOKEN,
456 symmetric=True):
457 key = FusedRMSQuantKey(fused_add=True,
458 quant=QuantKey(dtype=quant_dtype,
459 static=False,
429 - per_tensor=per_tensor,
460 + group_shape=group_shape,
461 symmetric=symmetric))
462 super().__init__(epsilon, key)
463
@@ -566,16 +597,12 @@ class FusionPass(VllmInductorPass):
597 self.patterns, self.record_match)
598
599 # Fuse rms_norm + dynamic per-token fp8 quant
569 - RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE,
570 - per_tensor=False).register(
571 - self.patterns, self.record_match)
600 + RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
601 + self.patterns, self.record_match)
602
603 # Fuse fused_add_rms_norm + dynamic per-token fp8 quant
574 - FusedAddRMSNormDynamicQuantPattern(epsilon,
575 - FP8_DTYPE,
576 - per_tensor=False).register(
577 - self.patterns,
578 - self.record_match)
604 + FusedAddRMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
605 + self.patterns, self.record_match)
606
607 # WARNING: This is a hack to clear the pattern matcher cache
608 # and allow multiple values of epsilon.
vllm/compilation/fusion_attn.py new
+165
@@ -0,0 +1,165 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import torch
4 +import torch._inductor.pattern_matcher as pm
5 +from torch._higher_order_ops.auto_functionalize import auto_functionalized
6 +from torch._inductor.pattern_matcher import PatternMatcherPass
7 +from torch._subclasses.fake_tensor import (FakeTensorMode,
8 + unset_fake_temporarily)
9 +
10 +from vllm.attention import Attention
11 +from vllm.config import VllmConfig
12 +from vllm.logger import init_logger
13 +from vllm.platforms import current_platform
14 +
15 +from .fusion import QUANT_OPS, GroupShape, QuantKey, empty_bf16, empty_fp32
16 +from .vllm_inductor_pass import VllmInductorPass
17 +
18 +logger = init_logger(__name__)
19 +
20 +ATTN_OP = torch.ops.vllm.unified_attention_with_output.default
21 +RESHAPE_OP = torch.ops.aten.reshape.default
22 +
23 +
24 +class AttentionStaticQuantPattern:
25 +
26 + def __init__(
27 + self,
28 + layer_name: str,
29 + num_heads: int,
30 + head_size: int,
31 + quant_dtype: torch.dtype,
32 + symmetric=True,
33 + ):
34 + self.layer_name = layer_name
35 + self.num_heads = num_heads
36 + self.head_size = head_size
37 + self.quant_dtype = quant_dtype
38 + self.quant_key = QuantKey(dtype=quant_dtype,
39 + static=True,
40 + group_shape=GroupShape.PER_TENSOR,
41 + symmetric=symmetric)
42 + assert self.quant_key in QUANT_OPS, \
43 + f"unsupported quantization scheme {self.quant_key}"
44 + self.QUANT_OP = QUANT_OPS[self.quant_key]
45 +
46 + def empty_quant(self, *args, **kwargs):
47 + kwargs = {'dtype': self.quant_dtype, 'device': "cuda", **kwargs}
48 + return torch.empty(*args, **kwargs)
49 +
50 + def register_if_supported(self, pm_pass: PatternMatcherPass,
51 + layer: Attention):
52 + if layer.impl.fused_output_quant_supported(self.quant_dtype,
53 + self.quant_key.static,
54 + self.quant_key.group_shape):
55 + self._register(pm_pass)
56 +
57 + def _register(self, pm_pass: PatternMatcherPass):
58 +
59 + def pattern(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
60 + output_attn: torch.Tensor, output_quant: torch.Tensor,
61 + scale: torch.Tensor):
62 + view_7 = RESHAPE_OP(output_attn,
63 + [-1, self.num_heads, self.head_size])
64 +
65 + at1 = auto_functionalized(ATTN_OP,
66 + query=q,
67 + key=k,
68 + value=v,
69 + output=view_7,
70 + layer_name=self.layer_name,
71 + output_scale=None)
72 + attn_out_view = RESHAPE_OP(at1[1],
73 + [-1, self.num_heads * self.head_size])
74 +
75 + at2 = auto_functionalized(self.QUANT_OP,
76 + result=output_quant,
77 + input=attn_out_view,
78 + scale=scale)
79 + return at2[1]
80 +
81 + def replacement(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
82 + output_attn: torch.Tensor, output_quant: torch.Tensor,
83 + scale: torch.Tensor):
84 + view_7 = RESHAPE_OP(output_quant,
85 + [-1, self.num_heads, self.head_size])
86 +
87 + at1 = auto_functionalized(ATTN_OP,
88 + query=q,
89 + key=k,
90 + value=v,
91 + output=view_7,
92 + layer_name=self.layer_name,
93 + output_scale=scale)
94 +
95 + return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
96 +
97 + # Need custom fake mode, otherwise tracing happens with real tensors.
98 + # That would not work for the unified_attention custom op.
99 + with unset_fake_temporarily(), FakeTensorMode():
100 + inputs = [
101 + empty_bf16(5, self.num_heads, self.head_size), # q
102 + empty_bf16(5, self.num_heads, self.head_size), # k
103 + empty_bf16(5, self.num_heads, self.head_size), # v
104 + empty_bf16(5, self.num_heads * self.head_size), # attn_output
105 + self.empty_quant(5, self.num_heads *
106 + self.head_size), # quant_output
107 + empty_fp32(1, 1) # scale
108 + ]
109 +
110 + def wrap_trace_fn(process_fx, trace_fn):
111 +
112 + def wrapped(*args, **kwargs):
113 + return process_fx(trace_fn(*args, **kwargs))
114 +
115 + return wrapped
116 +
117 + def fx_view_to_reshape(gm: torch.fx.GraphModule):
118 + from torch._inductor.fx_passes.post_grad import view_to_reshape
119 + view_to_reshape(gm)
120 + return gm
121 +
122 + pm.register_replacement(
123 + pattern, replacement, inputs,
124 + wrap_trace_fn(fx_view_to_reshape, pm.fwd_only), pm_pass)
125 +
126 +
127 +class AttnFusionPass(VllmInductorPass):
128 + """
129 + This pass fuses post-attention quantization onto attention if supported.
130 +
131 + It uses the pattern matcher and matches each layer manually, as strings
132 + cannot be wildcarded. This also lets us check support on attention layers
133 + upon registration instead of during pattern matching.
134 +
135 + Currently, only static fp8 quant is supported, but patterns could easily be
136 + added for other quant schemes and dtypes. The bigger hurdle for wider
137 + support are attention kernels, which need to support fusing output quant.
138 + """
139 +
140 + def __init__(self, config: VllmConfig):
141 + super().__init__(config)
142 + self.static_fwd_ctx = config.compilation_config.static_forward_context
143 +
144 + self.patterns = PatternMatcherPass(pass_name="attn_fusion_pass")
145 +
146 + for key, layer in self.static_fwd_ctx.items():
147 + pattern = AttentionStaticQuantPattern(key, layer.num_heads,
148 + layer.head_size,
149 + current_platform.fp8_dtype())
150 + pattern.register_if_supported(self.patterns, layer)
151 + if len(self.static_fwd_ctx) == 0:
152 + logger.warning(
153 + "Attention + quant fusion is enabled, but "
154 + "CompilationConfig.static_forward_context is empty. "
155 + "Cannot access attention layers so no fusion "
156 + "patterns were registered.")
157 +
158 + def __call__(self, graph: torch.fx.graph.Graph) -> None:
159 + self.begin()
160 + self.dump_graph(graph, "before_attn_fusion")
161 +
162 + count = self.patterns.apply(graph)
163 + logger.debug("Fused quantization onto %s attention nodes", count)
164 + self.dump_graph(graph, "after_attn_fusion")
165 + self.end_and_log()
vllm/compilation/fx_utils.py
+23 -1
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import operator
5 -from collections.abc import Iterable
5 +from collections.abc import Iterable, Iterator
6 from typing import Optional
7
8 from torch import fx
@@ -14,6 +14,10 @@ def is_func(node: fx.Node, target) -> bool:
14 return node.op == "call_function" and node.target == target
15
16
17 +def is_auto_func(node: fx.Node, op: OpOverload) -> bool:
18 + return is_func(node, auto_functionalized) and node.args[0] == op
19 +
20 +
21 # Returns the first specified node with the given op (if it exists)
22 def find_specified_fn_maybe(nodes: Iterable[fx.Node],
23 op: OpOverload) -> Optional[fx.Node]:
@@ -60,3 +64,21 @@ def find_getitem(node: fx.Node, idx: int) -> fx.Node:
64 ret = find_getitem_maybe(node, idx)
65 assert ret is not None, f"Could not find getitem {idx} in node {node}"
66 return ret
67 +
68 +
69 +# An auto-functionalization-aware utility for finding nodes with a specific op
70 +def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]:
71 + if not op._schema.is_mutable:
72 + yield from graph.find_nodes(op="call_function", target=op)
73 +
74 + for n in graph.find_nodes(op="call_function", target=auto_functionalized):
75 + if n.args[0] == op:
76 + yield n
77 +
78 +
79 +# Asserts that the node only has one user and returns it
80 +# Even if a node has only 1 user, it might share storage with another node,
81 +# which might need to be taken into account.
82 +def get_only_user(node: fx.Node) -> fx.Node:
83 + assert len(node.users) == 1
84 + return next(iter(node.users))
vllm/compilation/noop_elimination.py
+29 -1
@@ -23,7 +23,23 @@ class NoOpEliminationPass(VllmInductorPass):
23 in the 2D-case. Additionally, torch internal no-op elimination pass does
24 not handle certain slice variants.
25
26 + Cases handled:
27 + 1. A chain of reshapes is equivalent to the last reshape called on the
28 + base tensor (input of the first reshape).
29 + 2. A reshape that produces the shape of the input is redundant
30 + 3. A slice that produces the shape of the input is redundant
31 +
32 Example graph 1:
33 + mul_1: "f16[s0, 4096]" = ...
34 + view_1: "f16[s0, 128, 32]" = torch.reshape(mul_1, [-1, 128, 32])
35 + view_2: "f16[s0, 4096]" = torch.reshape(view_2, [-1, 4096])
36 + view_3: "f16[s0, 128, 32]" = torch.reshape(view_3, [-1, 128, 32])
37 +
38 + Can be replaced with:
39 + mul_1: "f16[s0, 4096]" = ...
40 + view_3: "f16[s0, 128, 32]" = ...
41 +
42 + Example graph 2:
43 getitem_1: "f16[s0, 4096]" = ...
44 view_1: "f16[s0, 4096]" = torch.reshape(getitem_1, [-1, 4096])
45 at = auto_functionalized(static_scaled_fp8_quant, input = view_1, ...)
@@ -34,7 +50,7 @@ class NoOpEliminationPass(VllmInductorPass):
50 at = auto_functionalized(static_scaled_fp8_quant, input = getitem_1, ...)
51 out: "f8e4m3fn[s0, 4096]" = at[1]
52
37 - Example graph 2:
53 + Example graph 3:
54 arg0: "s0" = SymInt(s0)
55 scaled_mm: "f16[s0, 4096]" = ...
56 slice_1: "f16[s0, 4096]" = torch.slice(scaled_mm, -1, 0, arg0)
@@ -58,6 +74,18 @@ class NoOpEliminationPass(VllmInductorPass):
74 # Remove no-op reshapes/views:
75 for node in graph.nodes:
76 if is_func(node, torch.ops.aten.reshape.default):
77 + # Case 1: rewrite reshape chains to reshapes on the base tensor
78 + input = node.args[0]
79 + # If the input is a reshape, rebind to that node
80 + if is_func(input, torch.ops.aten.reshape.default):
81 + # The new input is guaranteed not to be a reshape,
82 + # because we process nodes in order
83 + node.update_arg(0, input.args[0])
84 + if len(input.users) == 0:
85 + graph.erase_node(input)
86 + count += 1
87 +
88 + # Case 2: remove this reshape if it produces the original shape
89 input, shape = node.args[:2]
90 input_shape = input.meta["val"].shape
91 if len(shape) != len(input_shape):
vllm/compilation/pass_manager.py
+4
@@ -10,6 +10,7 @@ from .activation_quant_fusion import ActivationQuantFusionPass
10 from .collective_fusion import AsyncTPPass
11 from .fix_functionalization import FixFunctionalizationPass
12 from .fusion import FusionPass
13 +from .fusion_attn import AttnFusionPass
14 from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
15 from .noop_elimination import NoOpEliminationPass
16 from .sequence_parallelism import SequenceParallelismPass
@@ -59,6 +60,9 @@ class PostGradPassManager(CustomGraphPass):
60 if self.pass_config.enable_async_tp:
61 self.passes += [AsyncTPPass(config)]
62
63 + if self.pass_config.enable_attn_fusion:
64 + self.passes += [AttnFusionPass(config)]
65 +
66 self.fix_functionalization = FixFunctionalizationPass(config)
67
68 def add(self, pass_: InductorPass):
vllm/compilation/vllm_inductor_pass.py
+3
@@ -4,6 +4,7 @@
4 import time
5
6 import torch
7 +from torch._dynamo.utils import lazy_format_graph_code
8
9 from vllm.config import PassConfig, VllmConfig
10 # yapf: disable
@@ -34,6 +35,8 @@ class VllmInductorPass(InductorPass):
35 self.pass_name = self.__class__.__name__
36
37 def dump_graph(self, graph: torch.fx.Graph, stage: str, always=False):
38 + lazy_format_graph_code(stage, graph.owning_module)
39 +
40 if stage in self.pass_config.dump_graph_stages or always:
41 # Make sure filename includes rank in the distributed setting
42 parallel = p_is_init() and get_tp_world_size() > 1
vllm/config.py
+16 -11
@@ -3804,9 +3804,10 @@ class PassConfig:
3804 its own stages (before, after, maybe in-between)."""
3805 dump_graph_dir: Path = Path(".")
3806 """Directory to dump the graphs."""
3807 - # TODO(luka) better pass enabling system.
3807 enable_fusion: bool = True
3809 - """Whether to enable the custom fusion pass."""
3808 + """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""
3809 + enable_attn_fusion: bool = False
3810 + """Whether to enable the custom attention+quant fusion pass."""
3811 enable_noop: bool = True
3812 """Whether to enable the custom no-op elimination pass."""
3813 enable_sequence_parallelism: bool = False
@@ -3814,6 +3815,8 @@ class PassConfig:
3815 enable_async_tp: bool = False
3816 """Whether to enable async TP."""
3817
3818 + # TODO(luka) better pass enabling system.
3819 +
3820 def uuid(self):
3821 """
3822 Produces a hash unique to the pass configuration.
@@ -3821,18 +3824,20 @@ class PassConfig:
3824 Do not include dump_graph_* in the hash - they don't affect
3825 compilation.
3826 """
3824 - include = {
3825 - "enable_fusion", "enable_noop", "enable_sequence_parallelism",
3826 - "enable_async_tp"
3827 - }
3828 - dict_ = {k: v for k, v in asdict(self).items() if k in include}
3827 + exclude = {"dump_graph_stages", "dump_graph_dir"}
3828 + dict_ = {k: v for k, v in asdict(self).items() if k not in exclude}
3829 return InductorPass.hash_dict(dict_)
3830
3831 def __post_init__(self) -> None:
3832 - if not self.enable_noop and self.enable_fusion:
3833 - logger.warning_once(
3834 - "Fusion enabled but reshape elimination disabled. "
3835 - "RMSNorm + quant (fp8) fusion might not work")
3832 + if not self.enable_noop:
3833 + if self.enable_fusion:
3834 + logger.warning_once(
3835 + "Fusion enabled but reshape elimination disabled. "
3836 + "RMSNorm/SiluMul + quant (fp8) fusion might not work")
3837 + if self.enable_attn_fusion:
3838 + logger.warning_once(
3839 + "Fusion enabled but reshape elimination disabled. "
3840 + "Attention + quant (fp8) fusion might not work")
3841
3842
3843 @config
vllm/envs.py
+1 -1
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
15 VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60
16 VLLM_NCCL_SO_PATH: Optional[str] = None
17 LD_LIBRARY_PATH: Optional[str] = None
18 - VLLM_USE_TRITON_FLASH_ATTN: bool = False
18 + VLLM_USE_TRITON_FLASH_ATTN: bool = True
19 VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False
20 VLLM_FLASH_ATTN_VERSION: Optional[int] = None
21 LOCAL_RANK: int = 0
vllm/v1/attention/backends/flash_attn.py
+6
@@ -569,6 +569,7 @@ class FlashAttentionImpl(AttentionImpl):
569 kv_cache: torch.Tensor,
570 attn_metadata: FlashAttentionMetadata,
571 output: Optional[torch.Tensor] = None,
572 + output_scale: Optional[torch.Tensor] = None,
573 ) -> torch.Tensor:
574 """Forward pass with FlashAttention.
575
@@ -586,6 +587,11 @@ class FlashAttentionImpl(AttentionImpl):
587 """
588 assert output is not None, "Output tensor must be provided."
589
590 + if output_scale is not None:
591 + raise NotImplementedError(
592 + "fused output quantization is not yet supported"
593 + " for FlashAttentionImpl")
594 +
595 if attn_metadata is None:
596 # Profiling run.
597 return output
vllm/v1/attention/backends/flashinfer.py
+6
@@ -547,6 +547,7 @@ class FlashInferImpl(AttentionImpl):
547 kv_cache: torch.Tensor,
548 attn_metadata: FlashInferMetadata,
549 output: Optional[torch.Tensor] = None,
550 + output_scale: Optional[torch.Tensor] = None,
551 ) -> torch.Tensor:
552 """Forward pass with FlashInfer.
553
@@ -561,6 +562,11 @@ class FlashInferImpl(AttentionImpl):
562 """
563 assert output is not None, "Output tensor must be provided."
564
565 + if output_scale is not None:
566 + raise NotImplementedError(
567 + "fused output quantization is not yet supported"
568 + " for FlashInferImpl")
569 +
570 if attn_metadata is None:
571 # Profiling run.
572 return output
vllm/v1/attention/backends/flex_attention.py
+7
@@ -414,6 +414,7 @@ class FlexAttentionImpl(AttentionImpl):
414 kv_cache: torch.Tensor,
415 attn_metadata: FlexAttentionMetadata,
416 output: Optional[torch.Tensor] = None,
417 + output_scale: Optional[torch.Tensor] = None,
418 ) -> torch.Tensor:
419 """Forward pass with FLexAttention.
420
@@ -427,6 +428,12 @@ class FlexAttentionImpl(AttentionImpl):
428 shape = [num_tokens, num_heads * head_size]
429 """
430 assert output is not None, "Output tensor must be provided."
431 +
432 + if output_scale is not None:
433 + raise NotImplementedError(
434 + "fused output quantization is not yet supported"
435 + " for FlexAttentionImpl")
436 +
437 enable_gqa = self.num_kv_heads != self.num_heads
438
439 if attn_metadata is None:
vllm/v1/attention/backends/mla/common.py
+6
@@ -865,10 +865,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
865 kv_cache: torch.Tensor,
866 attn_metadata: M,
867 output: Optional[torch.Tensor] = None,
868 + output_scale: Optional[torch.Tensor] = None,
869 ) -> torch.Tensor:
870
871 assert output is not None, "Output tensor must be provided."
872
873 + if output_scale is not None:
874 + raise NotImplementedError(
875 + "fused output quantization is not yet supported"
876 + " for MLACommonImpl")
877 +
878 if attn_metadata is None:
879 # The zero fill is required when used with DP + EP
880 # to ensure all ranks within a DP group compute the
vllm/v1/attention/backends/pallas.py
+6
@@ -161,6 +161,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
161 kv_cache: torch.Tensor,
162 attn_metadata: PallasMetadata,
163 output: Optional[torch.Tensor] = None,
164 + output_scale: Optional[torch.Tensor] = None,
165 ) -> torch.Tensor:
166 """Forward pass with Pallas attention.
167
@@ -173,6 +174,11 @@ class PallasAttentionBackendImpl(AttentionImpl):
174 Returns:
175 shape = [num_tokens, num_heads * head_size]
176 """
177 + if output_scale is not None:
178 + raise NotImplementedError(
179 + "fused output quantization is not yet supported"
180 + " for PallasAttentionBackendImpl")
181 +
182 # For determine_available_memory case.
183 if kv_cache.numel() == 0:
184 if output is None:
vllm/v1/attention/backends/triton_attn.py
+6
@@ -142,6 +142,7 @@ class TritonAttentionImpl(AttentionImpl):
142 kv_cache: torch.Tensor,
143 attn_metadata: FlashAttentionMetadata,
144 output: Optional[torch.Tensor] = None,
145 + output_scale: Optional[torch.Tensor] = None,
146 ) -> torch.Tensor:
147 """Forward pass with FlashAttention.
148
@@ -156,6 +157,11 @@ class TritonAttentionImpl(AttentionImpl):
157 """
158 assert output is not None, "Output tensor must be provided."
159
160 + if output_scale is not None:
161 + raise NotImplementedError(
162 + "fused output quantization is not yet supported"
163 + " for TritonAttentionImpl")
164 +
165 if attn_metadata is None:
166 # Profiling run.
167 return output