recipes/docker-images/glm53-kindling/experimental/fixes/recoverssm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""GLM-5.3-Flash KDA speculative verify with one checkpoint state per request.

A port of Kimi-K3's RecoverSSM (models/kimi_k3/nvidia/ops/recoverssm.py and
kda_metadata.py), enabled with VLLM_GLM5NEXT_RECOVERSSM=1.

The stock spec path keeps 1 + num_spec recurrent states per request and writes
one full state per verified token. Here the verify kernel reads the request's
checkpoint, writes the outputs and records each token's raw inputs. After
sampling, the commit replays only the accepted tokens from the checkpoint.

Where this differs from Kimi-K3's version:
- The records are the raw k, gate, v and beta, and the commit replays them
  through the stock fused_recurrent_kda kernel, so the committed state is the
  one the stock spec path would have stored. Kimi-K3 stores fp32 corrections
  and folds them in closed form, which rounds differently. A copy of the stock
  recurrence also rounded differently (1-2 ulp): it read its state and records
  through integer addresses, so Triton chose other layouts and reduced sums in
  another order.
- The records live in a per-row scratch outside the KV pool: they only live
  from verify to commit within one step. Inside the mamba page they would
  make it larger than the 2304-token attention page and double the block size.
"""

from dataclasses import dataclass, field
from functools import cache
from typing import Any

import torch

from vllm.config import VllmConfig
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
from vllm.models.glm5next.nvidia.ops.third_party.kda.fused_recurrent import (
    token_stride as _token_stride,
)
from vllm.third_party.flash_linear_attention.ops.op import exp
from vllm.triton_utils import tl, triton
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.attention.backend import CommonAttentionMetadata
from vllm.v1.attention.backends.gdn_attn import (
    GDNAttentionBackend,
    GDNAttentionMetadata,
    GDNAttentionMetadataBuilder,
)
from vllm.v1.attention.backends.recoverssm_metadata import (
    RecoverSSMMetadata,
    RecoverSSMPostprocessMetadata,
)
from vllm.v1.attention.backends.utils import (
    NULL_BLOCK_ID,
    compute_causal_conv1d_metadata,
    mamba_get_block_table_tensor,
)
from vllm.v1.kv_cache_interface import MambaSpec

# Tile shape and warps of the stock fused_recurrent_kda launch. The verify
# kernel uses the same so its fp32 arithmetic compiles the same way.
_BV = 8
_NUM_WARPS = 1
_NUM_STAGES = 3


@triton.jit
def _verify_kernel(
    q,
    k,
    v,
    g,
    beta,
    o,
    h0,
    rec_kv,
    rec_g,
    rec_beta,
    cu_seqlens,
    state_indices,
    a_log,
    g_bias,
    scale,
    stride_q_t,
    stride_k_t,
    stride_v_t,
    stride_beta_t,
    stride_state_indices,
    stride_state: tl.constexpr,
    H: tl.constexpr,
    K: tl.constexpr,
    V: tl.constexpr,
    Q: tl.constexpr,
    BK: tl.constexpr,
    BV: tl.constexpr,
    LOWER_BOUND: tl.constexpr,
):
    # The per-token body is fused_recurrent_gated_delta_rule_fwd_kernel's
    # (IS_KDA, COMPUTE_GATE, SIGMOID_BETA, USE_QK_L2NORM_IN_KERNEL). Keep the
    # two in step: the equivalence test compares them bit for bit.
    i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    i_n, i_h = i_nh // H, i_nh % H
    bos = tl.load(cu_seqlens + i_n).to(tl.int64)
    eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)
    T = eos - bos
    if T == 0:
        return

    o_k = i_k * BK + tl.arange(0, BK)
    o_v = i_v * BV + tl.arange(0, BV)

    p_q = q + bos * stride_q_t + i_h * K + o_k
    p_k = k + bos * stride_k_t + i_h * K + o_k
    p_v = v + bos * stride_v_t + i_h * V + o_v
    p_beta = beta + bos * stride_beta_t + i_h
    p_gk = g + (bos * H + i_h) * K + o_k
    b_a_log = tl.exp(tl.load(a_log + i_h).to(tl.float32))
    p_o = o + (bos * H + i_h) * V + o_v

    mask_k = o_k < K
    mask_v = o_v < V
    mask_h = mask_v[:, None] & mask_k[None, :]

    state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64)
    if state_idx <= 0:
        return
    b_h = tl.zeros([BV, BK], dtype=tl.float32)
    p_h0 = h0 + state_idx * stride_state + i_h * V * K + o_v[:, None] * K + o_k[None, :]
    b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)

    # Row i_n owns record tokens [i_n * Q, i_n * Q + Q), laid out the way the
    # stock kernel reads q/k/v, g and beta. The commit reads them in the same
    # step over the same rows, so they never outlive a step.
    tok = i_n * Q
    p_rk = rec_kv + (tok * 2 * H + i_h) * K + o_k
    p_rv = rec_kv + ((tok * 2 + 1) * H + i_h) * V + o_v
    p_rg = rec_g + (tok * H + i_h) * K + o_k
    p_rbeta = rec_beta + tok * H + i_h

    for i_t in range(0, T):
        b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
        r_k = tl.load(p_k, mask=mask_k, other=0)
        b_k = r_k.to(tl.float32)
        r_v = tl.load(p_v, mask=mask_v, other=0)
        b_v = r_v.to(tl.float32)

        b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
        b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
        b_q = b_q * scale
        r_g = tl.load(p_gk)
        b_gk = r_g.to(tl.float32)
        b_gk += tl.load(g_bias + i_h * K + o_k, mask=mask_k, other=0.0).to(tl.float32)
        b_gk = LOWER_BOUND / (1.0 + tl.exp(-(b_a_log * b_gk)))
        b_h *= exp(b_gk[None, :])
        b_v -= tl.sum(b_h * b_k[None, :], 1)
        r_beta = tl.load(p_beta)
        b_beta = r_beta.to(tl.float32)
        b_beta = tl.sigmoid(b_beta)
        b_v *= b_beta
        b_h += b_v[:, None] * b_k[None, :]
        b_o = tl.sum(b_h * b_q[None, :], 1)
        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)

        tl.store(p_rv + i_t * 2 * H * V, r_v, mask=mask_v)
        if i_v == 0:
            tl.store(p_rk + i_t * 2 * H * K, r_k, mask=mask_k)
            tl.store(p_rg + i_t * H * K, r_g, mask=mask_k)
            tl.store(p_rbeta + i_t * H, r_beta)

        p_q += stride_q_t
        p_k += stride_k_t
        p_o += H * V
        p_v += stride_v_t
        p_gk += H * K
        p_beta += stride_beta_t


@triton.heuristics(
    {
        "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None,
        "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None,
    }
)
@triton.jit
def _prepare_commit_plan_kernel(
    num_accepted_ptr,
    request_indices_ptr,
    state_indices_ptr,
    query_start_loc_ptr,
    block_table_ptr,
    num_computed_ptr,
    commit_lens_ptr,
    final_state_indices_ptr,
    boundary_state_indices_ptr,
    boundary_recovery_lens_ptr,
    null_block_id,
    mamba_block_size,
    block_table_width,
    stride_num_accepted,
    stride_request_indices,
    stride_state_indices,
    stride_query_start_loc,
    stride_block_table_row,
    stride_block_table_col,
    stride_num_computed,
    SPEC_QUERY_LEN: tl.constexpr,
    HAS_REQUEST_INDICES: tl.constexpr,
    ALIGN_MODE: tl.constexpr,
):
    # Kimi-K3's RecoverSSM, except for the final state's column. Kimi-K3 uses
    # n // block_size, where n is the committed count. When n ends exactly on a
    # boundary, that is the next block, which align mode has not allocated
    # yet, so the state was lost. This uses (n - 1) // block_size, the block
    # holding the last token, as the stock align path does; the boundary block
    # is then the final block itself. _record_state_block_kernel records the
    # same column for the next step's pre-step copy.
    spec_idx = tl.program_id(0)
    source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to(
        tl.int64
    )
    request_idx = spec_idx
    if HAS_REQUEST_INDICES:
        request_idx = tl.load(
            request_indices_ptr + spec_idx * stride_request_indices
        ).to(tl.int64)
    num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to(
        tl.int32
    )
    bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64)
    eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to(
        tl.int64
    )
    query_len = (eos - bos).to(tl.int32)
    commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len)
    commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN)

    final_state_idx = source_state_idx
    boundary_state_idx = null_block_id
    boundary_recovery_len = 0
    if ALIGN_MODE:
        num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to(
            tl.int32
        )
        final_num_computed = num_computed + commit_len
        final_state_col = tl.minimum(
            tl.maximum(final_num_computed - 1, 0) // mamba_block_size,
            block_table_width - 1,
        )
        final_state_idx = tl.load(
            block_table_ptr
            + request_idx * stride_block_table_row
            + final_state_col * stride_block_table_col
        ).to(tl.int64)
        next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size
        crosses_boundary = final_num_computed >= next_boundary
        boundary_recovery_len = next_boundary - num_computed
        boundary_state_idx = tl.load(
            block_table_ptr
            + request_idx * stride_block_table_row
            + (next_boundary // mamba_block_size - 1) * stride_block_table_col,
            mask=crosses_boundary,
            other=null_block_id,
        ).to(tl.int64)
    valid = (source_state_idx > null_block_id) & (commit_len > 0)
    tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0))
    tl.store(
        final_state_indices_ptr + spec_idx,
        tl.where(valid, final_state_idx, null_block_id),
    )
    tl.store(
        boundary_state_indices_ptr + spec_idx,
        tl.where(valid, boundary_state_idx, null_block_id),
    )
    tl.store(
        boundary_recovery_lens_ptr + spec_idx,
        tl.where(valid, boundary_recovery_len, 0),
    )


@triton.jit
def _compact_conv_state_kernel(
    conv_state_ref_ptr,
    conv_state_base_addrs_ptr,
    conv_state_block_strides_ptr,
    conv_state_dim_strides_ptr,
    conv_state_token_strides_ptr,
    state_indices_ptr,
    commit_lens_ptr,
    final_state_indices_ptr,
    boundary_state_indices_ptr,
    boundary_recovery_lens_ptr,
    null_block_id,
    conv_dim,
    conv_history_len,
    stride_state_indices,
    BLOCK_D: tl.constexpr,
    BLOCK_HISTORY: tl.constexpr,
    ALIGN_MODE: tl.constexpr,
):
    # Unchanged from Kimi-K3's RecoverSSM. After the verify the conv state is
    # [history[1:], x_0 .. x_{q-1}]; the window after n accepted tokens starts
    # at n - 1, and moves to column 0 so the next step reads it with
    # num_accepted_tokens == 1.
    pid_d = tl.program_id(0)
    pid_b = tl.program_id(1)
    pid_l = tl.program_id(2)
    source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(
        tl.int64
    )
    if source_state_idx <= null_block_id:
        return

    commit_len = tl.load(commit_lens_ptr + pid_b)
    if commit_len == 0:
        return
    final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64)
    boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64)
    boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b)

    if final_state_idx <= null_block_id:
        return

    base_addr = tl.load(conv_state_base_addrs_ptr + pid_l)
    block_stride = tl.load(conv_state_block_strides_ptr + pid_l)
    dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l)
    token_stride = tl.load(conv_state_token_strides_ptr + pid_l)
    conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty))
    source_ptr = conv_state_ptr + source_state_idx * block_stride
    final_ptr = conv_state_ptr + final_state_idx * block_stride

    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
    offs_h = tl.arange(0, BLOCK_HISTORY)
    mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len)
    final_values = tl.load(
        source_ptr
        + offs_d[:, None] * dim_stride
        + (commit_len - 1 + offs_h[None, :]) * token_stride,
        mask=mask,
    )
    if ALIGN_MODE:
        boundary_values = tl.load(
            source_ptr
            + offs_d[:, None] * dim_stride
            + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride,
            mask=mask & (boundary_state_idx > null_block_id),
        )
        boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride
        tl.store(
            boundary_ptr
            + offs_d[:, None] * dim_stride
            + offs_h[None, :] * token_stride,
            boundary_values,
            mask=mask & (boundary_state_idx > null_block_id),
        )
    tl.store(
        final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride,
        final_values,
        mask=mask,
    )


@triton.jit
def _replay_indices_kernel(
    state_indices_ptr,
    commit_lens_ptr,
    final_state_indices_ptr,
    boundary_state_indices_ptr,
    boundary_recovery_lens_ptr,
    replay_indices_ptr,
    null_block_id,
    stride_state_indices,
    Q: tl.constexpr,
    BLOCK: tl.constexpr,
):
    # One row of the stock kernel's ssm_state_indices: it reads the initial
    # state at column num_accepted - 1 (Q here) and stores the state after
    # token t to column t when that is not the null block.
    i_n = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    source = tl.load(state_indices_ptr + i_n * stride_state_indices).to(tl.int32)
    commit_len = tl.load(commit_lens_ptr + i_n)
    final = tl.load(final_state_indices_ptr + i_n).to(tl.int32)
    boundary = tl.load(boundary_state_indices_ptr + i_n).to(tl.int32)
    boundary_len = tl.load(boundary_recovery_lens_ptr + i_n)
    row = tl.where(cols == commit_len - 1, final, null_block_id)
    # A boundary at the last token is the final state's own block.
    at_boundary = (
        (cols == boundary_len - 1)
        & (boundary_len < commit_len)
        & (boundary > null_block_id)
    )
    row = tl.where(at_boundary, boundary, row)
    row = tl.where(cols == Q, source, row)
    valid = (source > null_block_id) & (commit_len > 0) & (final > null_block_id)
    row = tl.where(valid, row, null_block_id)
    tl.store(replay_indices_ptr + i_n * (Q + 1) + cols, row, mask=cols < Q + 1)


@dataclass
class Records:
    """One layer's verify inputs for up to ``rows`` rows of ``q`` tokens.

    Token-major so the stock recurrent kernel reads them as q/k/v, g, beta:
    ``kv[t, 0]`` is k and ``kv[t, 1]`` is v for record token ``t``.
    """

    kv: torch.Tensor  # [rows * q, 2, heads, head_dim]
    g: torch.Tensor  # [rows * q, heads, head_dim]
    beta: torch.Tensor  # [rows * q, heads]
    rows: int
    q: int


def allocate_records(
    num_rows: int,
    num_heads: int,
    head_dim: int,
    spec_query_len: int,
    dtype: torch.dtype,
    device: torch.device,
) -> Records:
    tokens = num_rows * spec_query_len
    return Records(
        kv=torch.empty((tokens, 2, num_heads, head_dim), dtype=dtype, device=device),
        g=torch.empty((tokens, num_heads, head_dim), dtype=dtype, device=device),
        beta=torch.empty((tokens, num_heads), dtype=dtype, device=device),
        rows=num_rows,
        q=spec_query_len,
    )


def kda_recoverssm_verify(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    a_log: torch.Tensor,
    g_bias: torch.Tensor,
    lower_bound: float,
    checkpoint_state: torch.Tensor,
    records: Records,
    query_start_loc: torch.Tensor,
    state_indices: torch.Tensor,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Verify a speculative window from each row's checkpoint without writing it.

    Takes what kda.py passes fused_recurrent_kda on its spec path (raw gate
    logits and raw beta), with ``state_indices`` holding one checkpoint block
    per row. Writes the outputs and fills ``records`` for the commit.
    """
    _, total_tokens, num_heads, key_dim = k.shape
    value_dim = v.shape[-1]
    num_rows = state_indices.shape[0]
    if query_start_loc.shape[0] != num_rows + 1:
        raise ValueError("KDA RecoverSSM query metadata is incompatible")
    if num_rows > records.rows:
        raise ValueError("KDA RecoverSSM batch exceeds its record rows")
    if records.kv.shape[1:] != (2, num_heads, key_dim) or value_dim != key_dim:
        raise ValueError("KDA RecoverSSM records do not match the layer")
    if records.kv.dtype != k.dtype or records.beta.dtype != beta.dtype:
        raise ValueError("KDA RecoverSSM records must use the activation dtype")
    # No token-count check: FULL graphs pad the token count, and the builder
    # already routes rows longer than the record window to the prefill path.
    if out is None:
        out = torch.empty(k.shape, dtype=k.dtype, device=k.device)
    else:
        assert out.shape == k.shape and out.dtype == k.dtype
        assert out.is_contiguous()
    if total_tokens == 0 or num_rows == 0:
        return out

    g = g.contiguous()
    BK = triton.next_power_of_2(key_dim)
    BV = min(triton.next_power_of_2(value_dim), _BV)
    assert BK == key_dim, "KDA RecoverSSM needs a power-of-two head dim"
    grid = (1, triton.cdiv(value_dim, BV), num_rows * num_heads)
    _verify_kernel[grid](
        q,
        k,
        v,
        g,
        beta,
        out,
        checkpoint_state,
        records.kv,
        records.g,
        records.beta,
        query_start_loc,
        state_indices,
        a_log.reshape(-1).contiguous(),
        g_bias.reshape(-1).contiguous(),
        key_dim**-0.5,
        _token_stride(q),
        _token_stride(k),
        _token_stride(v),
        _token_stride(beta),
        state_indices.stride(0),
        stride_state=checkpoint_state.stride(0),
        H=num_heads,
        K=key_dim,
        V=value_dim,
        Q=records.q,
        BK=BK,
        BV=BV,
        LOWER_BOUND=lower_bound,
        num_warps=_NUM_WARPS,
        num_stages=_NUM_STAGES,
    )
    return out


