test_attention_backends.py 25 KB
Newer Older
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for v1 attention backends without GPUModelRunner dependency."""
4

5
from functools import partial
6
7
8

import pytest
import torch
9
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
10

11
12
13
14
15
from tests.v1.attention.utils import (
    BatchSpec,
    create_common_attn_metadata,
    create_standard_kv_cache_spec,
    create_vllm_config,
16
    try_backend_includes_kv_cache_update,
17
    try_get_attention_backend,
18
)
19
20
from vllm.config import ModelConfig
from vllm.platforms import current_platform
21
from vllm.utils.math_utils import cdiv
22
23
24
25
26
from vllm.utils.torch_utils import (
    STR_DTYPE_TO_TORCH_DTYPE,
    is_torch_equal_or_newer,
    set_random_seed,
)
27
from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
28
from vllm.v1.attention.backends.registry import AttentionBackendEnum
29
30
31
from vllm.v1.attention.backends.utils import (
    set_kv_cache_layout,
)
32
33
34
from vllm.v1.kv_cache_interface import FullAttentionSpec

BACKENDS_TO_TEST = [
35
36
37
38
39
    AttentionBackendEnum.FLASH_ATTN,
    AttentionBackendEnum.FLASHINFER,
    AttentionBackendEnum.FLEX_ATTENTION,
    AttentionBackendEnum.TRITON_ATTN,
    AttentionBackendEnum.TREE_ATTN,
40
    "FLEX_ATTENTION_SLOW",
41
42
43
44
45
46
]

# Remove flashinfer from the list if it's not available
try:
    import flashinfer  # noqa: F401
except ImportError:
47
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER)
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66


def _convert_dtype_to_torch(dtype):
    """Convert ModelDType to torch.dtype."""
    if isinstance(dtype, str):
        if dtype == "auto":
            return torch.float16  # Default dtype for testing
        elif dtype in STR_DTYPE_TO_TORCH_DTYPE:
            return STR_DTYPE_TO_TORCH_DTYPE[dtype]
        else:
            raise ValueError(f"Unknown dtype: {dtype}")
    elif isinstance(dtype, torch.dtype):
        return dtype
    else:
        raise ValueError(f"Unknown dtype: {dtype}")


# Define common batch configurations
BATCH_SPECS = {
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    "small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
    "small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
    "mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
    "medium_decode": BatchSpec(
        seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
        query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
    ),
    "medium_prefill": BatchSpec(
        seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]
    ),
    "mixed_medium": BatchSpec(
        seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]
    ),
    "large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
    "large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
82
83
84
    "mixed_large": BatchSpec(
        seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]
    ),
85
86
    "single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
    "single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
87
88
89
90
91
92
93
    # encoder-only
    "small_encoder_prefill": BatchSpec(
        seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]
    ),
    "medium_encoder_prefill": BatchSpec(
        seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]
    ),
94
95
96
97
}


def create_and_prepopulate_kv_cache(
98
99
100
101
102
103
104
105
106
107
108
    k_contexts: list[torch.Tensor],
    v_contexts: list[torch.Tensor],
    block_size: int,
    num_kv_heads: int,
    head_size: int,
    dtype: torch.dtype,
    device: torch.device,
    num_blocks: int,
    common_attn_metadata: CommonAttentionMetadata,
    randomize_blocks: bool = True,
) -> torch.Tensor:
109
    """Create and prepopulate a KV cache with context data.
110

111
112
113
114
115
116
117
118
119
120
121
    Args:
        k_contexts: List of key context tensors for each sequence
        v_contexts: List of value context tensors for each sequence
        seq_lens: List of sequence lengths
        block_size: Size of each block
        num_kv_heads: Number of KV heads
        head_size: Size of each head
        dtype: Data type for the cache
        device: Device to create the cache on
        num_blocks: Total number of blocks in the cache
        block_table: Block table tensor to populate
122
        randomize_blocks: Whether to randomly permute blocks
123
                          or use sequential order
124

