vllm.model_executor.models.zaya

Inference-only ZAYA1 model.

Classes:

• ZayaCCA –

  Compressed Convolutional Attention q/k/v projection.

• ZayaRouter –

  Picks an expert per token; an extra "skip" class bypasses the experts.

ZayaCCA

Bases: MambaBase, Module

Compressed Convolutional Attention q/k/v projection.

q/k mix over a causal conv of the raw projections, and half of v comes from the previous token, so each request keeps that history in Mamba state.

Source code in vllm/model_executor/models/zaya.py

class ZayaCCA(MambaBase, nn.Module):
    """Compressed Convolutional Attention q/k/v projection.

    q/k mix over a causal conv of the raw projections, and half of v comes from the
    previous token, so each request keeps that history in Mamba state.
    """

    def __init__(
        self,
        config: ZayaConfig,
        quant_config: QuantizationConfig | None,
        prefix: str,
    ):
        super().__init__()
        self.config = config
        self.prefix = prefix
        self.num_q_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.kv_groups = self.num_q_heads // self.num_kv_heads
        self.history_len = config.cca_time0 + config.cca_time1 - 2
        self.q_dim = self.num_q_heads * self.head_dim
        kv_dim = self.num_kv_heads * self.head_dim

        def linear(name: str, out_dim: int) -> ReplicatedLinear:
            return ReplicatedLinear(
                config.hidden_size,
                out_dim,
                bias=config.attention_bias,
                quant_config=quant_config,
                return_bias=False,
                prefix=f"{prefix}.{name}",
            )

        self.q_proj = linear("q_proj", self.q_dim)
        self.k_proj = linear("k_proj", kv_dim)
        self.v_proj_current = linear("v_proj_current", kv_dim // 2)
        self.v_proj_delayed = linear("v_proj_delayed", kv_dim // 2)
        channels = self.q_dim + kv_dim
        self.conv_qk_depthwise = nn.Conv1d(
            channels, channels, config.cca_time0, groups=channels
        )
        self.conv_qk_grouped = nn.Conv1d(
            channels,
            channels,
            config.cca_time1,
            groups=self.num_q_heads + self.num_kv_heads,
        )

        self.vllm_config = get_current_vllm_config()
        self.vllm_config.compilation_config.static_forward_context[prefix] = self
        self.kv_cache = (torch.tensor([]), torch.tensor([]))

    def forward(self, hidden_states: torch.Tensor, output: torch.Tensor) -> None:
        torch.ops.vllm.zaya_cca(hidden_states, output, self.prefix)

    def _conv(self, qk: torch.Tensor) -> torch.Tensor:
        return self.conv_qk_grouped(self.conv_qk_depthwise(qk))

    def project(self, hidden_states: torch.Tensor, output: torch.Tensor) -> None:
        num_tokens = hidden_states.shape[0]
        queries = self.q_proj(hidden_states)
        keys = self.k_proj(hidden_states)
        qk = torch.cat([queries, keys], dim=-1)
        delayed = self.v_proj_delayed(hidden_states)
        delayed_init = self.v_proj_delayed(
            hidden_states.new_zeros(hidden_states.shape[-1])
        )

        attn_metadata = get_forward_context().attn_metadata
        metadata = None
        if attn_metadata is not None:
            assert isinstance(attn_metadata, dict)
            metadata = attn_metadata.get(self.prefix)
        if metadata is None:  # Profile run
            conv_out = self._conv(F.pad(qk.T[None], (self.history_len, 0)))[0].T
            delayed_prev = delayed
        else:
            assert isinstance(metadata, ShortConvAttentionMetadata)
            conv_out, delayed_prev = self._mix_with_state(
                metadata, qk, delayed, delayed_init
            )

        heads = (num_tokens, -1, self.head_dim)
        query_residual = 0.5 * (
            queries.view(heads)
            + keys.view(heads).repeat_interleave(self.kv_groups, dim=1)
        )
        key_residual = query_residual.view(
            num_tokens, -1, self.kv_groups, self.head_dim
        ).mean(dim=2)
        query = conv_out[:, : self.q_dim].view(heads) + query_residual
        key = conv_out[:, self.q_dim :].view(heads) + key_residual
        value = torch.cat([self.v_proj_current(hidden_states), delayed_prev], dim=-1)
        output[:] = torch.cat([query, key, value.view(heads)], dim=1).flatten(1)

    def _mix_with_state(
        self,
        metadata: ShortConvAttentionMetadata,
        qk: torch.Tensor,
        delayed: torch.Tensor,
        delayed_init: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Causal conv of `qk` and the previous token's `delayed` value, reading
        and updating each request's state without syncing with the host.

        Decodes are one-token requests that always continue from their state.
        """
        device = qk.device
        num_decodes = metadata.num_decode_tokens
        num_tokens = num_decodes + metadata.num_prefill_tokens
        slots, starts, lens, has_state = [], [], [], []
        if num_decodes:
            assert metadata.state_indices_tensor_d is not None
            slots.append(metadata.state_indices_tensor_d[:, 0])
            starts.append(torch.arange(num_decodes, device=device))
            lens.append(torch.ones(num_decodes, dtype=torch.long, device=device))
            has_state.append(torch.ones(num_decodes, dtype=torch.bool, device=device))
        if metadata.num_prefills:
            query_start_loc = metadata.query_start_loc_p
            assert query_start_loc is not None
            assert metadata.state_indices_tensor_p is not None
            assert metadata.has_initial_states_p is not None
            slots.append(metadata.state_indices_tensor_p)
            starts.append(num_decodes + query_start_loc[:-1])
            lens.append(query_start_loc.diff())
            has_state.append(metadata.has_initial_states_p)
        # Padding rows of a captured graph use the unused null block
        slot = torch.cat(slots).clamp_min(NULL_BLOCK_ID)
        start, length = torch.cat(starts), torch.cat(lens)
        has_history = torch.cat(has_state)

        conv_state, delayed_state = self.kv_cache
        n = self.history_len
        history = conv_state[slot].to(qk.dtype) * has_history[:, None, None]
        first_delayed = torch.where(
            has_history[:, None], delayed_state[slot].to(delayed.dtype), delayed_init
        )

        # Tokens at least `n` into their request only see that request
        x = qk[:num_tokens]
        conv_out = qk.new_empty(qk.shape[0] + 1, qk.shape[1])
        conv_out[:num_tokens] = self._conv(F.pad(x.T[None], (n, 0)))[0].T
        # The first `n` tokens of each request also see its saved history; the
        # spare last row of `conv_out` absorbs those of requests shorter than `n`
        offset = torch.arange(n, device=device)
        first_tokens = start[:, None] + offset
        heads = self._conv(
            torch.cat([history, x[first_tokens.clamp(max=num_tokens - 1)].mT], -1)
        )
        valid = offset < length[:, None]
        conv_out[torch.where(valid, first_tokens, qk.shape[0])] = heads.mT

        delayed_prev = delayed.new_empty(delayed.shape)
        delayed_prev[1:num_tokens] = delayed[: num_tokens - 1]
        delayed_prev[start] = first_delayed

        # The new history is the last `n` inputs of the saved history and tokens
        last = length[:, None] - n + offset
        new_history = torch.where(
            (last >= 0)[..., None],
            x[(start[:, None] + last).clamp(min=0)],
            history.mT.gather(
                1, (last + n).clamp(0, n - 1)[..., None].expand(-1, -1, x.shape[1])
            ),
        )
        conv_state[slot] = new_history.mT.to(conv_state.dtype)
        delayed_state[slot] = delayed[start + length - 1].to(delayed_state.dtype)
        return conv_out[:-1], delayed_prev

    def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
        return _cca_state_shapes(self.config)

    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
        return _cca_state_dtypes(self.vllm_config)

    @property
    def is_kv_cache_tp_replicated(self) -> bool:
        return True

    @property
    def mamba_type(self) -> MambaAttentionBackendEnum:
        # CCA needs only the per-request slots, which is all ShortConv's metadata has
        return MambaAttentionBackendEnum.SHORT_CONV

_mix_with_state(metadata, qk, delayed, delayed_init)

Causal conv of qk and the previous token's delayed value, reading and updating each request's state without syncing with the host.

Decodes are one-token requests that always continue from their state.

Source code in vllm/model_executor/models/zaya.py

def _mix_with_state(
    self,
    metadata: ShortConvAttentionMetadata,
    qk: torch.Tensor,
    delayed: torch.Tensor,
    delayed_init: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Causal conv of `qk` and the previous token's `delayed` value, reading
    and updating each request's state without syncing with the host.

    Decodes are one-token requests that always continue from their state.
    """
    device = qk.device
    num_decodes = metadata.num_decode_tokens
    num_tokens = num_decodes + metadata.num_prefill_tokens
    slots, starts, lens, has_state = [], [], [], []
    if num_decodes:
        assert metadata.state_indices_tensor_d is not None
        slots.append(metadata.state_indices_tensor_d[:, 0])
        starts.append(torch.arange(num_decodes, device=device))
        lens.append(torch.ones(num_decodes, dtype=torch.long, device=device))
        has_state.append(torch.ones(num_decodes, dtype=torch.bool, device=device))
    if metadata.num_prefills:
        query_start_loc = metadata.query_start_loc_p
        assert query_start_loc is not None
        assert metadata.state_indices_tensor_p is not None
        assert metadata.has_initial_states_p is not None
        slots.append(metadata.state_indices_tensor_p)
        starts.append(num_decodes + query_start_loc[:-1])
        lens.append(query_start_loc.diff())
        has_state.append(metadata.has_initial_states_p)
    # Padding rows of a captured graph use the unused null block
    slot = torch.cat(slots).clamp_min(NULL_BLOCK_ID)
    start, length = torch.cat(starts), torch.cat(lens)
    has_history = torch.cat(has_state)

    conv_state, delayed_state = self.kv_cache
    n = self.history_len
    history = conv_state[slot].to(qk.dtype) * has_history[:, None, None]
    first_delayed = torch.where(
        has_history[:, None], delayed_state[slot].to(delayed.dtype), delayed_init
    )

    # Tokens at least `n` into their request only see that request
    x = qk[:num_tokens]
    conv_out = qk.new_empty(qk.shape[0] + 1, qk.shape[1])
    conv_out[:num_tokens] = self._conv(F.pad(x.T[None], (n, 0)))[0].T
    # The first `n` tokens of each request also see its saved history; the
    # spare last row of `conv_out` absorbs those of requests shorter than `n`
    offset = torch.arange(n, device=device)
    first_tokens = start[:, None] + offset
    heads = self._conv(
        torch.cat([history, x[first_tokens.clamp(max=num_tokens - 1)].mT], -1)
    )
    valid = offset < length[:, None]
    conv_out[torch.where(valid, first_tokens, qk.shape[0])] = heads.mT

    delayed_prev = delayed.new_empty(delayed.shape)
    delayed_prev[1:num_tokens] = delayed[: num_tokens - 1]
    delayed_prev[start] = first_delayed

    # The new history is the last `n` inputs of the saved history and tokens
    last = length[:, None] - n + offset
    new_history = torch.where(
        (last >= 0)[..., None],
        x[(start[:, None] + last).clamp(min=0)],
        history.mT.gather(
            1, (last + n).clamp(0, n - 1)[..., None].expand(-1, -1, x.shape[1])
        ),
    )
    conv_state[slot] = new_history.mT.to(conv_state.dtype)
    delayed_state[slot] = delayed[start + length - 1].to(delayed_state.dtype)
    return conv_out[:-1], delayed_prev

ZayaRouter

Bases: Module

Picks an expert per token; an extra "skip" class bypasses the experts.

Source code in vllm/model_executor/models/zaya.py

class ZayaRouter(nn.Module):
    """Picks an expert per token; an extra "skip" class bypasses the experts."""

    def __init__(self, config: ZayaConfig, layer_idx: int, prefix: str):
        super().__init__()
        self.skip_id = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.down_proj = ReplicatedLinear(
            config.hidden_size,
            config.router_hidden_size,
            return_bias=False,
            prefix=f"{prefix}.down_proj",
        )
        self.use_eda = layer_idx != 0
        if self.use_eda:
            self.router_states_scale = nn.Parameter(
                torch.ones(config.router_hidden_size)
            )
        self.router_mlp = ZayaRouterMLP(
            config, config.num_experts + 1, f"{prefix}.router_mlp"
        )
        self.balancing_biases = nn.Buffer(
            torch.zeros(config.num_experts + 1, dtype=torch.float32)
        )

    def forward(
        self, hidden_states: torch.Tensor, router_states: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        router_hidden = self.down_proj(hidden_states)
        if self.use_eda and router_states is not None:
            router_hidden = router_hidden + router_states * self.router_states_scale
        probs = self.router_mlp(router_hidden).softmax(dim=-1).float()
        _, indices = (probs + self.balancing_biases).topk(self.top_k, dim=-1)
        weights = probs.gather(1, indices)
        skip = indices == self.skip_id
        routing = torch.cat(
            [weights.masked_fill(skip, 0), indices.masked_fill(skip, 0).float()], dim=-1
        )
        return routing, router_hidden

_cca_state_shapes(config)

(raw q/k history, previous delayed value) of one CCA layer.

Source code in vllm/model_executor/models/zaya.py

def _cca_state_shapes(config: ZayaConfig) -> tuple[tuple[int, int], tuple[int]]:
    """`(raw q/k history, previous delayed value)` of one CCA layer."""
    heads = config.num_attention_heads + config.num_key_value_heads
    history = config.cca_time0 + config.cca_time1 - 2
    delayed_dim = config.num_key_value_heads * config.head_dim // 2
    return (heads * config.head_dim, history), (delayed_dim,)