Skip to content

vllm.model_executor.layers.mamba.linear.base

LinearAttention

Bases: PluggableLayer, MambaBase

Base class for Linear attention layer.

Source code in vllm/model_executor/layers/mamba/linear/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
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
class LinearAttention(PluggableLayer, MambaBase):
    """Base class for Linear attention layer."""

    def __init__(
        self, config: PretrainedConfig, vllm_config: VllmConfig, prefix: str = ""
    ):
        super().__init__()
        self.layer_idx = extract_layer_index(prefix)
        self.prefix = prefix
        self.model_config = vllm_config.model_config
        self.cache_config = vllm_config.cache_config
        self.quant_config = vllm_config.quant_config

        self.BLOCK = getattr(config, "block", 256)
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_hidden_layers = config.num_hidden_layers
        self.head_dim = (
            config.head_dim
            if hasattr(config, "head_dim")
            else config.hidden_size // self.num_heads
        )
        self.hidden_inner_size = self.head_dim * self.num_heads

        self.tp_size = get_tensor_model_parallel_world_size()
        self.tp_rank = get_tensor_model_parallel_rank()
        assert self.num_heads % self.tp_size == 0

    @property
    def mamba_type(self) -> MambaAttentionBackendEnum:
        return MambaAttentionBackendEnum.LINEAR

    def get_state_dtype(self) -> tuple[torch.dtype]:
        assert self.model_config is not None
        assert self.cache_config is not None
        return MambaStateDtypeCalculator.linear_attention_state_dtype(
            self.model_config.dtype,
            self.cache_config.mamba_cache_dtype,
        )

    def get_state_shape(self) -> tuple[tuple[int, int, int], ...]:
        return MambaStateShapeCalculator.linear_attention_state_shape(
            num_heads=self.num_heads, tp_size=self.tp_size, head_dim=self.head_dim
        )