vllm.models.beam.nvidia.attention_ops

Fused Q/K normalization, position transforms, and attention output gating.

Functions:

• beam_output_gate –

  attn_output * sigmoid(gate + gate_bias) with activation-dtype rounding

• beam_qk_norm_rope –

  Round norm before position transforms and optionally insert BF16 K/V.

beam_output_gate(attn_output, gate, gate_bias, *, num_heads, head_dim)

attn_output * sigmoid(gate + gate_bias) with activation-dtype rounding after the bias add and the sigmoid.

Source code in vllm/models/beam/nvidia/attention_ops.py

def beam_output_gate(
    attn_output: torch.Tensor,
    gate: torch.Tensor,
    gate_bias: torch.Tensor,
    *,
    num_heads: int,
    head_dim: int,
) -> torch.Tensor:
    """``attn_output * sigmoid(gate + gate_bias)`` with activation-dtype rounding
    after the bias add and the sigmoid."""
    output = torch.empty_like(attn_output)
    if output.numel():
        if (
            attn_output.dtype == gate.dtype == gate_bias.dtype == torch.bfloat16
            and current_platform.is_device_capability_family(100)
        ):
            from .attention_ops_cutedsl import run

            run(attn_output, gate, gate_bias, num_heads, head_dim, output)
            return output
        _output_gate_kernel[(triton.cdiv(output.numel(), 256),)](
            attn_output,
            gate,
            gate_bias,
            output,
            output.numel(),
            gate.stride(0),
            num_heads,
            head_dim,
            256,
            enable_fp_fusion=False,
        )
    return output

beam_qk_norm_rope(q, k, q_weight, k_weight, eps, positions, cos_sin_cache, head_dim, rotary_dim=0, is_neox_style=True, *, ssmax=False, value=None, kv_cache=None, slot_mapping=None)

Round norm before position transforms and optionally insert BF16 K/V.

Source code in vllm/models/beam/nvidia/attention_ops.py

def beam_qk_norm_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
    positions: torch.Tensor,
    cos_sin_cache: torch.Tensor | None,
    head_dim: int,
    rotary_dim: int = 0,
    is_neox_style: bool = True,
    *,
    ssmax: bool = False,
    value: torch.Tensor | None = None,
    kv_cache: torch.Tensor | None = None,
    slot_mapping: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Round norm before position transforms and optionally insert BF16 K/V."""
    if kv_cache is not None:
        assert value is not None and slot_mapping is not None
        assert kv_cache.dtype == q.dtype == torch.bfloat16
        assert kv_cache.ndim == 4 and kv_cache.stride(-1) == 1
    q_out, k_out = torch.empty_like(q), torch.empty_like(k)
    q_heads, k_heads = q.shape[-1] // head_dim, k.shape[-1] // head_dim
    if q.shape[0]:
        if rotary_dim or q.shape[0] < 512:
            tile, warps = 4, 4
        elif q.shape[0] <= 2048:
            tile, warps = 16, 8
        else:
            tile, warps = 32, 4
        _qk_norm_rope_kernel[(triton.cdiv(q.shape[0], tile), q_heads + k_heads)](
            q,
            k,
            q_weight,
            k_weight,
            positions,
            cos_sin_cache,
            q_out,
            k_out,
            q.shape[0],
            value,
            kv_cache,
            slot_mapping,
            slot_mapping.numel() if slot_mapping is not None else 0,
            kv_cache is not None,
            value.stride(0) if value is not None else 0,
            kv_cache.stride(0) if kv_cache is not None else 0,
            kv_cache.stride(1) if kv_cache is not None else 0,
            kv_cache.stride(2) if kv_cache is not None else 0,
            kv_cache.shape[2] if kv_cache is not None else 1,
            tile,
            q_heads,
            k_heads,
            head_dim,
            q.stride(0),
            k.stride(0),
            eps,
            rotary_dim,
            is_neox_style,
            ssmax,
            triton.next_power_of_2(head_dim),
            num_warps=warps,
            enable_fp_fusion=False,
        )
    return q_out, k_out