vllm.models.step5.nvidia.ops.indexer_norm_rope

Indexer normalization, rotary embedding and FP8 query preparation.

Functions:

• indexer_norm_rope –

  Prepare indexer Q/K/weights and update the decode CSA state.

indexer_norm_rope(query, key, weights, q_weight, k_weight, k_bias, cos_sin_cache, positions, eps, zero_centered, native_rope=False, *, num_groups=1, decode_cache=None, decode_ring=None, decode_z=None, decode_slots=None)

Prepare indexer Q/K/weights and update the decode CSA state.

Source code in vllm/models/step5/nvidia/ops/indexer_norm_rope.py

def indexer_norm_rope(
    query: torch.Tensor,
    key: torch.Tensor,
    weights: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    k_bias: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    positions: torch.Tensor,
    eps: float,
    zero_centered: bool,
    native_rope: bool = False,
    *,
    num_groups: int = 1,
    decode_cache=None,
    decode_ring=None,
    decode_z: torch.Tensor | None = None,
    decode_slots: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Prepare indexer Q/K/weights and update the decode CSA state."""
    tokens, width = query.shape
    if (decode_cache is None) != (decode_ring is None):
        raise ValueError("CSA fusion needs both the summary cache and the ring")
    if decode_cache is not None:
        if decode_z is None or decode_slots is None:
            raise ValueError("CSA fusion needs decode z and slots")
    elif decode_z is not None or decode_slots is not None:
        raise ValueError("CSA decode inputs need a summary cache")
    if decode_slots is not None and decode_slots.numel() > tokens:
        raise ValueError("CSA decode rows must fit the indexer token count")
    if num_groups < 1 or width % (num_groups * 256):
        raise ValueError("Indexer query heads must divide evenly into provider groups")
    heads = width // (num_groups * 256)
    padded_heads = max(heads, 8)
    q = torch.empty(
        (num_groups, tokens, padded_heads, 256),
        device=query.device,
        dtype=torch.float8_e4m3fn,
    )
    k = torch.empty((tokens, 1, 256), device=key.device, dtype=key.dtype)
    w = torch.empty(
        (num_groups, tokens, padded_heads), device=weights.device, dtype=torch.float32
    )
    if tokens:
        _indexer_norm_rope[(tokens, num_groups * padded_heads + 1)](
            query,
            key,
            weights,
            q_weight,
            k_weight,
            k_bias,
            cos_sin_cache,
            positions,
            q,
            k,
            w,
            decode_z,
            decode_slots,
            decode_cache.mean if decode_cache is not None else None,
            decode_ring.rows if decode_ring is not None else None,
            decode_ring.metadata().base if decode_ring is not None else None,
            decode_cache is not None,
            decode_slots.numel() if decode_slots is not None else 0,
            decode_z.stride(0) if decode_z is not None else 0,
            decode_slots.stride(0) if decode_slots is not None else 0,
            decode_ring.capacity if decode_ring is not None else 8,
            *(
                summary_layout(decode_cache.mean)
                if decode_cache is not None
                else (0, 8)
            ),
            query.stride(0),
            key.stride(0),
            weights.stride(0),
            heads,
            padded_heads,
            num_groups,
            tokens,
            cos_sin_cache.shape[1],
            cos_sin_cache.stride(0),
            eps,
            float(zero_centered),
            native_rope,
            num_warps=4,
            enable_fp_fusion=False,
        )
    return q.transpose(0, 1), k, w.transpose(0, 1)