test_fusion_attn.py 17.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
5
6
7

import pytest
import torch._dynamo

8
from tests.compile.backend import LazyInitPass, TestBackend
9
from tests.utils import TestFP8Layer, flat_product
10
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
13
14
15
16
from vllm.compilation.passes.fusion.attn_quant_fusion import ATTN_OP, AttnFusionPass
from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
from vllm.compilation.passes.fx_utils import find_op_nodes
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
17
from vllm.config import (
18
    AttentionConfig,
19
20
    CacheConfig,
    CompilationConfig,
21
    CompilationMode,
22
23
24
25
26
27
    ModelConfig,
    PassConfig,
    SchedulerConfig,
    VllmConfig,
    set_current_vllm_config,
)
28
from vllm.forward_context import get_forward_context, set_forward_context
29
from vllm.model_executor.layers.attention import Attention
30
from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
    QuantKey,
32
    kFp8StaticTensorSym,
33
    kNvfp4Dynamic,
34
)
35
from vllm.platforms import current_platform
36
from vllm.utils.flashinfer import has_flashinfer
37
38
from vllm.v1.attention.backend import AttentionMetadata
from vllm.v1.attention.backends.registry import AttentionBackendEnum
39
40
41
from vllm.v1.kv_cache_interface import AttentionSpec

FP8_DTYPE = current_platform.fp8_dtype()
42
FP4_DTYPE = torch.uint8
43
44


45
46
class AttentionQuantPatternModel(torch.nn.Module):
    """Base model for AttentionQuantPattern fusion."""
47

48
49
50
51
52
53
54
55
56
57
    def __init__(
        self,
        num_qo_heads: int,
        num_kv_heads: int,
        head_size: int,
        kv_cache_dtype: torch.dtype,
        device: torch.device,
        vllm_config: VllmConfig,
        **kwargs,
    ):
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
        super().__init__()
        self.num_qo_heads = num_qo_heads
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
        self.kv_cache_dtype = kv_cache_dtype
        self.device = device
        self.vllm_config = vllm_config

        self.attn = Attention(
            num_heads=self.num_qo_heads,
            head_size=self.head_size,
            scale=1.0 / (self.head_size**0.5),
            num_kv_heads=self.num_kv_heads,
            cache_config=vllm_config.cache_config,
            prefix="model.layers.0.self_attn.attn",
        )
74
75
        self.attn._k_scale = self.attn._k_scale.to(device)
        self.attn._v_scale = self.attn._v_scale.to(device)
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

        self.block_size = 16

        # Initialize attn MetadataBuilder
        self.builder = self.attn.attn_backend.get_builder_cls()(
            kv_cache_spec=AttentionSpec(
                block_size=self.block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                dtype=self.kv_cache_dtype,
            ),
            layer_names=[self.attn.layer_name],
            vllm_config=self.vllm_config,
            device=self.device,
        )

92
    def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
93
94
95
        """Initialize attention metadata."""

        # Create common attn metadata
96
        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
97
        common_attn_metadata = create_common_attn_metadata(
98
99
            batch_spec, self.block_size, self.device, arange_block_indices=True
        )
100

101
        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
102
        num_blocks = batch_size * max_blocks
103
        backend = self.attn.backend
104

105
        # TODO(luka) use get_kv_cache_stride_order
106
        # Create dummy KV cache for the selected backend
107
        if backend == AttentionBackendEnum.ROCM_ATTN:
108
            # k/v as 1st dimention
109
110
111
112
113
114
115
116
117
118
            # HND: [num_blocks, num_kv_heads, block_size, head_size]
            kv_cache = torch.zeros(
                2,
                num_blocks,
                self.num_kv_heads,
                self.block_size,
                self.head_size,
                dtype=self.kv_cache_dtype,
                device=self.device,
            )
119
        elif backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
120
121
122
123
124
125
126
127
128
129
130
            # k/v as 1st dimention
            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
            kv_cache = torch.zeros(
                2,
                num_blocks,
                self.block_size,
                self.num_kv_heads,
                self.head_size,
                dtype=self.kv_cache_dtype,
                device=self.device,
            )
131
        elif backend == AttentionBackendEnum.TRITON_ATTN:
132
            # k/v as 2nd dimention
133
134
135
136
137
138
139
140
141
142
            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
            kv_cache = torch.zeros(
                num_blocks,
                2,
                self.num_kv_heads,
                self.block_size,
                self.head_size,
                dtype=self.kv_cache_dtype,
                device=self.device,
            )
143
        elif backend == AttentionBackendEnum.FLASHINFER:
