vllm.models.kimi_k3.amd.ops.kda_chunk ¶

@cache
def _chunk_groups(chunks_per_seq: int, num_seqs: int, num_heads: int) -> int:
    """How many parallel chunk groups to cut each sequence into.

    The scan composes the group operators with the exact transfer
    ``M_g = prod_c (diag(d_c) - w_c^T kg_c)``, so the composition error does
    not grow with group length and the choice is purely a machine fit:

    * ``fill1``/``fill2``: the largest G whose pass-one / pass-two grid still
      fits in one scheduling batch. Rows are launched in blocks of at most
      ``kBV = 128``, so pass one (``kV + kK`` rows) needs two blocks per
      ``(n, h, g)`` and pass two one. Past a pass's fill point its block count
      grows as fast as its depth shrinks and the rung stops paying.
    * ``_DEEP_LEN``: unless the groups are still long, where pass two's halving
      is worth letting pass one spill to a second batch.
    * ``_MIN_LEN``: below this many chunks per group the per-group fixed cost
      exceeds what the shorter walk saves.

    ``G == 2`` is depth-neutral -- two passes of ``nt/2`` -- and only doubles
    staging, so the split is taken only from ``G >= 4``.
    """
    cus = _num_cus(torch.accelerator.current_device_index())
    nh = max(num_seqs * num_heads, 1)
    fill1 = cus // (_BLOCKS_P1 * nh)
    fill2 = cus // (_BLOCKS_P2 * nh)
    cand = fill1
    if fill2 > 0 and chunks_per_seq // fill2 >= _DEEP_LEN:
        cand = fill2
    cand = min(cand, chunks_per_seq // _MIN_LEN, _MAX_GROUPS)
    return cand if cand >= 4 else 1

def _kda_group_workspace(
    groups: int, nh: int, device: torch.device
) -> torch.Tensor | None:
    """One fp32 buffer the kernel carves into bg / sin_ / ag / mgT.

    Kept caller-owned so the memory stays inside the caching allocator and
    vLLM's memory profiling accounts for it; the kernel's carve order must
    match the layout here.
    """
    if groups <= 1:
        return None
    planes = groups * nh
    plane = HEAD_DIM * HEAD_DIM
    floats = 2 * planes * plane + planes * HEAD_DIM  # bg, sin_, ag
    floats += (planes * plane + 1) // 2  # mgT, bf16
    return torch.empty(floats, dtype=torch.float32, device=device)

def fused_kda_chunk(
    qg: torch.Tensor,
    w: torch.Tensor,
    u: torch.Tensor,
    kg_t: torch.Tensor,
    aqk: torch.Tensor,
    decay: torch.Tensor,
    out: torch.Tensor,
    scale: float,
    cu_seqlens: torch.Tensor,
    initial_state: torch.Tensor | None,
    output_final_state: bool,
    chunk_offsets: torch.Tensor | None = None,
    checkpoint_state: torch.Tensor | None = None,
    checkpoint_offsets: torch.Tensor | None = None,
    checkpoint_state_indices: torch.Tensor | None = None,
    state_cache: torch.Tensor | None = None,
    state_indices: torch.Tensor | None = None,
    has_initial_state: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Run the chunk recurrence and the output projection in one launch.

    Args:
        qg: ``q * exp2(gk_cumsum)``, ``[1, T, H, 128]``.
        kg_t: chunk-major transposed gated keys, ``[chunks, H, 128, 64]``.
        decay: ``exp2`` of each chunk's last gate row, ``[chunks, H, 128]``.
        out: output buffer, ``[1, T, H, 128]``; may alias ``u``'s source.
        checkpoint_state: destination for the mid-prefill state snapshots,
            fp32 ``[rows, H, 128, 128]``. Without
            ``checkpoint_state_indices`` it is a staging buffer indexed by
            sequence; with them it can be the paged state cache itself.
        checkpoint_offsets: int32 ``[N]`` token offsets, each relative to its
            sequence's first query token, at which to snapshot the recurrent
            state. ``0`` disables the export for that sequence. A non-zero
            offset must be a multiple of :data:`KDA_CHECKPOINT_ALIGNMENT`;
            the kernel skips any that is not, so the caller has to filter
            rather than rely on the store happening.
        checkpoint_state_indices: optional int32 ``[N]`` destination row per
            sequence. A negative row disables the export for that sequence.
        state_cache: the paged recurrent state, fp32
            ``[slots, H, 128, 128]``. When given, the walk reads sequence
            ``n``'s initial state from row ``state_indices[n]`` and writes its
            final state back to the same row, which removes the gather and the
            scatter that otherwise bracket this kernel. Replaces
            ``initial_state`` / ``output_final_state``, and the returned final
            state is then ``None`` because it is already in the cache.
        state_indices: int32 ``[N]`` cache row per sequence.
        has_initial_state: bool ``[N]``; a false entry starts that sequence
            from a zero state and its cache row is read only, never before the
            walk writes it.
    """
    if state_cache is not None:
        if initial_state is not None or output_final_state:
            raise ValueError("state_cache replaces initial_state/output_final_state")
        if state_indices is None or has_initial_state is None:
            raise ValueError("state_cache needs state_indices and has_initial_state")
    num_seqs = cu_seqlens.numel() - 1
    final_state = None
    if output_final_state:
        final_state = torch.empty(
            num_seqs,
            u.shape[2],
            u.shape[3],
            qg.shape[3],
            dtype=torch.float32,
            device=u.device,
        )
    cu_seqlens = cu_seqlens.to(torch.int32)
    if chunk_offsets is None:
        chunk_offsets = prepare_chunk_offsets(cu_seqlens, CHUNK_SIZE)

    # The chunk walk is serial in chunks. Splitting each sequence into `groups`
    # stretches that run in parallel costs a second pass and a scan, and only
    # pays once the chain is long enough to amortise both.
    chunks_per_seq = kg_t.shape[0] // max(num_seqs, 1)
    groups = _chunk_groups(chunks_per_seq, num_seqs, u.shape[2])
    group_state = _kda_group_workspace(groups, num_seqs * u.shape[2], u.device)

    torch.ops._C.fused_kda_chunk(
        qg,
        w,
        u,
        kg_t,
        aqk,
        decay,
        initial_state,
        final_state,
        out,
        cu_seqlens,
        chunk_offsets,
        scale,
        group_state,
        groups,
        checkpoint_state,
        None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32),
        None
        if checkpoint_state_indices is None
        else checkpoint_state_indices.to(torch.int32).contiguous(),
        state_cache,
        None if state_indices is None else state_indices.to(torch.int32).contiguous(),
        has_initial_state,
    )
    return out, final_state

def fused_kda_prologue(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    raw_g: torch.Tensor,
    raw_beta: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    scale: float,
    lower_bound: float,
    cu_seqlens: torch.Tensor,
    conv_weight: torch.Tensor | None = None,
    conv_state: torch.Tensor | None = None,
    conv_state_indices: torch.Tensor | None = None,
    conv_has_initial_state: torch.Tensor | None = None,
    chunk_indices: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
    """Run the whole chunk-path prologue in one launch.

    Replaces the two L2 norms, the gate cumsum, both intra-chunk passes and the
    w/u recompute. Returns the operands the fused chunk kernel consumes.

    The kernel can also apply the depthwise conv and its silu, in which case
    ``q``/``k``/``v`` are the three raw bands of the QKV projection and
    ``conv_state`` is updated in place. This path is currently not used.
    """
    # `beta` and `g` reach the layer as last-dim slices of the fused QKVGFAB
    # projection and carry its row stride. The kernel reads both with an
    # explicit per-token stride, so neither is copied here.
    _, t_total, num_heads, _ = q.shape
    if chunk_indices is None:
        chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)
    chunk_indices = chunk_indices.to(torch.int32)
    num_chunks = chunk_indices.shape[0]
    dev = q.device

    # q may be a strided band view, so the workspaces are sized rather than
    # cloned from it.
    def _like(last: int) -> torch.Tensor:
        return torch.empty(1, t_total, num_heads, last, dtype=q.dtype, device=q.device)

    ws = dict(
        qg=_like(HEAD_DIM),
        w=_like(HEAD_DIM),
        u=_like(HEAD_DIM),
        kg_t=torch.empty(
            num_chunks, num_heads, HEAD_DIM, CHUNK_SIZE, dtype=q.dtype, device=dev
        ),
        aqk=torch.empty(1, t_total, num_heads, CHUNK_SIZE, dtype=q.dtype, device=dev),
        decay=torch.empty(
            num_chunks, num_heads, HEAD_DIM, dtype=torch.float32, device=dev
        ),
    )
    torch.ops._C.fused_kda_prologue(
        q,
        k,
        v,
        raw_g,
        raw_beta,
        A_log.reshape(-1),
        dt_bias.reshape(-1),
        ws["qg"],
        ws["w"],
        ws["u"],
        ws["kg_t"],
        ws["aqk"],
        ws["decay"],
        cu_seqlens.to(torch.int32),
        chunk_indices,
        conv_weight,
        conv_state,
        conv_state_indices,
        conv_has_initial_state,
        scale,
        lower_bound,
    )
    return ws