utils.py 35.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
"""Kernel test utils"""

5
6
import itertools
import random
7
from collections.abc import Sequence
8
from numbers import Number
9
from typing import Any, NamedTuple
10
from unittest.mock import patch
11

12
import pytest
13
import torch
14
from torch._prims_common import TensorLikeType
15

bnellnm's avatar
bnellnm committed
16
from tests.kernels.quant_utils import native_w8a8_block_matmul
17
from vllm.attention import AttentionType
18
from vllm.model_executor.layers.activation import SiluAndMul
19
20
21
22
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
from vllm.utils import (
    STR_BACKEND_ENV_VAR,
)
23
from vllm.utils.torch_utils import make_tensor_with_pad
24

25
26
# For now, disable "test_aot_dispatch_dynamic" since there are some
# bugs related to this test in PyTorch 2.4.
27
DEFAULT_OPCHECK_TEST_UTILS: tuple[str, ...] = (
28
29
30
31
32
    "test_schema",
    "test_autograd_registration",
    "test_faketensor",
)

33
ALL_OPCHECK_TEST_UTILS: tuple[str, ...] = (
34
35
36
37
38
39
    "test_schema",
    "test_autograd_registration",
    "test_faketensor",
    "test_aot_dispatch_dynamic",
)

40

41
class QKVInputs(NamedTuple):
42
    """
43
    Data structure for representing unpacked attention inputs,
44
45
46
47
    query/key/values and their sequence lengths.

    Attributes:

48
        * {query,key,value}: unpacked (batch_size x padded_seq_len x
49
50
51
                             num_heads x head_size) attention inputs
        * q_seq_lens: query sequence lengths list
        * kv_seq_lens: shared key/value sequence lengths list
52
    """
53
54
55
56

    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
57
58
    q_seq_lens: list[int]
    kv_seq_lens: list[int]
59
60
61


class QKVO(NamedTuple):
62
    """
63
    Data structure for representing unpacked attention inputs,
64
65
66
67
    alongside unpacked known-correct attention output

    Attributes:

68
        * qkv: unpacked (batch_size x padded_seq_len x
69
                             num_heads x head_size) attention inputs
70
        * ideal_output: unpacked (batch_size x padded_seq_len x
71
                        num_heads x head_size) known-correct attention output
72
    """
73
74
75
76
77
78

    qkv: QKVInputs
    ideal_output: torch.Tensor


class PackedQKVInputs(NamedTuple):
79
    """
80
81
82
83
    Data structure for representing packed attention inputs

    Attributes:

84
        * {query,key,value}: packed (number_of_tokens x num_heads
85
86
87
88
89
90
                             x head_size) attention inputs
        * q_start_loc_list: list of query start locations within packed tensor
        * kv_start_loc_list: shared list of key/value start locations within
                             packed tensor
        * q_seq_lens: query sequence lengths list
        * kv_seq_lens: shared key/value sequence lengths list
91
    """
92
93
94
95

    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
96
97
98
99
    q_start_loc_list: list[int] | None
    kv_start_loc_list: list[int] | None
    q_seq_lens: list[int] | None
    kv_seq_lens: list[int] | None
100
101
102


class PackedQKVO(NamedTuple):
103
    """
104
    Data structure for representing packed attention inputs,
105
106
107
108
    alongside packed known-correct attention output

    Attributes:

109
        * packed_qkv: packed (number_of_tokens x num_heads
110
                      x head_size) attention inputs
111
        * ideal_output: packed (number_of_tokens x num_heads
112
                        x head_size) known-correct attention output
113
    """
114

115
    packed_qkv: PackedQKVInputs | None
116
117
118
119
    ideal_output: torch.Tensor


class KVMemoryMap(NamedTuple):
120
    """
121
122
123
124
125
126
    Data structure for encapsulating KV cache memory mapping.

    Attributes:

        * block_tables: KV cache block tables
        * slot_mapping: mapping of sequence offset to physical address
127
    """
128
129
130
131
132
133

    block_tables: torch.Tensor
    slot_mapping: torch.Tensor


class PhaseTestParameters(NamedTuple):
134
    """
135
136
137
138
139
140
    Data structure for encapsulating the test parameters
    for a given test "phase" (prefill or decode phase) and attention
    scenario (encoder, decoder-self, encoder/decoder-cross)

    Attributes:

141
        * packed_qkvo: packed (number_of_tokens x num_heads
142
143
144
145
                       x head_size) attention inputs & known-correct
                       output
        * kv_mmap: KV cache memory mapping, specific to this test phase &
                   attention scenario
146
    """
147
148

    packed_qkvo: PackedQKVO
149
    kv_mmap: KVMemoryMap | None
150
151
152


def maybe_make_int_tensor(
153
154
    _list: list[int] | None,
    device: torch.device | str,
155
) -> torch.Tensor:
156
    """
157
158
159
160
161
162
    Convert Python int list to a 1D int torch.Tensor on `device`

    Returns:

    * If _list is not None: 1D int torch.Tensor on `device`
    * None otherwise
163
164
165
166
    """
    return (
        None if _list is None else torch.tensor(_list, dtype=torch.int, device=device)
    )
