Skip to content

vllm.models.beam.amd.dflash2_context_ops

Functions:

transform_context_kv(kv, weight, positions, rope, eps, value_scale, is_neox)

Apply Beam K normalization/RoPE and V scaling without touching caches.

Source code in vllm/models/beam/amd/dflash2_context_ops.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def transform_context_kv(
    kv: torch.Tensor,
    weight: torch.Tensor,
    positions: torch.Tensor,
    rope: torch.Tensor,
    eps: float,
    value_scale: float,
    is_neox: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply Beam K normalization/RoPE and V scaling without touching caches."""
    rows, layers, _, heads, dim = kv.shape
    keys = torch.empty((layers, rows, heads, dim), dtype=kv.dtype, device=kv.device)
    values = torch.empty_like(keys)
    if rows:
        _transform_context_kv_kernel[(triton.cdiv(rows, 4), layers * heads)](
            kv,
            weight,
            positions,
            rope,
            keys,
            values,
            rows,
            layers,
            heads,
            dim,
            eps,
            value_scale,
            is_neox,
            triton.next_power_of_2(dim),
            num_warps=4,
            enable_fp_fusion=False,
        )
    return keys, values