125
126
127
128
    Returns:
        Tuple of (kv_cache, updated_block_table)
    """
    batch_size = len(k_contexts)
129
    seq_lens = common_attn_metadata.seq_lens.cpu()
130
131
132
133
    query_lens = (
        common_attn_metadata.query_start_loc_cpu[1:]
        - common_attn_metadata.query_start_loc_cpu[:-1]
    )
134
    context_lens = seq_lens - query_lens
135
136
137
138
    block_table = common_attn_metadata.block_table_tensor
    slot_mapping = common_attn_metadata.slot_mapping

    # Create KV cache
139
    kv_cache = torch.zeros(
140
141
        2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
    )
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
    kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)

    # Populate the cache with the context tokens
    # Start from block_id=1 since block_id=0 is considered the null block
    start_block_idx = 1
    for i in range(batch_size):
        k_context, v_context = k_contexts[i], v_contexts[i]
        start = start_block_idx * block_size
        end = start + k_context.shape[0]
        kv_cache_flat[0, start:end, ...] = k_context
        kv_cache_flat[1, start:end, ...] = v_context

        # Stay block aligned and allocate enough blocks for the new tokens
        start_block_idx += cdiv(int(seq_lens[i]), block_size)

    blocks_end = start_block_idx

    # Permute the context blocks (excluding block 0 which is null)
    if randomize_blocks:
Matthew Bonanni's avatar
Matthew Bonanni committed
161
162
        # Random permutation starting from block 1
        perm = torch.randperm(blocks_end - 1) + 1
163
    else:
Matthew Bonanni's avatar
Matthew Bonanni committed
164
165
        # Sequential order starting from block 1
        perm = torch.arange(1, blocks_end)
166
167

    inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
Matthew Bonanni's avatar
Matthew Bonanni committed
168
169
    # Add 1 to account for starting from block 1
    inv_perm[1:] = torch.argsort(perm) + 1
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
    kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]

    # Construct the right block table
    # Start from block_id=1 since block_id=0 is considered the null block
    start_block_idx = 1
    for i in range(batch_size):
        num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
        start = start_block_idx
        end = start + num_blocks_for_seq
        block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
        start_block_idx += num_blocks_for_seq

        # Create a realistic slot mapping that corresponds to the block table
    for i in range(batch_size):
        token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
        block_indices = token_offsets // block_size
        token_inter_block_offsets = token_offsets % block_size
        start = common_attn_metadata.query_start_loc_cpu[i]
        end = common_attn_metadata.query_start_loc_cpu[i + 1]
        slot_mapping[start:end] = block_table[
190
191
            i, block_indices
        ] * block_size + token_inter_block_offsets.to(device)
192
193
194
195
196
197
198
199
200
201
202
203

    return kv_cache


class MockAttentionLayer:
    """A mock attention layer for testing."""

    def __init__(self, device: torch.device):
        self._q_scale = torch.tensor(1.0, device=device)
        self._k_scale = torch.tensor(1.0, device=device)
        self._v_scale = torch.tensor(1.0, device=device)
        # Add float versions for flashinfer
204
        self._q_scale_float = 1.0
205
206
207
208
        self._k_scale_float = 1.0
        self._v_scale_float = 1.0


209
def run_attention_backend(
210
    backend: AttentionBackendEnum,
211
212
213
214
215
216
217
218
219
    kv_cache_spec: FullAttentionSpec,
    layer_names: list[str],
    vllm_config,
    device: torch.device,
    common_attn_metadata: CommonAttentionMetadata,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    kv_cache: torch.Tensor,
220
    attn_type: AttentionType = AttentionType.DECODER,
221
    sliding_window: int | None = None,
222
) -> torch.Tensor:
223
224
    """Run attention computation using the specified backend's AttentionImpl."""

225
226
227
228
229
    # Handle special case for FLEX_ATTENTION_SLOW
    actual_backend = backend

    use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
    if backend == "FLEX_ATTENTION_SLOW":
230
        actual_backend = AttentionBackendEnum.FLEX_ATTENTION
231
232
        use_direct_block_mask = False

233
    builder_cls, impl_cls = try_get_attention_backend(actual_backend)
234
235

    # Mock flashinfer's get_per_layer_parameters if needed
236
    if actual_backend == AttentionBackendEnum.FLASHINFER:
237
238
        import unittest.mock

239
        from vllm.v1.attention.backends.utils import PerLayerParameters
240

241
        def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
