Skip to content

vllm.v1.attention.backends.mamba1_attn

Mamba1AttentionBackend

Bases: AttentionBackend

Source code in vllm/v1/attention/backends/mamba1_attn.py 
17
18
19
20
class Mamba1AttentionBackend(AttentionBackend):
    @staticmethod
    def get_builder_cls() -> type["Mamba1AttentionMetadataBuilder"]:
        return Mamba1AttentionMetadataBuilder

get_builder_cls staticmethod 

get_builder_cls() -> type[Mamba1AttentionMetadataBuilder]
Source code in vllm/v1/attention/backends/mamba1_attn.py 
18
19
20
@staticmethod
def get_builder_cls() -> type["Mamba1AttentionMetadataBuilder"]:
    return Mamba1AttentionMetadataBuilder

Mamba1AttentionMetadata dataclass

Source code in vllm/v1/attention/backends/mamba1_attn.py 
23
24
25
26
27
28
29
30
31
32
33
@dataclass
class Mamba1AttentionMetadata:
    query_start_loc: torch.Tensor
    context_lens_tensor: torch.Tensor
    state_indices_tensor: torch.Tensor
    has_initial_states: torch.Tensor | None
    num_prefills: int
    num_prefill_tokens: int
    num_decodes: int
    num_decode_tokens: int
    num_padded_decodes: int

context_lens_tensor instance-attribute 

context_lens_tensor: Tensor

has_initial_states instance-attribute 

has_initial_states: Tensor | None

num_decode_tokens instance-attribute 

num_decode_tokens: int

num_decodes instance-attribute 

num_decodes: int

num_padded_decodes instance-attribute 

num_padded_decodes: int

num_prefill_tokens instance-attribute 

num_prefill_tokens: int

num_prefills instance-attribute 

num_prefills: int

query_start_loc instance-attribute 

query_start_loc: Tensor

state_indices_tensor instance-attribute 

state_indices_tensor: Tensor

__init__ 

__init__(
    query_start_loc: Tensor,
    context_lens_tensor: Tensor,
    state_indices_tensor: Tensor,
    has_initial_states: Tensor | None,
    num_prefills: int,
    num_prefill_tokens: int,
    num_decodes: int,
    num_decode_tokens: int,
    num_padded_decodes: int,
) -> None

Mamba1AttentionMetadataBuilder

Bases: BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]

Source code in vllm/v1/attention/backends/mamba1_attn.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
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
77
78
79
80
81
82
83
84
85
86
class Mamba1AttentionMetadataBuilder(
    BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]
):
    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> Mamba1AttentionMetadata:
        query_start_loc = common_attn_metadata.query_start_loc

        state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
        context_lens_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
            query_start_loc.device
        )

        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
            split_decodes_and_prefills(
                common_attn_metadata, decode_threshold=self.reorder_batch_threshold
            )
        )

        has_initial_states = None
        padded_decodes = num_decodes

        if num_prefills > 0:
            has_initial_states = context_lens_tensor > 0
        elif (
            num_decodes > 0
            and num_decodes <= self.decode_cudagraph_max_bs
            and self.compilation_config.full_cuda_graph
        ):
            state_indices_for_decode = state_indices_tensor[:num_decodes]
            padded_decodes = self.vllm_config.pad_for_cudagraph(num_decodes)
            self.state_indices_tensor[:num_decodes].copy_(
                state_indices_for_decode, non_blocking=True
            )
            state_indices_tensor = self.state_indices_tensor[:padded_decodes]
            state_indices_tensor[num_decodes:] = PAD_SLOT_ID

        return Mamba1AttentionMetadata(
            query_start_loc=query_start_loc,
            context_lens_tensor=context_lens_tensor,
            has_initial_states=has_initial_states,
            state_indices_tensor=state_indices_tensor,
            num_prefills=num_prefills,
            num_prefill_tokens=num_prefill_tokens,
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_padded_decodes=padded_decodes,
        )

build 

build(
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> Mamba1AttentionMetadata
Source code in vllm/v1/attention/backends/mamba1_attn.py 
39
40
41
42
43
44
45
46
47
48
49
50
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
77
78
79
80
81
82
83
84
85
86
def build(
    self,
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> Mamba1AttentionMetadata:
    query_start_loc = common_attn_metadata.query_start_loc

    state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
    context_lens_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
        query_start_loc.device
    )

    num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
        split_decodes_and_prefills(
            common_attn_metadata, decode_threshold=self.reorder_batch_threshold
        )
    )

    has_initial_states = None
    padded_decodes = num_decodes

    if num_prefills > 0:
        has_initial_states = context_lens_tensor > 0
    elif (
        num_decodes > 0
        and num_decodes <= self.decode_cudagraph_max_bs
        and self.compilation_config.full_cuda_graph
    ):
        state_indices_for_decode = state_indices_tensor[:num_decodes]
        padded_decodes = self.vllm_config.pad_for_cudagraph(num_decodes)
        self.state_indices_tensor[:num_decodes].copy_(
            state_indices_for_decode, non_blocking=True
        )
        state_indices_tensor = self.state_indices_tensor[:padded_decodes]
        state_indices_tensor[num_decodes:] = PAD_SLOT_ID

    return Mamba1AttentionMetadata(
        query_start_loc=query_start_loc,
        context_lens_tensor=context_lens_tensor,
        has_initial_states=has_initial_states,
        state_indices_tensor=state_indices_tensor,
        num_prefills=num_prefills,
        num_prefill_tokens=num_prefill_tokens,
        num_decodes=num_decodes,
        num_decode_tokens=num_decode_tokens,
        num_padded_decodes=padded_decodes,
    )