144
145
146
147
148
149
150
151
152
153
154
            kv_cache = torch.zeros(
                num_blocks,
                2,
                self.num_kv_heads,
                self.block_size,
                self.head_size,
                dtype=self.kv_cache_dtype,
                device=self.device,
            ).permute(0, 1, 3, 2, 4)
        else:
            raise ValueError(f"Unsupported backend: {backend}")
155
156
157
158
        self.attn.kv_cache = [kv_cache]

        # Build attn metadata
        self.attn_metadata = self.builder.build(
159
160
            common_prefix_len=0, common_attn_metadata=common_attn_metadata
        )
161
162
163

        return self.attn_metadata

164
165
166
167
168
169
170
171
172
173

class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
    """Test model for AttentionFp8StaticQuantPattern fusion."""

    quant_key = kFp8StaticTensorSym

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        hidden_size = self.num_qo_heads * self.head_size
174
175
176
177
178
        self.fp8_linear = TestFP8Layer(
            weight_shape=(hidden_size, hidden_size),
            activation_quant_key=self.quant_key,
            weight_quant_key=self.quant_key,
            device=self.device,
179
        )
180

181
182
183
184
185
186
187
188
189
190
191
192
        w = kwargs.get("w")
        if w is not None:
            self.fp8_linear.weight = w["weight"]
            self.fp8_linear.weight_scale = w["wscale"]
            self.fp8_linear.input_scale = w["scale"]

        self.w = {
            "weight": self.fp8_linear.weight,
            "wscale": self.fp8_linear.weight_scale,
            "scale": self.fp8_linear.input_scale,
        }

193
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
194
195
        """Forward pass that creates the pattern to be fused."""
        attn_output = self.attn(q, k, v)
196
        return self.fp8_linear(attn_output)
197
198
199
200
201


class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
    """Test model for AttentionNvfp4QuantPattern fusion."""

202
    quant_key = kNvfp4Dynamic
203
204
205
206
207
208

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        hidden_size = self.num_qo_heads * self.head_size
        self.w = kwargs.get(
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
            "w",
            {
                "weight": torch.randint(
                    256,
                    (hidden_size, hidden_size // 2),
                    dtype=FP4_DTYPE,
                    device=self.device,
                ),
                "wscale_swizzled": torch.randn(hidden_size, hidden_size // 16).to(
                    dtype=FP8_DTYPE, device=self.device
                ),
                "wscale": torch.tensor([500], dtype=torch.float32, device=self.device),
                "scale": torch.tensor([0.002], dtype=torch.float32, device=self.device),
            },
        )
224
225
226
227
228

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        """Forward pass that creates the pattern to be fused."""
        attn_output = self.attn(q, k, v)
        quant_output, output_block_scale = scaled_fp4_quant(
229
230
231
232
233
234
235
236
237
238
            attn_output, 1 / self.w["scale"]
        )
        return cutlass_scaled_fp4_mm(
            a=quant_output,
            b=self.w["weight"],
            block_scale_a=output_block_scale,
            block_scale_b=self.w["wscale_swizzled"],
            alpha=self.w["scale"] * self.w["wscale"],
            out_dtype=attn_output.dtype,
        )
239
240


241
242
PATTERN_TEST_MODELS_FP8: list[tuple[str, type]] = []
PATTERN_TEST_MODELS_FP4: list[tuple[str, type]] = []
243
244
HEADS: list[tuple[int, int]] = []
SPLIT_ATTENTION: list[bool] = []
245
246
BACKENDS_FP8: list[AttentionBackendEnum] = []
BACKENDS_FP4: list[AttentionBackendEnum] = []
247

248
if current_platform.is_cuda():
249
    HEADS = [(64, 8), (40, 8)]
250
    PATTERN_TEST_MODELS_FP8 = [
251
        (
252
            "RedHatAI/Meta-Llama-3.1-8B-FP8",
253
            TestAttentionFp8StaticQuantPatternModel,
254
255
        )
    ]
256
    PATTERN_TEST_MODELS_FP4 = [
257
        (
258
            "nvidia/Llama-3.1-8B-Instruct-NVFP4",
259
            TestAttentionNvfp4QuantPatternModel,
260
        )
261
    ]
262
263
    BACKENDS_FP8 = [AttentionBackendEnum.TRITON_ATTN, AttentionBackendEnum.FLASHINFER]
    BACKENDS_FP4 = [AttentionBackendEnum.FLASHINFER]
264

265
elif current_platform.is_rocm():
266
    HEADS = [(32, 8), (40, 8)]
267
    PATTERN_TEST_MODELS_FP8 = [
268
269
        ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)
    ]