242
243
244
            # Return mock parameters for a single layer
            head_size = vllm_config.model_config.get_head_size()
            return {
245
                layer_name: PerLayerParameters(
246
247
                    window_left=-1,  # No sliding window
                    logits_soft_cap=0.0,  # No soft cap
248
                    sm_scale=1.0 / (head_size**0.5),  # Standard scale
249
                )
250
                for layer_name in layer_names
251
252
253
            }

        with unittest.mock.patch(
254
255
256
257
            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
            mock_get_per_layer_parameters,
        ):
            builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
258
259
260
261
262
263
            attn_metadata = builder.build(
                common_prefix_len=0,
                common_attn_metadata=common_attn_metadata,
            )
    else:
        # Build metadata
264
        builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
265
        if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
266
            builder.direct_build = use_direct_block_mask
267
268
269
270
271
272
273
        attn_metadata = builder.build(
            common_prefix_len=0,
            common_attn_metadata=common_attn_metadata,
        )

    # Instantiate implementation
    num_heads = vllm_config.model_config.get_num_attention_heads(
274
275
        vllm_config.parallel_config
    )
276
    num_kv_heads = vllm_config.model_config.get_num_kv_heads(
277
278
        vllm_config.parallel_config
    )
279
280
281
282
283
284
285
286
    head_size = vllm_config.model_config.get_head_size()
    scale = 1.0 / (head_size**0.5)
    impl = impl_cls(
        num_heads=num_heads,
        head_size=head_size,
        scale=scale,
        num_kv_heads=num_kv_heads,
        alibi_slopes=None,
287
        sliding_window=sliding_window,
288
        attn_type=attn_type,
289
290
291
292
293
294
295
296
297
298
        kv_cache_dtype="auto",
    )

    # Create mock layer and output buffer
    mock_layer = MockAttentionLayer(device)
    output = torch.empty_like(query)

    # Run forward pass
    # NOTE: The query, key, and value are already shaped correctly
    # in the calling test function.
299
300
301
302
    if not try_backend_includes_kv_cache_update(actual_backend):
        impl.do_kv_cache_update(
            mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
        )
303
304
305
    output = impl.forward(
        mock_layer, query, key, value, kv_cache, attn_metadata, output=output
    )
306
307
308
309

    return output


310
311
312
def _test_backend_correctness(
    batch_spec: BatchSpec,
    model: str,
313
    backend_to_test: list[AttentionBackendEnum | str],
314
315
    mask_mod,
    *,
316
    attn_type: AttentionType = AttentionType.DECODER,
317
318
319
    block_size: int = 16,
    atol: float = 1e-2,
    rtol: float = 1e-2,
320
    tensor_parallel_size: int = 1,
321
):
322
323
324
325
326
327
328
329
330
331
332
333
334
335
    """
    Test that all backends produce similar outputs to a reference implementation
    using torch.nn.functional.scaled_dot_product_attention.

    This test works by:
    1. Generating a batch of sequences with specified context and query lengths.
    2. Computing a ground-truth attention output using torch.sdpa on
       contiguous Q, K, and V tensors.
    3. Simulating vLLM's paged KV cache: It takes the context portion of the
       K/V tensors and manually places them into a paged buffer according to
       the test's (randomly generated) block table.
    4. Running each vLLM attention backend with the new queries and the
       simulated paged KV cache.
    5. Comparing the vLLM backend's output to the ground-truth SDPA output.
336
337
338
339
340

    Note: When tensor_parallel_size > 1, we simulate the head partitioning
    by overriding the model config to use fewer heads, without requiring
    multiple GPUs. This tests that backends work correctly with different
    head counts.
341
    """
342
    set_random_seed(42)
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360

    hf_config_override = None
    if tensor_parallel_size > 1:
        from vllm.config import ModelConfig

        temp_config = ModelConfig(model=model, max_model_len=1)
        original_num_heads = temp_config.hf_text_config.num_attention_heads
        original_num_kv_heads = getattr(
            temp_config.hf_text_config, "num_key_value_heads", None
        )
        hf_config_override = {
            "num_attention_heads": original_num_heads // tensor_parallel_size,
        }
        if original_num_kv_heads is not None:
            hf_config_override["num_key_value_heads"] = max(
                1, original_num_kv_heads // tensor_parallel_size
            )

