Skip to content

vllm.attention.ops.rocm_aiter_mla

tags module-attribute 

tags = ()

aiter_mla_decode_fwd 

aiter_mla_decode_fwd(
    q: Tensor,
    kv_buffer: Tensor,
    o: Tensor,
    sm_scale: float,
    qo_indptr: Tensor,
    max_seqlen_qo: int,
    kv_indptr: Tensor | None = None,
    kv_indices: Tensor | None = None,
    kv_last_page_lens: Tensor | None = None,
    logit_cap: float = 0.0,
)
Source code in vllm/attention/ops/rocm_aiter_mla.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def aiter_mla_decode_fwd(
    q: torch.Tensor,
    kv_buffer: torch.Tensor,
    o: torch.Tensor,
    sm_scale: float,
    qo_indptr: torch.Tensor,
    max_seqlen_qo: int,
    kv_indptr: torch.Tensor | None = None,
    kv_indices: torch.Tensor | None = None,
    kv_last_page_lens: torch.Tensor | None = None,
    logit_cap: float = 0.0,
):
    torch.ops.vllm.rocm_aiter_mla_decode_fwd(
        q,
        kv_buffer.view(-1, 1, 1, q.shape[-1]),
        o,
        qo_indptr,
        max_seqlen_qo,
        kv_indptr,
        kv_indices,
        kv_last_page_lens,
        sm_scale=sm_scale,
        logit_cap=logit_cap,
    )

get_aiter_mla_metadata 

get_aiter_mla_metadata(
    max_batch_size: int,
    block_size: int,
    max_block_per_batch: int,
    device: device,
) -> tuple[Tensor, ...]
Source code in vllm/attention/ops/rocm_aiter_mla.py 
11
12
13
14
15
16
17
18
19
20
21
22
def get_aiter_mla_metadata(
    max_batch_size: int, block_size: int, max_block_per_batch: int, device: torch.device
) -> tuple[torch.Tensor, ...]:
    paged_kv_indices = torch.zeros(
        max_batch_size * max_block_per_batch, dtype=torch.int32, device=device
    )
    paged_kv_indptr = torch.zeros(max_batch_size + 1, dtype=torch.int32, device=device)
    paged_kv_last_page_lens = torch.full(
        (max_batch_size,), block_size, dtype=torch.int32
    )
    qo_indptr = torch.zeros(max_batch_size + 1, dtype=torch.int, device=device)
    return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens, qo_indptr

mla_decode_fwd_fake 

mla_decode_fwd_fake(
    q: Tensor,
    kv_buffer: Tensor,
    o: Tensor,
    qo_indptr: Tensor,
    max_seqlen_qo: int,
    kv_indptr: Tensor | None = None,
    kv_indices: Tensor | None = None,
    kv_last_page_lens: Tensor | None = None,
    sm_scale: float = 1.0,
    logit_cap: float = 0.0,
) -> None
Source code in vllm/attention/ops/rocm_aiter_mla.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
def mla_decode_fwd_fake(
    q: torch.Tensor,
    kv_buffer: torch.Tensor,
    o: torch.Tensor,
    qo_indptr: torch.Tensor,
    max_seqlen_qo: int,
    kv_indptr: torch.Tensor | None = None,
    kv_indices: torch.Tensor | None = None,
    kv_last_page_lens: torch.Tensor | None = None,
    sm_scale: float = 1.0,
    logit_cap: float = 0.0,
) -> None:
    pass

mla_decode_fwd_impl 

mla_decode_fwd_impl(
    q: Tensor,
    kv_buffer: Tensor,
    o: Tensor,
    qo_indptr: Tensor,
    max_seqlen_qo: int,
    kv_indptr: Tensor | None = None,
    kv_indices: Tensor | None = None,
    kv_last_page_lens: Tensor | None = None,
    sm_scale: float = 1.0,
    logit_cap: float = 0.0,
) -> None
Source code in vllm/attention/ops/rocm_aiter_mla.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def mla_decode_fwd_impl(
    q: torch.Tensor,
    kv_buffer: torch.Tensor,
    o: torch.Tensor,
    qo_indptr: torch.Tensor,
    max_seqlen_qo: int,
    kv_indptr: torch.Tensor | None = None,
    kv_indices: torch.Tensor | None = None,
    kv_last_page_lens: torch.Tensor | None = None,
    sm_scale: float = 1.0,
    logit_cap: float = 0.0,
) -> None:
    from aiter.mla import mla_decode_fwd

    mla_decode_fwd(
        q,
        kv_buffer.view(-1, 1, 1, q.shape[-1]),
        o,
        qo_indptr,
        kv_indptr,
        kv_indices,
        kv_last_page_lens,
        max_seqlen_qo,
        sm_scale=sm_scale,
        logit_cap=logit_cap,
    )