167
168
169


def maybe_make_long_tensor(
170
171
    _list: list[int] | None,
    device: torch.device | str,
172
) -> torch.Tensor:
173
    """
174
175
176
177
178
179
    Convert Python int list to a 1D long torch.Tensor on `device`

    Returns:

    * If _list is not None: 1D long torch.Tensor on `device`
    * None otherwise
180
181
182
183
    """
    return (
        None if _list is None else torch.tensor(_list, dtype=torch.long, device=device)
    )
184
185


186
def maybe_max(_list: list | None) -> Number | None:
187
    """
188
189
190
191
    Returns:

    * If _list is not None: max(_list)
    * None otherwise
192
    """
193
194
195
196
197
198
199
    return None if _list is None else max(_list)


def make_causal_mask(
    q_max_seq_len: int,
    kv_max_seq_len: int,
) -> torch.Tensor:
200
    """
201
202
203
    Create a q_max_seq_len x kv_max_seq_len causal mask

    Arguments:
204

205
206
207
208
209
210
    * q_max_seq_len: query max seq len
    * kv_max_seq_len: key/value max seq len

    Returns:

    * 2D tensor, q_max_seq_len x kv_max_seq_len
211
    """
212
213
214
215

    # Create a matrix where entry (i, j) is True if i >= j
    mask = torch.triu(torch.ones(q_max_seq_len, kv_max_seq_len), diagonal=1)
    # Replace True with float('-inf') and False with 0
216
    mask = mask.masked_fill(mask == 1, float("-inf")).masked_fill(mask == 0, 0.0)
217
218
219
    return mask


220
221
222
223
def override_backend_env_variable(
    mpatch: pytest.MonkeyPatch, backend_name: str
) -> None:
    """
224
225
226
227
228
229
230
231
    Override the environment variable indicating the vLLM backend temporarily,
    using pytest monkeypatch to ensure that the env vars get
    reset once the test context exits.

    Arguments:

    * mpatch: pytest monkeypatch instance
    * backend_name: attention backend name to force
232
    """
233
    mpatch.setenv(STR_BACKEND_ENV_VAR, backend_name)
234
235


236
237
238
239
240
def ref_masked_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: float,
241
242
243
    custom_mask: torch.Tensor | None = None,
    q_seq_lens: list | None = None,
    kv_seq_lens: list | None = None,
244
245
) -> torch.Tensor:
    """
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
    "Golden" masked attention reference. Supports two types of masking:

    * Basic attention mask, utilizing {q,kv}_seq_lens args to mask out
      padding elements
    * Custom attention mask, which can force an arbitrary mask tensor, i.e.
      causal

    Arguments:

    * query: batch_size x q_padded_seq_len x num_heads x head_size
    * key: batch_size x kv_padded_seq_len x num_heads x head_size
    * value: batch_size x kv_padded_seq_len x num_heads x head_size
    * scale: Attention scale factor
    * custom_mask: custom attention mask; good place to inject a causal
      attention mask
    * q_seq_lens: list of unpadded query seq_lens for each batch index
    * kv_seq_lens: list of unpadded key/value seq_lens for each batch index

    Returns:

    * Attention result, batch_size x q_padded_seq_len x num_heads x head_size
267
    """
268
269
270
271
272

    assert q_seq_lens is not None
    assert kv_seq_lens is not None

    batch_size = query.shape[0]
273
274
    assert len(q_seq_lens) == batch_size
    assert len(kv_seq_lens) == batch_size
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301

    attn_weights = scale * torch.einsum("bqhd,bkhd->bhqk", query, key).float()

    # Basic attention mask, derived from seq lens
    if (q_seq_lens is not None) or (kv_seq_lens is not None):
        attn_mask = torch.zeros_like(attn_weights)
        if q_seq_lens is not None:
            for bdx, plen in enumerate(q_seq_lens):
                attn_mask[bdx, :, plen:, :] = -torch.inf
        if kv_seq_lens is not None:
            for bdx, plen in enumerate(kv_seq_lens):
                attn_mask[bdx, :, :, plen:] = -torch.inf

        attn_weights = attn_weights + attn_mask.float()

    # Custom attention mask
    if custom_mask is not None:
        attn_weights = attn_weights + custom_mask.float()

    attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
    out = torch.einsum("bhqk,bkhd->bqhd", attn_weights, value)
    return out