@dataclass
class Glm5NextRecoverSSMCommitContext:
    """Per-KV-group commit state: every layer's pages, records and gate params."""

    conv_states: tuple[torch.Tensor, ...]
    conv_state_base_addrs: torch.Tensor
    conv_state_block_strides: torch.Tensor
    conv_state_dim_strides: torch.Tensor
    conv_state_token_strides: torch.Tensor
    conv_history_len: int
    layers: tuple[Any, ...]
    lower_bound: float
    spec_query_len: int
    commit_lens: torch.Tensor
    final_state_indices: torch.Tensor
    boundary_state_indices: torch.Tensor
    boundary_recovery_lens: torch.Tensor
    replay_indices: torch.Tensor
    replay_query_start_loc: torch.Tensor
    replay_num_accepted: torch.Tensor
    replay_out: torch.Tensor

    @classmethod
    def create(
        cls, layers: list[Any], *, spec_query_len: int, max_num_reqs: int
    ) -> "Glm5NextRecoverSSMCommitContext":
        """``layers``: KDA layers with ``kv_cache`` (conv, state), ``A_log``,
        ``dt_bias``, ``kda_lower_bound`` and ``recoverssm_records``."""
        if not layers:
            raise ValueError("KDA RecoverSSM commit requires at least one layer")
        conv_states = [layer.kv_cache[0] for layer in layers]
        if not is_conv_state_dim_first():
            conv_states = [state.transpose(-1, -2) for state in conv_states]
        checkpoints = [layer.kv_cache[1] for layer in layers]
        lower_bounds = {layer.kda_lower_bound for layer in layers}
        if len(lower_bounds) != 1:
            raise ValueError("KDA RecoverSSM layers need matching gate bounds")

        state_ref = checkpoints[0]
        _, num_heads, value_dim, key_dim = state_ref.shape
        for state in checkpoints:
            if (
                state.shape[1:] != state_ref.shape[1:]
                or state.dtype != state_ref.dtype
                or state.stride()[1:] != (value_dim * key_dim, key_dim, 1)
            ):
                raise ValueError("KDA RecoverSSM layers need matching state layout")
        record_rows = layers[0].recoverssm_records.rows
        for layer in layers:
            records = layer.recoverssm_records
            if records.q != spec_query_len or records.rows != record_rows:
                raise ValueError("KDA RecoverSSM layers need matching records")
            if records.rows < max_num_reqs:
                raise ValueError("KDA RecoverSSM records do not cover max_num_seqs")
            if records.kv.shape[1:] != (2, num_heads, key_dim):
                raise ValueError("KDA RecoverSSM records do not match the state")

        conv_ref = conv_states[0]
        conv_history_len = conv_ref.shape[2] - spec_query_len + 1
        if conv_history_len <= 0:
            raise ValueError("KDA RecoverSSM conv state is shorter than its window")
        for conv_state in conv_states:
            if conv_state.shape[1:] != conv_ref.shape[1:] or (
                conv_state.dtype != conv_ref.dtype
            ):
                raise ValueError("KDA RecoverSSM layers need matching conv state")

        device = state_ref.device
        rows = max_num_reqs

        def _i64(values: list[int]) -> torch.Tensor:
            return torch.tensor(values, dtype=torch.int64, device=device)

        def _i32(n: int) -> torch.Tensor:
            return torch.empty(n, dtype=torch.int32, device=device)

        return cls(
            conv_states=tuple(conv_states),
            conv_state_base_addrs=_i64([s.data_ptr() for s in conv_states]),
            conv_state_block_strides=_i64([s.stride(0) for s in conv_states]),
            conv_state_dim_strides=_i64([s.stride(1) for s in conv_states]),
            conv_state_token_strides=_i64([s.stride(2) for s in conv_states]),
            conv_history_len=conv_history_len,
            layers=tuple(layers),
            lower_bound=lower_bounds.pop(),
            spec_query_len=spec_query_len,
            commit_lens=_i32(rows),
            final_state_indices=_i32(rows),
            boundary_state_indices=_i32(rows),
            boundary_recovery_lens=_i32(rows),
            replay_indices=torch.zeros(
                (rows, spec_query_len + 1), dtype=torch.int32, device=device
            ),
            replay_query_start_loc=torch.arange(
                0,
                (rows + 1) * spec_query_len,
                spec_query_len,
                dtype=torch.int32,
                device=device,
            ),
            replay_num_accepted=torch.full(
                (rows,), spec_query_len + 1, dtype=torch.int32, device=device
            ),
            # Shaped like the records' k, as the stock kernel requires.
            replay_out=torch.empty(
                (1, record_rows * spec_query_len, num_heads, value_dim),
                dtype=layers[0].recoverssm_records.kv.dtype,
                device=device,
            ),
        )

    def commit(
        self,
        num_accepted_tokens: torch.Tensor,
        state_indices: torch.Tensor,
        query_start_loc: torch.Tensor,
        request_indices: torch.Tensor | None = None,
        block_table: torch.Tensor | None = None,
        num_computed_tokens: torch.Tensor | None = None,
        mamba_block_size: int | None = None,
    ) -> None:
        """Fold each row's accepted tokens into its conv and recurrent state.

        ``num_accepted_tokens`` is per batch row (the runner's num_sampled);
        ``request_indices`` maps spec rows to batch rows when they differ.
        """
        batch = state_indices.shape[0]
        if batch == 0:
            return
        if batch > self.commit_lens.shape[0]:
            raise ValueError("KDA RecoverSSM commit batch exceeds its plan capacity")
        if query_start_loc.shape[0] != batch + 1:
            raise ValueError("KDA RecoverSSM commit metadata is incompatible")
        align_args = (block_table, num_computed_tokens, mamba_block_size)
        if any(a is not None for a in align_args) and any(
            a is None for a in align_args
        ):
            raise ValueError("KDA RecoverSSM align metadata is incomplete")
        if mamba_block_size is not None and mamba_block_size < self.spec_query_len:
            raise ValueError(
                "KDA RecoverSSM align block size must cover one speculative window"
            )
        align = block_table is not None
        block_table_stride = block_table.stride() if align else (0, 0)

        _prepare_commit_plan_kernel[(batch,)](
            num_accepted_tokens,
            request_indices,
            state_indices,
            query_start_loc,
            block_table,
            num_computed_tokens,
            self.commit_lens,
            self.final_state_indices,
            self.boundary_state_indices,
            self.boundary_recovery_lens,
            NULL_BLOCK_ID,
            mamba_block_size or 1,
            block_table.shape[1] if align else 1,
            num_accepted_tokens.stride(0),
            request_indices.stride(0) if request_indices is not None else 0,
            state_indices.stride(0),
            query_start_loc.stride(0),
            block_table_stride[0],
            block_table_stride[1],
            num_computed_tokens.stride(0) if align else 0,
            SPEC_QUERY_LEN=self.spec_query_len,
            num_warps=1,
        )

        num_layers = len(self.layers)
        conv_dim = self.conv_states[0].shape[1]
        _compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)](
            self.conv_states[0],
            self.conv_state_base_addrs,
            self.conv_state_block_strides,
            self.conv_state_dim_strides,
            self.conv_state_token_strides,
            state_indices,
            self.commit_lens,
            self.final_state_indices,
            self.boundary_state_indices,
            self.boundary_recovery_lens,
            NULL_BLOCK_ID,
            conv_dim,
            self.conv_history_len,
            state_indices.stride(0),
            BLOCK_D=256,
            BLOCK_HISTORY=triton.next_power_of_2(self.conv_history_len),
            ALIGN_MODE=align,
            num_warps=4,
        )

        q_len = self.spec_query_len
        _replay_indices_kernel[(batch,)](
            state_indices,
            self.commit_lens,
            self.final_state_indices,
            self.boundary_state_indices,
            self.boundary_recovery_lens,
            self.replay_indices,
            NULL_BLOCK_ID,
            state_indices.stride(0),
            Q=q_len,
            BLOCK=triton.next_power_of_2(q_len + 1),
            num_warps=1,
        )
        # The stock kernel runs all Q record tokens of each row and writes a
        # state only to the columns set above. Same kernel, same inputs, so
        # those states match the stock spec path's bit for bit.
        cu_seqlens = self.replay_query_start_loc[: batch + 1]
        indices = self.replay_indices[:batch]
        num_accepted = self.replay_num_accepted[:batch]
        for layer in self.layers:
            records = layer.recoverssm_records
            k = records.kv[:, 0].unsqueeze(0)
            fused_recurrent_kda(
                q=k,
                k=k,
                v=records.kv[:, 1].unsqueeze(0),
                g=records.g.unsqueeze(0),
                beta=records.beta.unsqueeze(0),
                initial_state=layer.kv_cache[1],
                use_qk_l2norm_in_kernel=True,
                cu_seqlens=cu_seqlens,
                ssm_state_indices=indices,
                num_accepted_tokens=num_accepted,
                out=self.replay_out,
                sigmoid_beta=True,
                a_log=layer.A_log,
                g_bias=layer.dt_bias,
                compute_gate=True,
                lower_bound=self.lower_bound,
            )


