Skip to content

vllm.model_executor.layers.fla.ops.chunk_scaled_dot_kkt

Functions:

chunk_scaled_dot_kkt_fwd(k, g=None, beta=None, cu_seqlens=None, chunk_indices=None, chunk_size=FLA_CHUNK_SIZE, output_dtype=torch.float32)

Compute beta * K * K^T.

Parameters:

  • k

    (Tensor) –

    The key tensor of shape [B, T, H, K].

  • beta

    (Tensor, default: None ) –

    The beta tensor of shape [B, T, H].

  • g

    (Tensor, default: None ) –

    The cumulative sum of the gate tensor of shape [B, T, H]. Default: None.

  • cu_seqlens

    (Tensor, default: None ) –

    The cumulative sequence lengths of the input tensor. Default: None

  • chunk_indices

    (Tensor, default: None ) –

    Pre-computed chunk indices. If None and cu_seqlens is provided, computed internally. Default: None

  • chunk_size

    (int, default: FLA_CHUNK_SIZE ) –

    The chunk size. Default: 64.

  • output_dtype

    (dtype, default: float32 ) –

    The dtype of the output tensor. Default: torch.float32

Returns:

  • Tensor –

    beta * K * K^T of shape [B, T, H, BT] where BT is the chunk size.

Source code in vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def chunk_scaled_dot_kkt_fwd(
    k: torch.Tensor,
    g: torch.Tensor | None = None,
    beta: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None,
    chunk_indices: torch.Tensor | None = None,
    chunk_size: int = FLA_CHUNK_SIZE,
    output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    r"""
    Compute beta * K * K^T.

    Args:
        k (torch.Tensor):
            The key tensor of shape `[B, T, H, K]`.
        beta (torch.Tensor):
            The beta tensor of shape `[B, T, H]`.
        g (torch.Tensor):
            The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
        cu_seqlens (torch.Tensor):
            The cumulative sequence lengths of the input tensor.
            Default: None
        chunk_indices (torch.Tensor):
            Pre-computed chunk indices. If None and cu_seqlens is provided,
            computed internally. Default: None
        chunk_size (int):
            The chunk size. Default: 64.
        output_dtype (torch.dtype):
            The dtype of the output tensor. Default: `torch.float32`

    Returns:
        beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
    """
    # This kernel is slightly different from fla to support Q/K with different head numbers.
    # In fla, Q/K always have the same head number, so Hg is always equal to H.
    B, T, Hg, K = k.shape
    H = beta.shape[-1]
    BT = chunk_size
    if chunk_indices is None and cu_seqlens is not None:
        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

    A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
    chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](
        k=k,
        g=g,
        beta=beta,
        A=A,
        cu_seqlens=cu_seqlens,
        chunk_indices=chunk_indices,
        T=T,
        H=H,
        Hg=Hg,
        K=K,
        BT=BT,
    )
    return A