Skip to content

vllm.models.beam.nvidia.dflash2_context_ops

Functions:

  • store_context_kv –

    Normalize, apply NeoX RoPE, and scatter all draft layers' context K/V.

store_context_kv(kv, weight, positions, rope, cache_ptrs, slots, cache, eps, scale_ptrs=None)

Normalize, apply NeoX RoPE, and scatter all draft layers' context K/V.

Source code in vllm/models/beam/nvidia/dflash2_context_ops.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def store_context_kv(
    kv: torch.Tensor,
    weight: torch.Tensor,
    positions: torch.Tensor,
    rope: torch.Tensor,
    cache_ptrs: torch.Tensor,
    slots: torch.Tensor,
    cache: torch.Tensor,
    eps: float,
    scale_ptrs: torch.Tensor | None = None,
) -> None:
    """Normalize, apply NeoX RoPE, and scatter all draft layers' context K/V."""
    rows, layers, _, heads, dim = kv.shape
    fp8 = cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
    assert cache.dtype == kv.dtype or fp8
    assert not fp8 or scale_ptrs is not None
    if rows:
        _context_kv_kernel[(triton.cdiv(rows, 4), layers * heads)](
            kv,
            weight,
            positions,
            rope,
            cache_ptrs,
            cache,
            scale_ptrs,
            slots,
            rows,
            slots.numel(),
            layers,
            heads,
            dim,
            cache.stride(0),
            cache.stride(1),
            cache.stride(2),
            cache.shape[2],
            eps,
            fp8,
            torch.finfo(cache.dtype).max if fp8 else 0.0,
            triton.next_power_of_2(dim),
            num_warps=4,
            enable_fp_fusion=False,
        )