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, load_vec / 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`,
    `load_vec` / `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)))