Skip to content

vllm.v1.attention.ops.flydsl_kernels.buffer_utils

Buffer-resource (V#) views shared by the FlyDSL decode kernels.

The descriptor sits in SGPRs and each access adds a 32-bit voffset, so offsets into one view must stay below 4 GiB; move the base with base_byte_offset for anything further away.

Functions:

  • buffer_view –

    Buffer view of a kernel tensor argument, indexed in elem elements.

  • load_vec –

    Load n elements at element index of a unit_elems=n view.

  • store_vec –

    Store the n-element vector value at element index.

buffer_view(tensor, elem, unit_elems=1, base_byte_offset=None)

Buffer view of a kernel tensor argument, indexed in elem elements.

view[i] loads or stores one element. With unit_elems > 1, :func:load_vec / :func:store_vec access unit_elems consecutive elements starting at any element index i.

Source code in vllm/v1/attention/ops/flydsl_kernels/buffer_utils.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def buffer_view(tensor, elem, unit_elems=1, base_byte_offset=None):
    """Buffer view of a kernel tensor argument, indexed in ``elem`` elements.

    ``view[i]`` loads or stores one element. With ``unit_elems > 1``,
    :func:`load_vec` / :func:`store_vec` access ``unit_elems`` consecutive
    elements starting at any element index ``i``.
    """
    it = fx.get_iter(tensor)
    if base_byte_offset is not None:
        it = fx.add_offset(fx.recast_iter(fx.Uint8, it), base_byte_offset)
    elem_ptr = fx.PointerType.get(elem.ir_type, fx.AddressSpace.Global, elem.width // 8)
    it = fx.recast_iter(elem_ptr, it)
    if unit_elems == 1:
        layout = fx.make_layout(_VIEW_ELEMS, 1)
    else:
        layout = fx.make_layout((_VIEW_ELEMS, unit_elems), (1, 1))
    return fx.rocdl.make_buffer_tensor(fx.make_view(it, layout))

load_vec(view, index, n, cache_modifier=0)

Load n elements at element index of a unit_elems=n view.

Source code in vllm/v1/attention/ops/flydsl_kernels/buffer_utils.py 
40
41
42
43
44
45
def load_vec(view, index, n, cache_modifier=0):
    """Load ``n`` elements at element ``index`` of a ``unit_elems=n`` view."""
    elem = view.element_type
    frag = fx.make_rmem_tensor(n, elem)
    fx.copy(_copy_atom(elem, n, cache_modifier), fx.slice(view, (index, None)), frag)
    return fx.Vector(frag.load())

store_vec(view, index, value, n)

Store the n-element vector value at element index.

Source code in vllm/v1/attention/ops/flydsl_kernels/buffer_utils.py 
48
49
50
51
52
53
def store_vec(view, index, value, n):
    """Store the ``n``-element vector ``value`` at element ``index``."""
    elem = view.element_type
    frag = fx.make_rmem_tensor(n, elem)
    frag.store(value)
    fx.copy(_copy_atom(elem, n), frag, fx.slice(view, (index, None)))