vllm.models.flint.ops.qk_norm_rope

Fused per-head q/k RMSNorm + NeoX RoPE + paged KV-cache write.

Numerics follow the unfused vLLM ops it replaces: RMSNorm computes bf16(x * rsqrt(mean(x^2) + eps) * w) in fp32, RoPE computes x * cos - y * sin in fp32 from the fp32 cos/sin cache (as the reference implementation does; vLLM's rotary op would round the cache to bf16 first) and rounds to bf16, and K/V are stored unchanged in the bf16 cache.

Functions:

• qk_norm_rope_cache –

  Normalize and rotate q/k from qkv and write k/v into the cache.

qk_norm_rope_cache(qkv, positions, q_weight, k_weight, eps, cos_sin_cache, k_cache, v_cache, slot_mapping, num_q_heads, num_kv_heads, head_dim, launch_dependents=False)

Normalize and rotate q/k from qkv and write k/v into the cache.

Parameters:

• qkv

  (Tensor) –

  [T, (Hq + 2 * Hkv) * D] bf16; k is normalized in place.

• positions

  (Tensor) –

  [T] token positions.

• q_weight

  (Tensor) –

  [D] q RMSNorm weight.

• k_weight

  (Tensor) –

  [D] k RMSNorm weight.

• eps

  (float) –

  RMSNorm epsilon.

• cos_sin_cache

  (Tensor | None) –

  [max_pos, D] fp32 (cos | sin), or None for no RoPE.

• k_cache

  (Tensor | None) –

  [num_blocks, block_size, Hkv, D] view (any strides with a contiguous last dim), or None to skip the cache write.

• v_cache

  (Tensor | None) –

  Same shape and strides as k_cache.

• slot_mapping

  (Tensor | None) –

  [<= T] int64; negative slots are skipped.

• num_q_heads

  (int) –

  Hq.

• num_kv_heads

  (int) –

  Hkv.

• head_dim

  (int) –

  D.

• launch_dependents

  (bool, default: False ) –

  Let a PDL successor that waits on the grid dependency before reading q / K / V (Flint decode attention) start its prologue early. Must stay False before other kernels.

Returns:

• q ( Tensor ) –

  [T, Hq * D] bf16.

Source code in vllm/models/flint/ops/qk_norm_rope.py

def qk_norm_rope_cache(
    qkv: torch.Tensor,
    positions: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
    cos_sin_cache: torch.Tensor | None,
    k_cache: torch.Tensor | None,
    v_cache: torch.Tensor | None,
    slot_mapping: torch.Tensor | None,
    num_q_heads: int,
    num_kv_heads: int,
    head_dim: int,
    launch_dependents: bool = False,
) -> torch.Tensor:
    """Normalize and rotate q/k from ``qkv`` and write k/v into the cache.

    Args:
        qkv: [T, (Hq + 2 * Hkv) * D] bf16; k is normalized in place.
        positions: [T] token positions.
        q_weight: [D] q RMSNorm weight.
        k_weight: [D] k RMSNorm weight.
        eps: RMSNorm epsilon.
        cos_sin_cache: [max_pos, D] fp32 (cos | sin), or None for no RoPE.
        k_cache: [num_blocks, block_size, Hkv, D] view (any strides with a
            contiguous last dim), or None to skip the cache write.
        v_cache: Same shape and strides as ``k_cache``.
        slot_mapping: [<= T] int64; negative slots are skipped.
        num_q_heads: Hq.
        num_kv_heads: Hkv.
        head_dim: D.
        launch_dependents: Let a PDL successor that waits on the grid
            dependency before reading q / K / V (Flint decode attention) start
            its prologue early. Must stay False before other kernels.

    Returns:
        q: [T, Hq * D] bf16.

    """
    num_tokens = qkv.shape[0]
    q = torch.empty(
        (num_tokens, num_q_heads * head_dim), dtype=qkv.dtype, device=qkv.device
    )
    if num_tokens == 0:
        return q
    if k_cache is not None and v_cache is not None and slot_mapping is not None:
        assert k_cache.stride() == v_cache.stride() and k_cache.stride(-1) == 1
        strides = k_cache.stride()[:3]
        block_size = k_cache.shape[1]
        num_slots = slot_mapping.shape[0]
    else:
        k_cache = v_cache = slot_mapping = qkv
        strides = (0, 0, 0)
        block_size = 1
        num_slots = 0
    grid = (triton.cdiv(num_tokens, _BLOCK_T), num_q_heads + 2 * num_kv_heads)
    _qk_norm_rope_cache_kernel[grid](
        qkv,
        qkv.stride(0),
        q,
        positions,
        q_weight,
        k_weight,
        qkv if cos_sin_cache is None else cos_sin_cache,
        0 if cos_sin_cache is None else cos_sin_cache.stride(0),
        slot_mapping,
        num_tokens,
        num_slots,
        k_cache,
        v_cache,
        *strides,
        block_size,
        eps,
        NUM_Q_HEADS=num_q_heads,
        NUM_KV_HEADS=num_kv_heads,
        HEAD_DIM=head_dim,
        APPLY_ROPE=cos_sin_cache is not None,
        BLOCK_T=_BLOCK_T,
        LAUNCH_DEPENDENTS=launch_dependents,
        num_warps=4,
    )
    return q