361
362
    vllm_config = create_vllm_config(
        model_name=model,
363
        tensor_parallel_size=1,  # Always use TP=1 to avoid multi-GPU requirements
364
365
366
        max_model_len=max(batch_spec.seq_lens),
        block_size=block_size,
        num_gpu_blocks=8192,
367
        hf_config_override=hf_config_override,
368
    )
369
370
371
372
373
374
375
376
377
    device = torch.device("cuda:0")

    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

    # 1. Setup
    batch_size = batch_spec.batch_size
    seq_lens = batch_spec.seq_lens
    query_lens = batch_spec.query_lens
    num_q_heads = vllm_config.model_config.get_num_attention_heads(
378
379
        vllm_config.parallel_config
    )
380
    num_kv_heads = vllm_config.model_config.get_num_kv_heads(
381
382
        vllm_config.parallel_config
    )
383
    head_size = vllm_config.model_config.get_head_size()
384
    sliding_window = vllm_config.model_config.get_sliding_window()
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
    dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
    block_size = vllm_config.cache_config.block_size
    scale = 1.0 / (head_size**0.5)

    # 2. Generate data and compute SDPA reference output
    all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
    all_sdpa_outputs = []
    k_contexts, v_contexts = [], []

    for i in range(batch_size):
        s_len = seq_lens[i]
        q_len = query_lens[i]
        context_len = s_len - q_len

        # Generate Q, K, V for the whole sequence to be used in SDPA
400
401
402
        q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
        k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
        v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
403
404
405
406
407
408
409
410
411

        # SDPA expects (N, H, L, D), so unsqueeze batch and permute
        q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
        k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
        v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)

        if num_q_heads != num_kv_heads:
            assert num_q_heads % num_kv_heads == 0, (
                f"num_q_heads ({num_q_heads}) must be divisible by "
412
413
                f"num_kv_heads ({num_kv_heads})"
            )
414
415
416
417
418
419
420
            repeats = num_q_heads // num_kv_heads
            k_sdpa_in = k_sdpa_in.repeat_interleave(repeats, dim=1)
            v_sdpa_in = v_sdpa_in.repeat_interleave(repeats, dim=1)

        # Create causal mask: query token i attends to positions 0 to
        #  (context_len + i)
        kv_len = s_len
421
422

        final_mask_mod = partial(mask_mod, context_len=context_len)
423
424
425
426
427
428
429
430
431
432
433
        block_mask = create_block_mask(
            final_mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device
        )
        sdpa_out_i = flex_attention(
            q_sdpa_in,
            k_sdpa_in,
            v_sdpa_in,
            block_mask=block_mask,
            scale=scale,
            enable_gqa=True,
        )
434

435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
        all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))

        # Inputs for vLLM backends are just the new tokens
        all_q_vllm.append(q)
        all_k_vllm.append(k_full[context_len:])
        all_v_vllm.append(v_full[context_len:])

        # Contextual K/V data used to populate the paged cache
        k_contexts.append(k_full[:context_len])
        v_contexts.append(v_full[:context_len])

    query_vllm = torch.cat(all_q_vllm, dim=0)
    key_vllm = torch.cat(all_k_vllm, dim=0)
    value_vllm = torch.cat(all_v_vllm, dim=0)
    sdpa_output = torch.cat(all_sdpa_outputs, dim=0)

    common_attn_metadata = create_common_attn_metadata(
452
453
        batch_spec, vllm_config.cache_config.block_size, device
    )
454
455
456
    if attn_type == AttentionType.ENCODER_ONLY:
        # For encoder-only, all tokens are prefill tokens
        common_attn_metadata.causal = False
457
458
459
460
461
462
463
464
465
466
467
468

    # 3. Simulate Paged KV Cache and a realistic slot_mapping
    kv_cache = create_and_prepopulate_kv_cache(
        k_contexts=k_contexts,
        v_contexts=v_contexts,
        block_size=block_size,
        num_kv_heads=num_kv_heads,
        head_size=head_size,
        dtype=dtype,
        device=device,
        num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
        common_attn_metadata=common_attn_metadata,
469
470
        randomize_blocks=True,
    )
471
472
473
474

    # 4. Run vLLM backends and compare
    # Note: flex_attention has known Triton kernel compatibility issues
    # with test infrastructures
475
    for backend_name in backend_to_test:
476
477
        # FlashAttentionm + FlexAttention:
        #   [2, num_blocks, block_size, num_kv_heads, head_size]