270
    BACKENDS_FP8 = [
271
272
273
        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
        AttentionBackendEnum.ROCM_ATTN,
        AttentionBackendEnum.TRITON_ATTN,
274
    ]
275
276
277


@pytest.mark.parametrize("num_qo_heads, num_kv_heads", HEADS)
278
@pytest.mark.parametrize("head_size", [128])
279
280
281
@pytest.mark.parametrize(
    "batch_size", [7, 256, 533] if current_platform.is_cuda() else [8]
)
282
283
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize(
284
285
    "backend, model_name, model_class, custom_ops",
    # Test attention+quant_fp8 fusion with custom and torch impls of QuantFP8
286
287
288
289
290
    list(
        flat_product(
            BACKENDS_FP8, PATTERN_TEST_MODELS_FP8, ["+quant_fp8", "-quant_fp8"]
        )
    )
291
    # quant_fp4 only has the custom impl
292
    + list(flat_product(BACKENDS_FP4, PATTERN_TEST_MODELS_FP4, [""])),
293
294
295
296
)
@pytest.mark.skipif(
    not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"
)
297
@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
298
299
300
301
302
303
def test_attention_quant_pattern(
    num_qo_heads: int,
    num_kv_heads: int,
    head_size: int,
    batch_size: int,
    dtype: torch.dtype,
304
    custom_ops: str,
305
306
    model_name: str,
    model_class: type[AttentionQuantPatternModel],
307
    backend: AttentionBackendEnum,
308
    dist_init,
309
310
    monkeypatch,
    use_fresh_inductor_cache,
311
):
312
    """Test AttentionStaticQuantPattern fusion pass"""
313
314
    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

315
    if backend == AttentionBackendEnum.FLASHINFER and (
316
317
        not current_platform.is_device_capability((10, 0)) or not has_flashinfer()
    ):
318
        # This also captures the FP4 case
319
        pytest.skip("FlashInfer attn fusion requires Blackwell and flashinfer")
320

321
    custom_ops_list = custom_ops.split(",") if custom_ops else []
322

323
    device = torch.device("cuda:0")
324
    torch.set_default_dtype(dtype)
325
326
    torch.manual_seed(42)

327
328
329
330
331
    model_config = ModelConfig(
        model=model_name,
        max_model_len=2048,
        dtype=dtype,
    )
332
    vllm_config = VllmConfig(
333
334
335
336
337
        model_config=model_config,
        scheduler_config=SchedulerConfig(
            max_num_seqs=1024,
            max_model_len=model_config.max_model_len,
            is_encoder_decoder=model_config.is_encoder_decoder,
338
339
        ),
        compilation_config=CompilationConfig(
340
            mode=CompilationMode.VLLM_COMPILE,
341
            custom_ops=custom_ops_list,
342
        ),
343
        cache_config=CacheConfig(cache_dtype="fp8"),
344
        attention_config=AttentionConfig(backend=backend),
345
    )
346
347

    # Create test inputs
348
349
350
    q = torch.randn(batch_size, num_qo_heads * head_size, dtype=dtype, device=device)
    k = torch.randn(batch_size, num_kv_heads * head_size, dtype=dtype, device=device)
    v = torch.randn(batch_size, num_kv_heads * head_size, dtype=dtype, device=device)
351
352
353
354
355
356
357
358

    # Mark first dimension as dynamic for realistic testing
    torch._dynamo.mark_dynamic(q, 0)
    torch._dynamo.mark_dynamic(k, 0)
    torch._dynamo.mark_dynamic(v, 0)

    # Run model directly without compilation and fusion
    vllm_config_unfused = copy.deepcopy(vllm_config)
359
360
361
362
363
364
365
366
367
368
369
370
    with (
        set_current_vllm_config(vllm_config_unfused),
        set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused),
    ):
        model_unfused = model_class(
            num_qo_heads=num_qo_heads,
            num_kv_heads=num_kv_heads,
            head_size=head_size,
            kv_cache_dtype=FP8_DTYPE,
            device=device,
            vllm_config=vllm_config_unfused,
        )
371
        model_unfused = model_unfused.to(device)
372
        result_unfused_0 = model_unfused(q, k, v)  # noqa: F841  HACK: See #131044
373
374

        forward_ctx = get_forward_context()
375
        forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)
376

377
378
        # Run model directly without fusion
        # Still compile so query QuantFP8 has closer numerics
379
380
        compiled_unfused = torch.compile(model_unfused, fullgraph=True)
        result_unfused = compiled_unfused(q, k, v)