@dataclass
class _AlignCommit:
    block_table: torch.Tensor
    num_computed_tokens: torch.Tensor
    block_size: int


@dataclass
class _Commit:
    state_indices: torch.Tensor
    query_start_loc: torch.Tensor
    request_indices: torch.Tensor | None
    align: _AlignCommit | None


@dataclass
class Glm5NextRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
    recoverssm_commit: _Commit | None = None
    recoverssm_context: Glm5NextRecoverSSMCommitContext | None = field(
        default=None, repr=False, compare=False
    )

    def commit_recoverssm_state(
        self, num_accepted_tokens: torch.Tensor
    ) -> RecoverSSMPostprocessMetadata | None:
        commit = self.recoverssm_commit
        if commit is None:
            return None
        assert self.recoverssm_context is not None
        align = commit.align
        self.recoverssm_context.commit(
            num_accepted_tokens,
            commit.state_indices[: self.num_spec_decodes, 0],
            commit.query_start_loc[: self.num_spec_decodes + 1],
            request_indices=commit.request_indices,
            block_table=align.block_table if align is not None else None,
            num_computed_tokens=(
                align.num_computed_tokens if align is not None else None
            ),
            mamba_block_size=align.block_size if align is not None else None,
        )
        if align is None:
            return None
        return RecoverSSMPostprocessMetadata(
            num_spec_decodes=self.num_spec_decodes,
            request_indices=commit.request_indices,
            block_table=align.block_table,
            num_computed_tokens=align.num_computed_tokens,
            block_size=align.block_size,
        )