def make_qkv(
    batch_size: int,
    max_q_seq_len: int,
302
    max_kv_seq_len: int | None,
303
304
    num_heads: int,
    head_size: int,
305
306
    device: torch.device | str,
    force_kv_seq_lens: list[int] | None = None,
307
308
    attn_type: AttentionType = AttentionType.ENCODER_DECODER,
    force_max_len: bool = False,
309
) -> tuple[QKVInputs, QKVInputs, QKVInputs]:
310
    """
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
    Construct QKV test tensors for self- and cross-attention.

    Generates three query/key/value triplets:

    * "Baseline" query/key/value (for input to reference attention function)
    * "Prefill" query/key/value (last sequence offset zero'd out, for use as
      input to prefill kernel)
    * "Decode" query/key/value (only the last sequence offset  from baseline,
      for use as input to decode kernel)

    Each Q/K/V triplet is associated with a list of q seqlens and a list of k/v
    seqlens

    Arguments:

    * batch_size
    * max_q_seq_len: max query seq len
    * max_kv_seq_len: max key/value seq len
    * num_heads
    * head_size
331
332
333
    * is_encoder_decoder_attn: if True, query seqlen may differ from
      key/value seqlen (as is often the case for cross-attention);
      o/w, query/key/value seqlens match at each batch index
334
335
336
337
338
339
340
341
342
343
344
345
346
      (max_kv_seq_len is unused)
    * force_kv_seq_lens: if not None, overrides kv sequence lengths
    * attn_type: encoder, decoder self, or enc/dec cross attention
    * force_max_len: if True, all query seqlens are max_q_seq_len; o/w query
      seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens
      and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False
    * device: CPU or CUDA device

    Returns:

    * Overall QKVInputs structure (containing full unpacked Q/K/V tensors)
    * Prefill QKVInputs structure (containing all but the last sequence offset)
    * Decode QKVInputs structure (containing all only the last sequence offset)
347
    """
348
349
350
351

    if force_max_len:
        q_seq_lens = [max_q_seq_len for _ in range(batch_size)]
    else:
352
        q_seq_lens = [random.randint(2, max_q_seq_len) for _ in range(batch_size)]
353
354
355
356
357
358
359
360
361
362
363
364
    kv_seq_lens = None
    if force_kv_seq_lens is not None:
        kv_seq_lens = force_kv_seq_lens
    elif attn_type != AttentionType.ENCODER_DECODER:
        # K,V seq lens match Q for self-attention
        kv_seq_lens = q_seq_lens
    else:
        # K,V seq lens are distinct from Q seq lens & random
        assert max_kv_seq_len is not None
        if force_max_len:
            kv_seq_lens = [max_kv_seq_len] * batch_size
        else:
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
            kv_seq_lens = [random.randint(2, max_kv_seq_len) for _ in range(batch_size)]

    query = torch.rand((batch_size, max_q_seq_len, num_heads, head_size)).to(device)
    key = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)
    value = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)

    prefill_query = torch.zeros((batch_size, max_q_seq_len, num_heads, head_size)).to(
        device
    )
    prefill_key = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
        device
    )
    prefill_value = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
        device
    )

    decode_query = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
382
    decode_key = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
383
    decode_value = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
384

385
    for bdx, (q_seq_len, kv_seq_len) in enumerate(zip(q_seq_lens, kv_seq_lens)):
386
387
388
389
        query[bdx, q_seq_len:, :, :] = 0
        key[bdx, kv_seq_len:, :, :] = 0
        value[bdx, kv_seq_len:, :, :] = 0