478
        # FlashInfer + Triton:
479
480
481
        #   [num_blocks, 2, block_size, num_kv_heads, head_size]
        # Select the appropriate KV cache format for each backend
        kv_cache_for_backend = kv_cache
482
        reset_kv_cache_layout = False
483
484
485
486
        if backend_name in (
            AttentionBackendEnum.FLASHINFER,
            AttentionBackendEnum.TRITON_ATTN,
        ):
487
488
            kv_cache_for_backend = kv_cache.transpose(0, 1)

489
        if backend_name == AttentionBackendEnum.FLASHINFER:
490
            # For FlashInfer default to HND layout and
491
492
493
            kv_cache_for_backend = (
                kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
            )
494
            set_kv_cache_layout("HND")
495
            reset_kv_cache_layout = True
496
        elif backend_name == AttentionBackendEnum.TRITON_ATTN:
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
            kv_cache_for_backend = kv_cache_for_backend.contiguous()

        try:
            backend_output = run_attention_backend(
                backend_name,
                kv_cache_spec,
                ["placeholder"],
                vllm_config,
                device,
                common_attn_metadata,
                query_vllm,
                key_vllm,
                value_vllm,
                kv_cache_for_backend,
                sliding_window=sliding_window,
512
                attn_type=attn_type,
513
514
515
516
            )
        finally:
            if reset_kv_cache_layout:
                set_kv_cache_layout(None)
517
518
519
520

        # Check shape and dtype consistency
        assert backend_output.shape == sdpa_output.shape, (
            f"[{backend_name}] shape {backend_output.shape} != "
521
522
            f"SDPA shape {sdpa_output.shape}"
        )
523
524
        assert backend_output.dtype == sdpa_output.dtype, (
            f"[{backend_name}] dtype {backend_output.dtype} != "
525
526
            f"SDPA dtype {sdpa_output.dtype}"
        )
527
528

        assert torch.isfinite(backend_output).all(), (
529
530
            f"[{backend_name}] produced non-finite values"
        )
531
532

        # Check numerical similarity
533
        def error_msg(msg: str, backend_name: str):
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
            return f"[{backend_name}] output differs from SDPA baseline. {msg}"

        torch.testing.assert_close(
            backend_output,
            sdpa_output,
            rtol=rtol,
            atol=atol,
            msg=partial(error_msg, backend_name=backend_name),
        )


@pytest.mark.parametrize(
    "batch_spec_name",
    [
        "small_decode",
        "small_prefill",
        "mixed_small",
        "medium_decode",
        "medium_prefill",
        "mixed_medium",
        "large_decode",
        "large_prefill",
        "single_decode",
        "single_prefill",
    ],
)
560
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
561
562
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
def test_causal_backend_correctness(
563
    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
564
):
565
566
567
568
569
570
571
572
573
574
575
576
577
    """Test backend's correctness with causal attention."""

    def causal_mask_mod(
        b: torch.Tensor,
        h: torch.Tensor,
        q_idx: torch.Tensor,
        kv_idx: torch.Tensor,
        *,
        context_len: int,
    ):
        return (q_idx + context_len) >= kv_idx

    batch_spec = BATCH_SPECS[batch_spec_name]
578
    LARGE_BLOCK_BACKENDS = (
579
580
581
        [AttentionBackendEnum.FLEX_ATTENTION]
        if is_torch_equal_or_newer("2.9.0.dev0")
        else []
582
    )
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

    if current_platform.is_rocm():
        SMALL_BLOCK_BACKENDS = [
            x
            for x in BACKENDS_TO_TEST
            if (
                x not in LARGE_BLOCK_BACKENDS
                and x is not AttentionBackendEnum.FLASH_ATTN
            )
        ]
    else:
        SMALL_BLOCK_BACKENDS = [
            x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
        ]

598
599
600
601
602
603
604
    _test_backend_correctness(
        batch_spec,
        model,
        SMALL_BLOCK_BACKENDS,
        causal_mask_mod,
        tensor_parallel_size=tensor_parallel_size,
    )
605
606
607

    # Fast FlexAttention needs to run with block_size=128
    if LARGE_BLOCK_BACKENDS:
