vllm.models.flint.ops.qkv_fused

Fused decode qkv for Flint at M <= 8 tokens (CuTe DSL), two kernels:

1. fused_add_rms_norm (or rms_norm for the first layer) writing the normalized activations xn and the new residual; it releases its PDL dependents immediately so the second kernel's weights stream in meanwhile.
2. the qkv GEMV (same fragment math, K split and summation order as nvfp4_gemv at M <= 8) whose 16-row CTAs form one cluster per head and send their BF16 rows to rank 0 through DSMEM, which applies per-head q/k RMSNorm, NeoX RoPE (FP32 cos/sin table, FP32 math) and the paged KV-cache write like qk_norm_rope_cache.

Numerics are those of norm.fused_add_rms_norm / norm.rms_norm -> GEMV -> qk_norm_rope_cache; the per-head q/k norm sums its squares in a different order, so q/k can differ from that path by one BF16 ulp.

Functions:

• flint_qkv_fused –

  Residual-add RMSNorm + qkv GEMV + q/k norm, RoPE, cache write for

• no_rope_table –

  Placeholder fp32 cos|sin argument for kernels compiled without RoPE.

flint_qkv_fused(x, residual, norm_weight, eps, weight, weight_scale, global_scale, q_norm_weight, k_norm_weight, qk_eps, cos_sin_cache, positions, k_cache, v_cache, slot_mapping, num_q_heads, num_kv_heads)

Residual-add RMSNorm + qkv GEMV + q/k norm, RoPE, cache write for M <= 8 tokens.

Parameters:

• x

  (Tensor) –

  [M, 2048] bf16 attention input (contiguous).

• residual

  (Tensor | None) –

  [M, 2048] bf16 residual, or None for the first layer.

• norm_weight

  (Tensor) –

  Input RMSNorm weight.

• eps

  (float) –

  Input RMSNorm epsilon.

• weight

  (Tensor) –

  qkv weight, NVFP4 (uint8 [N, 1024]) or BF16 [N, 2048].

• weight_scale

  (Tensor | None) –

  NVFP4 128x4-swizzled block scales (None for BF16).

• global_scale

  (Tensor | None) –

  NVFP4 fp32 global scale (None for BF16).

• q_norm_weight

  (Tensor) –

  Per-head q RMSNorm weight.

• k_norm_weight

  (Tensor) –

  Per-head k RMSNorm weight.

• qk_eps

  (float) –

  q/k RMSNorm epsilon.

• cos_sin_cache

  (Tensor | None) –

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

• positions

  (Tensor) –

  [M] int64.

• k_cache

  (Tensor) –

  [blocks, block_size, kv_heads, 128] view.

• v_cache

  (Tensor) –

  Same shape and strides as k_cache.

• slot_mapping

  (Tensor) –

  [<= M] int64; negative slots and tokens beyond it are not written.

• num_q_heads

  (int) –

  Number of q heads.

• num_kv_heads

  (int) –

  Number of KV heads.

Returns:

• tuple[Tensor, Tensor] –

  (q [M, num_q_heads * 128] bf16, residual_out [M, 2048] bf16).

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

def flint_qkv_fused(
    x: torch.Tensor,
    residual: torch.Tensor | None,
    norm_weight: torch.Tensor,
    eps: float,
    weight: torch.Tensor,
    weight_scale: torch.Tensor | None,
    global_scale: torch.Tensor | None,
    q_norm_weight: torch.Tensor,
    k_norm_weight: torch.Tensor,
    qk_eps: float,
    cos_sin_cache: torch.Tensor | None,
    positions: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    num_q_heads: int,
    num_kv_heads: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Residual-add RMSNorm + qkv GEMV + q/k norm, RoPE, cache write for
    ``M <= 8`` tokens.

    Args:
        x: [M, 2048] bf16 attention input (contiguous).
        residual: [M, 2048] bf16 residual, or None for the first layer.
        norm_weight: Input RMSNorm weight.
        eps: Input RMSNorm epsilon.
        weight: qkv weight, NVFP4 (uint8 [N, 1024]) or BF16 [N, 2048].
        weight_scale: NVFP4 128x4-swizzled block scales (None for BF16).
        global_scale: NVFP4 fp32 global scale (None for BF16).
        q_norm_weight: Per-head q RMSNorm weight.
        k_norm_weight: Per-head k RMSNorm weight.
        qk_eps: q/k RMSNorm epsilon.
        cos_sin_cache: [max_pos, 128] fp32 (cos | sin), or None (no RoPE).
        positions: [M] int64.
        k_cache: [blocks, block_size, kv_heads, 128] view.
        v_cache: Same shape and strides as ``k_cache``.
        slot_mapping: [<= M] int64; negative slots and tokens beyond it are
            not written.
        num_q_heads: Number of q heads.
        num_kv_heads: Number of KV heads.

    Returns:
        (q [M, num_q_heads * 128] bf16, residual_out [M, 2048] bf16).

    """
    m = x.shape[0]
    assert 0 < m <= MAX_M and x.shape[1] == _K and x.is_contiguous()
    bf16w = weight.dtype == torch.bfloat16
    n = (num_q_heads + 2 * num_kv_heads) * _HEAD
    assert weight.shape[0] == n and weight.is_contiguous()
    assert k_cache.stride() == v_cache.stride() and k_cache.stride(-1) == 1
    has_residual = residual is not None
    if residual is not None:
        assert residual.is_contiguous()
        res_out = torch.empty_like(x)
    else:
        residual = res_out = x
    xn = torch.empty_like(x)
    _compile_norm(has_residual)(x, residual, res_out, xn, norm_weight, m, eps)
    q = torch.empty(m, num_q_heads * _HEAD, dtype=x.dtype, device=x.device)
    w = weight.view(torch.uint8).view(-1)
    s = w if weight_scale is None else weight_scale.view(torch.uint8).view(-1)
    gs = global_scale.reshape(1) if global_scale is not None else _one(x.device)
    rope = cos_sin_cache is not None
    if cos_sin_cache is not None:
        assert cos_sin_cache.dtype == torch.float32 and cos_sin_cache.shape[1] == _HEAD
        cs = cos_sin_cache
    else:
        cs = no_rope_table(x.device)
    _compile_qkv(bf16w, rope, num_q_heads, num_kv_heads)(
        xn, w, s, gs, q_norm_weight, k_norm_weight, cs, positions, slot_mapping,
        q, _base(k_cache), _base(v_cache), m, slot_mapping.shape[0], qk_eps,
        k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.shape[1],
    )  # fmt: skip
    return q, res_out

no_rope_table(device)

Placeholder fp32 cos|sin argument for kernels compiled without RoPE.

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

def no_rope_table(device: torch.device) -> torch.Tensor:
    """Placeholder fp32 cos|sin argument for kernels compiled without RoPE."""
    t = _NO_ROPE.get(device)
    if t is None:
        t = _NO_ROPE[device] = torch.zeros(1, _HEAD, device=device)
    return t