390
391
392
393
394
395
396
397
398
399
400
401
402
        prefill_query[bdx, 0 : (q_seq_len - 1), :, :] = query[
            bdx, 0 : (q_seq_len - 1), :, :
        ]
        prefill_key[bdx, 0 : (kv_seq_len - 1), :, :] = key[
            bdx, 0 : (kv_seq_len - 1), :, :
        ]
        prefill_value[bdx, 0 : (kv_seq_len - 1), :, :] = value[
            bdx, 0 : (kv_seq_len - 1), :, :
        ]

        decode_query[bdx, :, :, :] = query[bdx, (q_seq_len - 1) : q_seq_len, :, :]
        decode_key[bdx, :, :, :] = key[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
        decode_value[bdx, :, :, :] = value[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
403
404
405
406
407
408
409
410
411
412
413
414
415

    prefill_q_seq_lens = [plen - 1 for plen in q_seq_lens]
    prefill_kv_seq_lens = [plen - 1 for plen in kv_seq_lens]

    decode_q_seq_lens = [1 for _ in q_seq_lens]
    decode_kv_seq_lens = [1 for _ in kv_seq_lens]

    return (
        QKVInputs(
            query,  # Overall QKV inputs
            key,
            value,
            q_seq_lens,
416
417
            kv_seq_lens,
        ),
418
419
420
421
422
        QKVInputs(
            prefill_query,  # Prefill subset of QKV sequences
            prefill_key,
            prefill_value,
            prefill_q_seq_lens,
423
424
            prefill_kv_seq_lens,
        ),
425
426
427
428
429
        QKVInputs(
            decode_query,  # Decode subset of KV sequences
            decode_key,
            decode_value,
            decode_q_seq_lens,
430
431
432
            decode_kv_seq_lens,
        ),
    )
433
434
435


def pack_tensor(
436
    unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str
437
438
) -> tuple[torch.Tensor, list[int]]:
    """
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
    Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an
    unpadded number_of_tokens x num_heads x head_size tensor, where
    number_of_tokens = sum(seq_lens)

    Arguments:

    * unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size
    * seq_lens: list of token counts for each seq
    * device: CPU or CUDA device

    Returns

    * packed_tensor: number_of_tokens x num_heads x head_size
    * start_loc_list: start idx of each batch elt in packed_tensor; [0] +
      list(itertools.accumulate(seq_lens))
454
    """
455
456
457
458
459
460
461
462

    num_tok = sum(seq_lens)
    num_heads = unpacked_tensor.shape[-2]
    head_size = unpacked_tensor.shape[-1]
    start_loc_list = [0] + list(itertools.accumulate(seq_lens))
    packed_tensor = torch.zeros((num_tok, num_heads, head_size), device=device)

    for bdx, (seq_len, start_loc) in enumerate(zip(seq_lens, start_loc_list)):
463
464
465
        packed_tensor[start_loc : (start_loc + seq_len), :, :] = unpacked_tensor[
            bdx, :seq_len, :, :
        ]
466
467
468
469

    return packed_tensor, start_loc_list


470
def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
471
    """
472
473
474
    Individually pack each of Q, K and V, each with dimensions batch_size x
    padded_seq_len x num_heads x head_size, into respective number_of_tokens x
    num_heads x head_size tensors.
475

476
477
478
479
480
481
482
483
484
485
486
487
488
489
    For Q, number_of_tokens = sum(q_seq_lens).

    For K and V, number_of_tokens = sum(kv_seq_lens)

    Arguments:

    * qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size)
           attention inputs
    * device: CPU or CUDA device

    Returns

    * Packed (number_of_tokens x num_heads x head_size) QKV inputs
      derived from unpacked inputs
490
    """
491
492
493
494
495

    if qkv.query is None:
        packed_query = None
        q_start_loc_list = None
    else:
496
497
498
499
        packed_query, q_start_loc_list = pack_tensor(
            qkv.query, qkv.q_seq_lens, device=device
        )
    packed_key, kv_start_loc_list = pack_tensor(qkv.key, qkv.kv_seq_lens, device=device)
500
501
    packed_value, _ = pack_tensor(qkv.value, qkv.kv_seq_lens, device=device)
    return PackedQKVInputs(
502
503
504
505
        packed_query,
        packed_key,
        packed_value,
        q_start_loc_list,
506
507
        kv_start_loc_list,
        (None if q_start_loc_list is None else qkv.q_seq_lens),
508
509
        qkv.kv_seq_lens,
    )
510
511


512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
def make_alibi_bias(
    alibi_slopes: torch.Tensor,
    num_kv_heads: int,
    dtype: torch.dtype,
    seq_lens: list[int],
) -> list[Any]:
    """Create ALiBi biases compatible with xFormers attention tests."""
    from xformers.ops.fmha.attn_bias import LowerTriangularMaskWithTensorBias

    if alibi_slopes is None:
        return [None for _ in seq_lens]

    attn_biases: list[Any] = []
    num_heads = alibi_slopes.shape[0]
    assert num_heads >= num_kv_heads, (
527
528
        "ALiBi slopes expect at least as many heads as KV heads"
    )
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548

    for seq_len in seq_lens:
        bias = torch.arange(seq_len, dtype=dtype, device=alibi_slopes.device)
        bias = bias[None, :] - bias[:, None]

        padded_len = (seq_len + 7) // 8 * 8
        bias_tensor = torch.empty(
            1,
            num_heads,
            seq_len,
            padded_len,
            device=alibi_slopes.device,
            dtype=dtype,
        )[:, :, :, :seq_len].copy_(bias)
        bias_tensor.mul_(alibi_slopes[:, None, None])
        attn_biases.append(LowerTriangularMaskWithTensorBias(bias_tensor))

    return attn_biases


549
def _make_metadata_tensors(
550
551
552
553
    seq_lens: list[int] | None,
    context_lens: list[int] | None,
    encoder_seq_lens: list[int] | None,
    device: torch.device | str,
554
555
556
557
558
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    Any,
    Any,
559
    torch.Tensor | None,
560
561
    torch.Tensor,
    torch.Tensor,
562
    int | None,
563
564
]:
    """
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
    Build scalar & tensor values required to build attention metadata structure.

    Arguments:

    * seq_lens: list of token-counts for each decoder input seq
    * context_lens: list of context length values for each seq
    * encoder_seq_lens: list of token-counts for each encoder input seq
    * device: CPU or CUDA device

    Returns:

    * seq_lens_tensor: decoder seq_lens list, as tensor
    * context_lens_tensor: context_lens list, as tensor
    * max_context_len: max(context_lens)
    * max_seq_len: max(seq_lens)
    * seq_start_loc: start idx of each sequence
581
582
    * encoder_seq_lens_tensor: encoder seq_lens list, as tensor
    * encoder_seq_start_loc: start idx of each encoder sequence
583
    * max_encoder_seq_len: encoder seq_lens list, as tensor
584
    """
585
586
587
588
589
590
    seq_lens_tensor = maybe_make_int_tensor(seq_lens, device)
    context_lens_tensor = maybe_make_int_tensor(context_lens, device)
    max_context_len = maybe_max(context_lens)
    max_seq_len = maybe_max(seq_lens)

    encoder_seq_lens_tensor = maybe_make_int_tensor(encoder_seq_lens, device)
591
    max_encoder_seq_len = None if encoder_seq_lens is None else max(encoder_seq_lens)
592
593
594

    seq_start_loc = None

595
    if seq_lens_tensor is not None:
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
        seq_start_loc = torch.zeros(
            seq_lens_tensor.shape[0] + 1,
            dtype=torch.int32,
            device=seq_lens_tensor.device,
        )
        torch.cumsum(
            seq_lens_tensor, dim=0, dtype=seq_start_loc.dtype, out=seq_start_loc[1:]
        )

    encoder_seq_start_loc = torch.zeros(
        encoder_seq_lens_tensor.shape[0] + 1,
        dtype=torch.int32,
        device=encoder_seq_lens_tensor.device,
    )
    torch.cumsum(
        encoder_seq_lens_tensor,
        dim=0,
        dtype=encoder_seq_start_loc.dtype,
        out=encoder_seq_start_loc[1:],
    )

    return (
        seq_lens_tensor,
        context_lens_tensor,
        max_context_len,
        max_seq_len,
        seq_start_loc,
        encoder_seq_lens_tensor,
        encoder_seq_start_loc,
        max_encoder_seq_len,
    )


def make_kv_cache(
    num_blocks: int,
    num_heads: int,
    head_size: int,
    block_size: int,
634
    device: torch.device | str,
635
636
637
638
    backend: str,
    default_val: float = 0.0,
) -> torch.Tensor:
    """
639
640
641
642
643
644
645
646
647
648
649
650
651
652
    Create a fake KV cache.

    Arguments:

    * num_blocks: number of blocks in the KV cache
    * num_heads: number of attention heads
    * head_size: head dimension
    * block_size: number of offsets within a block
    * device: CPU or CUDA device
    * default_val: initialization value for KV cache elements

    Returns:

    * kv_cache: 2 x num_blocks x (block_size * num_heads * head_size)
653
    *     for backend 'XFORMERS'
654
    * kv_cache: 2 x num_blocks x block_size x num_heads x head_size
655
    *     for backend 'FLASH_ATTN'
656
657
658
659
660
661
662
663
664
    """
    if backend == "XFORMERS":
        kv_cache = torch.rand((2, num_blocks, block_size * num_heads * head_size)).to(
            device
        )
    elif backend == "FLASH_ATTN":
        kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(
            device
        )
665
666
    else:
        raise ValueError(
667
668
            f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
        )
669
670
671
672
673
674
    if default_val is not None:
        kv_cache[:, :, :] = default_val
    return kv_cache


def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int:
675
    """
676
677
    Compute the minimum number of blocks required to hold num_tokens tokens,
    given block_size
678
    """
679
680
681
    return (num_tokens + block_size) // block_size


682
def make_empty_slot_mapping_tensor(device: torch.device | str):
683
684
685
    return maybe_make_long_tensor([], device)


686
def make_empty_block_tables_tensor(device: torch.device | str):
687
688
689
    return torch.tensor([], device=device)


690
691
692
def split_slot_mapping(
    slot_mapping_list: torch.Tensor,
    seq_lens: list[int],
693
    device: torch.device | str,
694
695
):
    """
696
697
698
699
700
    Split a slot mapping into valid prefill- and decode-phase slot mappings.

    Context:
    * Your goal is to test (1) prefill of N prompts, with prompt-lengths
      {K_i \\forall i \\in [0,N)}, followed by (2) decoding of a single token
701
      for all N prompts (N tokens total); the resultant sequence lengths
702
      after decode would be {K_i + 1 for i \\in [0,N)}
703
704
705
    * The test you want to do requires (1) having the prefill slot mapping
      for all tokens present during prefill, the number of which is
      M = \\sum_i{K_i}, and (2) having the decode slot mapping for all N
706
      decoded tokens
707
708

    This function consumes a single 1D slot mapping, which is the
709
710
711
712
713
    concatenation of N slot mappings each of length K_i + 1 (corresponding
    to the  sequence lengths after decode), with a total length of
    P = \\sum_i{K_i + 1} = M + N

    The prefill-phase slot mapping results from excising the (K_i + 1)-th entry
714
    from each of the N subsequences in the slot mapping (i.e. omitting the
715
716
717
718
719
720
    decoded token's mapping.)

    The N excised entries are appended to obtain the decode-phase slot mapping

    Arguments:

721
    * slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N
722
      post-decode sequences
723
    * seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
724
725
726
727
728
      description above)
    * device: cuda, cpu, etc.

    Returns:

729
    * prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
730
      reflecting all N prefill prompts
731
    * decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
732
      all N decoded tokens
733
    """
734
735
736
737
738
739

    prefill_slot_mapping = []
    decode_slot_mapping = []

    base_idx = 0
    for seq_len in seq_lens:
740
741
742
        prefill_slot_mapping.extend(
            slot_mapping_list[base_idx : (base_idx + seq_len - 1)]
        )
743
744
745
        decode_slot_mapping.append(slot_mapping_list[base_idx + seq_len - 1])
        base_idx += seq_len

746
747
748
749
    return (
        maybe_make_long_tensor(prefill_slot_mapping, device),
        maybe_make_long_tensor(decode_slot_mapping, device),
    )
750
751
752


def make_block_tables_slot_mapping(
753
754
    block_size: int,
    seq_lens: list[int],
755
    device: torch.device | str,
756
757
758
    block_base_addr: int = 0,
) -> tuple[torch.Tensor, list[int], int]:
    """
759
760
761
762
763
764
765
766
767
    Construct fake block tables & slot mappings.

    For a sequence with num_tokens tokens the minimum number
    of required KV cache blocks is

    num_blocks = (num_tokens + block_size) // block_size

    Then the minimum KV cache size in blocks is

768
    total_cache_blocks = sum(num_blocks for all seqs)
769
770
771
772
773
774
775
776

    Then, the blocktable mapping counts downward from

    block_base_addr + total_cache_blocks

    to

    block_base_addr
777

778
779
780
781
782
783
784
785
786
787
788
789
790
791

    The constructed block-tables and slot-mapping are sized to the
    lengths of the sequences in their entirety (as reflected by seq_lens),
    i.e. the total of prefill prompt tokens + decoded tokens.

    Arguments:

    * block_size: number of offsets per block
    * seq_lens: list of token-counts for each sequence
    * block_base_addr: the block table base address
    * device: CPU or CUDA device

    Return:

792
    * block_tables_tensor: block table for sequence
793
794
    * slot_mapping_list: slot mapping for sequence
    * max_block_idx: the highest block address within this block table
795
    """
796
797
798

    # Provision minimum number of KV cache blocks
    num_blocks_list = [
799
        _num_tokens_to_min_blocks(num_tokens, block_size) for num_tokens in seq_lens
800
801
802
803
804
805
806
807
808
809
810
811
    ]
    max_block_table_len = max(num_blocks_list)
    block_table_pad_tokens = 10

    block_tables = []
    slot_mapping_list = []
    # Compute uppermost address of block table
    total_cache_blocks = sum(num_blocks_list)
    block_base_idx = block_base_addr + total_cache_blocks
    max_block_idx = block_base_idx
    for sdx, num_tokens in enumerate(seq_lens):
        num_blocks = num_blocks_list[sdx]
812
        block_table = list(range(block_base_idx, block_base_idx - num_blocks, -1))
813
        for idx in range(num_tokens):
814
815
816
            mapping_value = (idx % block_size) + block_table[
                idx // block_size
            ] * block_size
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
            slot_mapping_list.append(mapping_value)

        block_base_idx -= num_blocks
        block_tables.append(block_table)

    block_tables_tensor = make_tensor_with_pad(
        block_tables,
        max_len=max_block_table_len + block_table_pad_tokens,
        pad=0,
        dtype=torch.int,
        device=device,
    )

    return (block_tables_tensor, slot_mapping_list, max_block_idx)


833
834
835
836
def assert_actual_matches_ideal(
    test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str
) -> None:
    """
837
838
839
840
841
842
843
    Assert that observed output matches the ideal output
    contained in the test parameters data structure.

    Arguments:

    * test_params: Test parameters including packed ideal output
    * output_under_test: actually observed output value
844
    """
845
    ideal_output = test_params.packed_qkvo.ideal_output
846
847
848
849
    if backend == "XFORMERS":
        torch.testing.assert_close(
            ideal_output, output_under_test.view_as(ideal_output)
        )
850

851
    elif backend == "FLASH_ATTN":
852
853
854
        # For FlashAttention override the accuracy thresholds to non default
        # values since we notice a higher difference between the ideal and
        # actual output.
855
856
857
        torch.testing.assert_close(
            ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
        )
858
859
    else:
        raise ValueError(
860
861
            f"Unknown backend value: '{backend}'. Expected 'XFORMERS' or 'FLASH_ATTN'."
        )
862
863


864
865
866
867
868
869
870
871
872
873
874
# Copied/modified from torch._refs.__init__.py
def fp8_allclose(
    a: TensorLikeType,
    b: TensorLikeType,
    rtol: float = 1e-05,
    atol: float = 1e-08,
    equal_nan: bool = False,
) -> bool:
    """
    Reference implementation of torch.allclose
    """
875
    torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol)
876
877
878

    return bool(
        torch.all(
879
880
881
882
883
            torch.isclose(
                a.double(), b.double(), rtol=rtol, atol=atol, equal_nan=equal_nan
            )
        ).item()
    )
884
885


886
887
888
# Marlin MoE test utils


889
def stack_and_dev(tensors: list[torch.Tensor]):
890
891
892
893
894
895
    dev = tensors[0].device
    return torch.stack(tensors, dim=0).to(dev)


def compute_max_diff(output, output_ref):
    return torch.mean(torch.abs(output - output_ref)) / torch.mean(
896
897
        torch.abs(output_ref)
    )
898
899


bnellnm's avatar
bnellnm committed
900
901
902
903
904
905
906
def torch_experts(
    a: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weight: torch.Tensor,
    topk_ids: torch.Tensor,
    global_num_experts: int = -1,
907
908
909
910
911
912
913
914
    b_bias1: torch.Tensor | None = None,
    b_bias2: torch.Tensor | None = None,
    expert_map: torch.Tensor | None = None,
    w1_scale: torch.Tensor | None = None,
    w2_scale: torch.Tensor | None = None,
    a1_scale: torch.Tensor | None = None,
    a2_scale: torch.Tensor | None = None,
    quant_dtype: torch.dtype | None = None,
bnellnm's avatar
bnellnm committed
915
    per_act_token_quant=False,
916
    block_shape: list[int] | None = None,
917
    apply_router_weights_on_input: bool = False,
bnellnm's avatar
bnellnm committed
918
) -> torch.Tensor:
919
920
921
922
923
    assert (
        global_num_experts == -1
        or (global_num_experts == w1.shape[0] and expert_map is None)
        or (expert_map is not None and global_num_experts == expert_map.shape[0])
    )
bnellnm's avatar
bnellnm committed
924
925

    M, K = a.shape
926
    topk = topk_ids.shape[1]
bnellnm's avatar
bnellnm committed
927

928
929
930
931
    if apply_router_weights_on_input:
        assert topk == 1
        a = a * topk_weight.to(a.dtype)

bnellnm's avatar
bnellnm committed
932
933
934
935
    a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)

    out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)