def _check_config(vllm_config: VllmConfig) -> None:
    """Fail at startup where the commit would otherwise be skipped silently."""
    from vllm.models.glm5next.common import model as glm_model

    if not getattr(glm_model, "_RECOVERSSM", False):
        raise RuntimeError(
            "VLLM_GLM5NEXT_RECOVERSSM=1 needs experimental/fixes/model.py "
            "mounted (recoverssm.yaml or sp.yaml): it installs the commit hook"
        )
    if not vllm_config.use_v2_model_runner:
        raise RuntimeError("KDA RecoverSSM commits from the V2 model runner only")
    if vllm_config.parallel_config.pipeline_parallel_size > 1:
        raise RuntimeError("KDA RecoverSSM requires pipeline_parallel_size=1")
    if vllm_config.cache_config.mamba_cache_mode not in ("none", "align"):
        raise RuntimeError("KDA RecoverSSM supports mamba cache modes none and align")


class Glm5NextRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder):
    """GDN metadata for GLM's KDA with one state slot per spec-decode row.

    Every running decode row goes down the spec path, drafts or not, so the
    conv state is always in the compacted form the commit leaves. Steps
    without such rows use the GDN builder unchanged.
    """

    def __init__(
        self,
        kv_cache_spec: MambaSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ) -> None:
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
        _check_config(vllm_config)
        assert self.use_spec_decode, "KDA RecoverSSM needs speculative decoding"
        self.spec_query_len = self.num_spec + 1
        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
        rows = max(self.decode_cudagraph_max_bs, self.max_num_reqs)
        self.rs_state_indices = torch.empty((rows, 1), dtype=torch.int32, device=device)
        # The conv kernel reads the window at num_accepted - 1; the commit
        # already moved it to column 0.
        self.rs_ones = torch.ones(rows, dtype=torch.int32, device=device)
        self.recoverssm_context: Glm5NextRecoverSSMCommitContext | None = None
        # Allocated here, before any capture: the eager KDA segment must not
        # allocate during a graph capture.
        layers = vllm_config.compilation_config.static_forward_context
        for name in layer_names:
            layer = layers[name]
            if getattr(layer, "recoverssm_records", None) is None:
                layer.recoverssm_records = allocate_records(
                    self.max_num_reqs,
                    layer.local_num_heads,
                    layer.head_dim,
                    self.spec_query_len,
                    vllm_config.model_config.dtype,
                    device,
                )

    def _get_context(self) -> Glm5NextRecoverSSMCommitContext:
        if self.recoverssm_context is None:
            layers = self.vllm_config.compilation_config.static_forward_context
            self.recoverssm_context = Glm5NextRecoverSSMCommitContext.create(
                [layers[name] for name in self.layer_names],
                spec_query_len=self.spec_query_len,
                max_num_reqs=self.max_num_reqs,
            )
        return self.recoverssm_context

    def build(  # type: ignore[override]
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        num_accepted_tokens: torch.Tensor | None = None,
        num_decode_draft_tokens_cpu: torch.Tensor | None = None,
        fast_build: bool = False,
    ) -> GDNAttentionMetadata:
        m = common_attn_metadata
        if num_decode_draft_tokens_cpu is None:
            return super().build(
                common_prefix_len, m, num_accepted_tokens, None, fast_build
            )
        query_start_loc = m.query_start_loc
        query_start_loc_cpu = m.query_start_loc_cpu
        query_lens_cpu = query_start_loc_cpu.diff()
        spec_mask_cpu = num_decode_draft_tokens_cpu >= 0
        if m.is_prefilling is not None:
            spec_mask_cpu |= (~m.is_prefilling) & (query_lens_cpu > 0)
        # Rows longer than the record window (adaptive probes) run as prefills.
        spec_mask_cpu &= query_lens_cpu <= self.spec_query_len
        num_spec_decodes = int(spec_mask_cpu.sum())
        if num_spec_decodes == 0:
            return super().build(
                common_prefix_len, m, num_accepted_tokens, None, fast_build
            )

        block_table = mamba_get_block_table_tensor(
            m.block_table_tensor,
            m.seq_lens,
            self.kv_cache_spec,
            self.vllm_config.cache_config.mamba_cache_mode,
        )
        device = query_start_loc.device
        num_query_tokens = int(query_start_loc_cpu[-1])
        active_non_spec_cpu = (~spec_mask_cpu) & (query_lens_cpu > 0)
        num_prefills = int(active_non_spec_cpu.sum())
        num_prefill_tokens = int(query_lens_cpu[active_non_spec_cpu].sum())
        num_spec_decode_tokens = num_query_tokens - num_prefill_tokens

        request_indices = None
        has_initial_state = None
        nums_dict = batch_ptr = token_chunk_offset_ptr = None
        if num_prefills == 0:
            # Real rows precede the cudagraph padding, and all of them verify.
            spec_token_indx = non_spec_token_indx = None
            spec_state_indices = block_table[:num_spec_decodes, :1]
            non_spec_state_indices = None
            spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
            non_spec_query_start_loc = None
        else:
            query_lens = query_start_loc.diff()
            spec_mask = async_tensor_h2d(spec_mask_cpu, device=device)
            request_indices = async_tensor_h2d(
                spec_mask_cpu.nonzero(as_tuple=True)[0],
                dtype=torch.int32,
                device=device,
            )
            spec_token_masks = torch.repeat_interleave(
                spec_mask, query_lens, output_size=num_query_tokens
            )
            # Stable, so each row's tokens stay in order in both halves.
            index = torch.argsort(spec_token_masks, stable=True)
            non_spec_token_indx = index[:num_prefill_tokens]
            spec_token_indx = index[num_prefill_tokens:]
            spec_state_indices = block_table[spec_mask_cpu, :1]
            non_spec_state_indices = block_table[active_non_spec_cpu, 0]
            spec_query_start_loc = torch.zeros(
                num_spec_decodes + 1, dtype=torch.int32, device=device
            )
            torch.cumsum(query_lens[spec_mask_cpu], dim=0, out=spec_query_start_loc[1:])
            non_spec_query_start_loc = torch.zeros(
                num_prefills + 1, dtype=torch.int32, device=device
            )
            torch.cumsum(
                query_lens[active_non_spec_cpu],
                dim=0,
                out=non_spec_query_start_loc[1:],
            )
            non_spec_query_start_loc_cpu = torch.zeros(
                num_prefills + 1, dtype=torch.int32
            )
            torch.cumsum(
                query_lens_cpu[active_non_spec_cpu],
                dim=0,
                out=non_spec_query_start_loc_cpu[1:],
            )
            has_initial_state = (m.compute_num_computed_tokens() > 0)[
                active_non_spec_cpu
            ]
            nums_dict, batch_ptr, token_chunk_offset_ptr = (
                compute_causal_conv1d_metadata(
                    non_spec_query_start_loc_cpu, device=device
                )
            )
        num_accepted = self.rs_ones[:num_spec_decodes]

        batch_size = m.num_reqs
        if (
            self.use_full_cuda_graph
            and num_prefills == 0
            and batch_size <= self.decode_cudagraph_max_bs
            and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
        ):
            self.rs_state_indices[:num_spec_decodes].copy_(
                spec_state_indices, non_blocking=True
            )
            spec_state_indices = self.rs_state_indices[:batch_size]
            spec_state_indices[num_spec_decodes:].fill_(NULL_BLOCK_ID)
            self.spec_query_start_loc[: num_spec_decodes + 1].copy_(
                spec_query_start_loc, non_blocking=True
            )
            spec_num_query_tokens = spec_query_start_loc[-1]
            spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
            spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens)
            num_accepted = self.rs_ones[:batch_size]

        align = None
        if self.kv_cache_spec.mamba_cache_mode == "align":
            align = _AlignCommit(
                block_table=m.block_table_tensor,
                num_computed_tokens=m.compute_num_computed_tokens(),
                block_size=self.kv_cache_spec.block_size,
            )
        return Glm5NextRecoverSSMMetadata(
            num_prefills=num_prefills,
            num_prefill_tokens=num_prefill_tokens,
            num_decodes=0,
            num_decode_tokens=0,
            num_spec_decodes=num_spec_decodes,
            num_spec_decode_tokens=num_spec_decode_tokens,
            num_actual_tokens=m.num_actual_tokens,
            has_initial_state=has_initial_state,
            spec_query_start_loc=spec_query_start_loc,
            non_spec_query_start_loc=non_spec_query_start_loc,
            spec_state_indices_tensor=spec_state_indices,
            non_spec_state_indices_tensor=non_spec_state_indices,
            spec_sequence_masks=None,
            spec_token_indx=spec_token_indx,
            non_spec_token_indx=non_spec_token_indx,
            num_accepted_tokens=num_accepted,
            nums_dict=nums_dict,
            batch_ptr=batch_ptr,
            token_chunk_offset_ptr=token_chunk_offset_ptr,
            recoverssm_commit=_Commit(
                state_indices=spec_state_indices,
                query_start_loc=spec_query_start_loc,
                request_indices=request_indices,
                align=align,
            ),
            recoverssm_context=self._get_context(),
        )