608
        _test_backend_correctness(
609
610
611
612
613
614
            batch_spec,
            model,
            LARGE_BLOCK_BACKENDS,
            causal_mask_mod,
            block_size=128,
            tensor_parallel_size=tensor_parallel_size,
615
        )
616
617


618
619
620
621
622
623
624
625
626
627
628
629
630
631
if current_platform.is_rocm():
    # FLASH_ATTN is not supported on ROCm
    SLIDING_WINDOW_BACKENDS_TO_TEST = [
        AttentionBackendEnum.FLEX_ATTENTION,
        AttentionBackendEnum.TRITON_ATTN,
        "FLEX_ATTENTION_SLOW",
    ]
else:
    SLIDING_WINDOW_BACKENDS_TO_TEST = [
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.FLEX_ATTENTION,
        AttentionBackendEnum.TRITON_ATTN,
        "FLEX_ATTENTION_SLOW",
    ]
632
633


634
635
@pytest.mark.parametrize(
    "batch_spec_name",
636
637
638
639
640
641
642
643
    [
        "small_decode",
        "small_prefill",
        "mixed_medium",
        "large_decode",
        "large_prefill",
        "mixed_large",
    ],
644
)
645
@pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
646
647
648
649
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
def test_sliding_window_backend_correctness(
    batch_spec_name: str, model: str, tensor_parallel_size: int
):
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
    """Test backend's correctness with sliding window attention."""

    def sliding_window_mask_mod(
        b: torch.Tensor,
        h: torch.Tensor,
        q_idx: torch.Tensor,
        kv_idx: torch.Tensor,
        *,
        context_len: int,
        sliding_window: int,
    ):
        causal_mask = q_idx + context_len >= kv_idx
        window_mask = q_idx + context_len - kv_idx < sliding_window
        return causal_mask & window_mask

    batch_spec = BATCH_SPECS[batch_spec_name]
666
    model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
667
    sliding_window = model_config.get_sliding_window()
668
669
670
    sliding_window_mask_mod_fn = partial(
        sliding_window_mask_mod, sliding_window=sliding_window
    )
671

672
    LARGE_BLOCK_BACKENDS = (
673
674
675
        [AttentionBackendEnum.FLEX_ATTENTION]
        if is_torch_equal_or_newer("2.9.0.dev0")
        else []
676
    )
677
    SMALL_BLOCK_BACKENDS = [
678
        x for x in SLIDING_WINDOW_BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
679
    ]
680
    _test_backend_correctness(
681
682
683
684
685
        batch_spec,
        model,
        SMALL_BLOCK_BACKENDS,
        sliding_window_mask_mod_fn,
        tensor_parallel_size=tensor_parallel_size,
686
    )
687
688
689

    # Fast FlexAttention needs to run with block_size=128
    if LARGE_BLOCK_BACKENDS:
690
691
692
693
694
695
        _test_backend_correctness(
            batch_spec,
            model,
            LARGE_BLOCK_BACKENDS,
            sliding_window_mask_mod_fn,
            block_size=128,
696
            tensor_parallel_size=tensor_parallel_size,
697
        )
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739


@pytest.mark.parametrize(
    "batch_spec_name",
    [
        "small_encoder_prefill",
        "medium_encoder_prefill",
    ],
)
@pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 2])
def test_sliding_window_encoder_backend_correctness(
    batch_spec_name: str, model: str, tensor_parallel_size: int
):
    """Test backend's correctness with sliding window attention."""

    def bidi_sliding_window_mask_mod(
        b: torch.Tensor,
        h: torch.Tensor,
        q_idx: torch.Tensor,
        kv_idx: torch.Tensor,
        *,
        context_len: int,
        sliding_window: int,
    ):
        return torch.abs(q_idx + context_len - kv_idx) < sliding_window

    batch_spec = BATCH_SPECS[batch_spec_name]
    model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
    sliding_window = model_config.get_sliding_window()
    sliding_window_mask_mod_fn = partial(
        bidi_sliding_window_mask_mod, sliding_window=sliding_window
    )

    _test_backend_correctness(
        batch_spec,
        model,
        SLIDING_WINDOW_BACKENDS_TO_TEST,
        sliding_window_mask_mod_fn,
        attn_type=AttentionType.ENCODER_ONLY,
        tensor_parallel_size=tensor_parallel_size,
    )