936
937
    if a1_scale:
        assert not per_act_token_quant and block_shape is None
938
939
940
    a, a_scale = moe_kernel_quantize_input(
        a, a1_scale, quant_dtype, per_act_token_quant, block_shape
    )
bnellnm's avatar
bnellnm committed
941
942
943

    num_experts = w1.shape[0]

944
    topk_ids = topk_ids.view(-1)
945
946
    if expert_map is not None:
        topk_ids = expert_map[topk_ids]
bnellnm's avatar
bnellnm committed
947

948
949
    f32 = torch.float32

bnellnm's avatar
bnellnm committed
950
    for i in range(num_experts):
951
952
        mask = topk_ids == i
        if mask.sum():
bnellnm's avatar
bnellnm committed
953
954
            if quant_dtype is None:
                tmp1 = a[mask] @ w1[i].transpose(0, 1)
955
956
                if b_bias1 is not None:
                    tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype)
bnellnm's avatar
bnellnm committed
957
958
                tmp2 = SiluAndMul()(tmp1)
                out[mask] = tmp2 @ w2[i].transpose(0, 1)
959
                if b_bias2 is not None:
960
                    out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype)
bnellnm's avatar
bnellnm committed
961
            elif block_shape is not None:
962
                # block quantized
963
964
965
966
967
968
969
970
                assert (
                    a_scale is not None
                    and w1_scale is not None
                    and w2_scale is not None
                )
                tmp1 = native_w8a8_block_matmul(
                    a[mask], w1[i], a_scale[mask], w1_scale[i], block_shape, out.dtype
                )