class Glm5NextRecoverSSMBackend(GDNAttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "GLM5NEXT_KDA_RECOVERSSM"

    @staticmethod
    def get_builder_cls() -> type[Glm5NextRecoverSSMMetadataBuilder]:
        return Glm5NextRecoverSSMMetadataBuilder


@triton.heuristics(
    {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None}
)
@triton.jit
def _record_state_block_kernel(
    idx_mapping_ptr,
    num_sampled_ptr,
    request_indices_ptr,
    num_computed_ptr,
    state_idx_ptr,
    num_accepted_ptr,
    HAS_REQUEST_INDICES: tl.constexpr,
    MAMBA_BLOCK_SIZE: tl.constexpr,
    BLOCK_TABLE_WIDTH: tl.constexpr,
):
    # The runner's _postprocess_recoverssm_align_kernel, except for the column:
    # the commit left the state in the block holding the last committed token.
    spec_idx = tl.program_id(0)
    batch_idx = spec_idx
    if HAS_REQUEST_INDICES:
        batch_idx = tl.load(request_indices_ptr + spec_idx)
    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
    if req_state_idx < 0:
        return
    num_sampled = tl.load(num_sampled_ptr + batch_idx)
    num_computed = tl.load(num_computed_ptr + batch_idx)
    last = tl.maximum(num_computed + num_sampled - 1, 0)
    tl.store(
        state_idx_ptr + req_state_idx,
        tl.minimum(last // MAMBA_BLOCK_SIZE, BLOCK_TABLE_WIDTH - 1),
    )
    tl.store(num_accepted_ptr + req_state_idx, 1)


def record_state_blocks(
    meta: RecoverSSMPostprocessMetadata,
    idx_mapping: torch.Tensor,
    num_sampled: torch.Tensor,
    state_indices: torch.Tensor,
    num_accepted_tokens: torch.Tensor,
) -> None:
    """Point each committed request's running state column at its state block.

    The next step's pre-step copy moves it from there when the step's tokens
    pass into the next block.
    """
    _record_state_block_kernel[(meta.num_spec_decodes,)](
        idx_mapping,
        num_sampled,
        meta.request_indices,
        meta.num_computed_tokens,
        state_indices,
        num_accepted_tokens,
        MAMBA_BLOCK_SIZE=meta.block_size,
        BLOCK_TABLE_WIDTH=meta.block_table.shape[1],
    )


@cache
def model_state_cls() -> type:
    """MambaHybridModelState with the runner's RecoverSSM commit hook.

    The stock state creates the hook only when cache_config.use_kda_recoverssm
    is set, which VllmConfig's validator allows for Kimi-K3 alone and resets
    on every re-validation (the DFlash loader re-validates after the target
    model loads). The hook records the state block with record_state_blocks.
    """
    from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
    from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState

    class Glm5NextRecoverSSMState(RecoverSSMState):
        def commit_step(
            self,
            num_sampled: torch.Tensor | int,
            idx_mapping: torch.Tensor,
            *,
            state_indices: torch.Tensor | None,
            num_accepted_tokens: torch.Tensor,
        ) -> None:
            step = self._step
            self._step = None
            if isinstance(num_sampled, int) or step is None:
                return
            for metadata in step:
                meta = metadata.commit_recoverssm_state(num_sampled)
                if meta is None:
                    continue
                assert state_indices is not None
                record_state_blocks(
                    meta, idx_mapping, num_sampled, state_indices, num_accepted_tokens
                )

    class Glm5NextRecoverSSMModelState(MambaHybridModelState):
        def __init__(self, *args, **kwargs) -> None:
            super().__init__(*args, **kwargs)
            self.recoverssm = Glm5NextRecoverSSMState()

    return Glm5NextRecoverSSMModelState