381
382
383

    # Run model with attn fusion enabled
    vllm_config.compilation_config.pass_config = PassConfig(
384
        fuse_attn_quant=True, eliminate_noops=True
385
386
387
388
389
390
391
392
393
394
395
396
397
398
    )
    with (
        set_current_vllm_config(vllm_config),
        set_forward_context(attn_metadata=None, vllm_config=vllm_config),
    ):
        model_fused = model_class(
            num_qo_heads=num_qo_heads,
            num_kv_heads=num_kv_heads,
            head_size=head_size,
            kv_cache_dtype=FP8_DTYPE,
            device=device,
            vllm_config=vllm_config,
            w=model_unfused.w,
        )
399
400
401
        model_fused = model_fused.to(device)

        forward_ctx = get_forward_context()
402
        forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)
403
404
405

        # Create test backend with fusion passes enabled
        noop_pass = NoOpEliminationPass(vllm_config)
406
407
408
409
        attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
        cleanup_pass = PostCleanupPass(vllm_config)

        test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
410
411
        # HACK: See https://github.com/vllm-project/vllm/issues/31044
        result_fused_0 = model_fused(q, k, v)  # noqa: F841
412
413

        # Compile model with fusion enabled
414
        compiled_fused = torch.compile(
415
416
            model_fused, backend=test_backend, fullgraph=True
        )
417
        assert compiled_fused.attn._o_scale_float is None
418

419
        result_fused = compiled_fused(q, k, v)
420

421
        if backend == AttentionBackendEnum.FLASHINFER:
422
423
424
425
            # With the Flashinfer backend after the 1st round of the forward
            # pass, output quant scale should be loaded into the attn layer's
            # _o_scale_float, the 2nd round should reuse the loaded
            # _o_scale_float
426
427
            assert compiled_fused.attn._o_scale_float is not None
            result_fused_2 = compiled_fused(q, k, v)
428

429
            assert compiled_fused.attn._o_scale_float is not None
430

431
432
433
            torch.testing.assert_close(
                result_unfused, result_fused_2, atol=1e-2, rtol=1e-2
            )
434
435

    # Check attn fusion support
436
    quant_key: QuantKey = model_class.quant_key
437
    attn_fusion_supported = [
438
439
        layer.impl.fused_output_quant_supported(quant_key)
        for key, layer in vllm_config.compilation_config.static_forward_context.items()
440
    ]
441
442
443
444
445
446
447
448
449
450
451
452
453
    assert sum(attn_fusion_supported) == len(attn_fusion_supported), (
        "All layers should support attention fusion"
    )

    # Check quantization ops in the graph before and after fusion
    quant_op = (
        torch.ops.aten.reciprocal
        if "-quant_fp8" in custom_ops_list
        else QUANT_OPS[quant_key]
    )

    # Note: for fp8, fully_replaced=False because query quant ops remain in graph.
    # Only output quant ops are fused into attention.
454
    test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
455

456
457
458
    # access the underlying `AttnFusionPass` on the `LazyInitPass`
    assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)

459
460
    # Check attention ops in the graph before and after fusion
    attn_nodes_pre = list(find_op_nodes(ATTN_OP, test_backend.graph_pre_pass))
461
    attn_nodes_post = list(find_op_nodes(ATTN_OP, test_backend.graph_post_pass))
462
463

    assert len(attn_nodes_pre) > 0, "Should have attention nodes before fusion"
464
    assert len(attn_nodes_pre) == len(attn_nodes_post), (
465
        "Should have same number of attention nodes before and after fusion"
466
467
    )
    assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
468
        "Attention should not have output_scale before fusion"
469
470
    )
    assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
471
        "Attention should have output_scale after fusion"
472
    )
473

474
    assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
475
        "Attention should not have output_block_scale before fusion"
476
    )
477
478
479
480
481
482
483
484
485
486
487

    kv_cache_dummy_dep_pre_is_none = (
        attn_nodes_pre[0].kwargs.get("kv_cache_dummy_dep") is None
    )
    kv_cache_dummy_dep_post_is_none = (
        attn_nodes_post[0].kwargs.get("kv_cache_dummy_dep") is None
    )
    assert not (kv_cache_dummy_dep_pre_is_none ^ kv_cache_dummy_dep_post_is_none), (
        "The kv_cache_dummy_dep should be consistent before and after fusion"
    )

488
    if quant_key.dtype == FP8_DTYPE:
489
        assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
490
            "Attention should not have output_block_scale after FP8 fusion"
491
        )
492
    elif quant_key.dtype == FP4_DTYPE:
493
494
        assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
            "Attention should have output_block_scale after FP4 fusion"
495
        )
496

497
    # Check that results are close
498
    torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)