971
972
                if b_bias1 is not None:
                    tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype)
bnellnm's avatar
bnellnm committed
973
974
                tmp2 = SiluAndMul()(tmp1)
                tmp2, b_scale = moe_kernel_quantize_input(
975
976
                    tmp2, a2_scale, quant_dtype, per_act_token_quant, block_shape
                )
bnellnm's avatar
bnellnm committed
977

978
979
980
                out[mask] = native_w8a8_block_matmul(
                    tmp2, w2[i], b_scale, w2_scale[i], block_shape, out.dtype
                )
981
                if b_bias2 is not None:
982
                    out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype)
bnellnm's avatar
bnellnm committed
983
            else:
984
985
986
987
988
                assert (
                    a_scale is not None
                    and w1_scale is not None
                    and w2_scale is not None
                )
bnellnm's avatar
bnellnm committed
989
                scales = a_scale if a_scale.numel() == 1 else a_scale[mask]
990

bnellnm's avatar
bnellnm committed
991
992
                tmp1 = a[mask].to(f32) * scales
                w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
993
                tmp1 = (tmp1 @ w1_dq).to(out.dtype)
994
995
                if b_bias1 is not None:
                    tmp1 = tmp1 + b_bias1[i].view(1, -1).to(out.dtype)
996
997
998
999

                tmp2 = SiluAndMul()(tmp1).to(out.dtype)

                tmp2, b_scale = moe_kernel_quantize_input(
1000
1001
                    tmp2, a2_scale, quant_dtype, per_act_token_quant, block_shape
                )
1002
1003
1004
                assert b_scale is not None

                tmp2 = tmp2.to(f32) * b_scale
bnellnm's avatar
bnellnm committed
1005
1006
                w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
                out[mask] = (tmp2 @ w2_dq).to(out.dtype)
1007
                if b_bias2 is not None:
1008
                    out[mask] = out[mask] + b_bias2[i].view(1, -1).to(out.dtype)
bnellnm's avatar
bnellnm committed
1009

1010
1011
1012
    if apply_router_weights_on_input:
        return out
    else:
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
        return (
            (out.view(M, -1, w2.shape[1]).to(f32) * topk_weight.view(M, -1, 1))
            .sum(dim=1)
            .to(out.dtype)
        )


def torch_moe(
    a: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    score: torch.Tensor,
    topk: int,
1026
1027
    b_bias1: torch.Tensor | None = None,
    b_bias2: torch.Tensor | None = None,
1028
    global_num_experts: int = -1,
1029
    expert_map: torch.Tensor | None = None,
1030
) -> torch.Tensor:
1031
1032
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
    return torch_experts(
        a,
        w1,
        w2,
        topk_weight,
        topk_ids,
        global_num_experts,
        b_bias1,
        b_bias2,
        expert_map,
    )
1044
1045


1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
def torch_moe_single(a, w, score, topk):
    B, D = a.shape
    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
    out = torch.zeros(B * topk, w.shape[1], dtype=a.dtype, device=a.device)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    _, topk_ids = torch.topk(score, topk)
    topk_ids = topk_ids.view(-1)
    for i in range(w.shape[0]):
        mask = topk_ids == i
        if mask.sum():
            out[mask] = a[mask] @ w[i].transpose(0, 1)
    return (out.view(B, -1, w.shape[1])).sum(dim=1)


1060
1061
# A special version of op check that has a restricted default set of test_utils
# and a patched version of allclose that supports fp8 types.
1062
def opcheck(
1063
1064
1065
    op: torch._ops.OpOverload
    | torch._ops.OpOverloadPacket
    | torch._library.custom_ops.CustomOpDef,
1066
    args: tuple[Any, ...],
1067
    kwargs: dict[str, Any] | None = None,
1068
    *,
1069
    test_utils: str | Sequence[str] = ALL_OPCHECK_TEST_UTILS,
1070
1071
1072
    raise_exception: bool = True,
    cond: bool = True,
) -> dict[str, str]:
1073
    with patch("torch.allclose", new=fp8_allclose):
1074
1075
1076
1077
1078
1079
1080
        return (
            torch.library.opcheck(
                op, args, kwargs, test_utils=test_utils, raise_exception=raise_exception
            )
            if cond
            else {}
        )
1081
1082
1083
1084
1085


# For testing quantized linear kernels
def to_fp8(tensor: torch.Tensor):
    finfo = torch.finfo(torch.float8_e4m3fn)
1086
1087
1088
    return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to(
        dtype=torch.float8_e4m3fn
    )
1089
1090
1091
1092
1093
1094


def to_int8(tensor: torch.Tensor):
    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)


1095
1096
1097
1098
1099
1100
def baseline_scaled_mm(
    a: torch.Tensor,
    b: torch.Tensor,
    scale_a: torch.Tensor,
    scale_b: torch.Tensor,
    out_dtype: type[torch.dtype],
1101
    bias: torch.Tensor | None = None,
1102
) -> torch.Tensor:
1103
    # We treat N-dimensional group scaling as extended numpy-style broadcasting
1104
    # in numpy simply stretches dimensions with an extent of 1 to match
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
    # the target shape by repeating the data along that dimension (broadcasting)
    # , we extend these semantics to say if the extent of a dimension in the
    # source shape is not 1 and does not match the target shape we repeat each
    # element along that dimension src_shape[dim] // target_shape[dim] times
    # example if we have:
    #       a = [[1, 2], and target_shape = (2, 4)
    #            [3, 4]]
    # then we would expand a to:
    #       a = [[1, 1, 2, 2],
    #            [3, 3, 4, 4]]
1115
    # NOTE this function does not explicitly broadcast dimensions
1116
1117
1118
1119
1120
    # with an extent of 1, since this can be done implicitly by pytorch
    def group_broadcast(t, shape):
        for i, s in enumerate(shape):
            if t.shape[i] != s and t.shape[i] != 1:
                assert s % t.shape[i] == 0
1121
1122
1123
1124
1125
                t = (
                    t.unsqueeze(i + 1)
                    .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
                    .flatten(i, i + 1)
                )
1126
1127
1128
1129
1130
        return t

    scale_a = group_broadcast(scale_a, a.shape)
    scale_b = group_broadcast(scale_b, b.shape)

1131
1132
1133
    output = torch.mm(
        (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
    ).to(out_dtype)
1134

1135
1136
1137
1138
    if bias is not None:
        output = output + bias

    return output