Skip to content

vllm.v1.attention.ops.flydsl_kernels.mma_utils

MMA-atom calls on register vectors for the FlyDSL decode kernels.

Functions:

  • mma –

    Issue atom once on per-lane operand vectors; return the new acc.

mma(atom, acc, a, b, **atom_state)

Issue atom once on per-lane operand vectors; return the new acc.

atom_state is forwarded to fx.gemm (e.g. scale_a/scale_b for scaled MFMA).

Source code in vllm/v1/attention/ops/flydsl_kernels/mma_utils.py 
15
16
17
18
19
20
21
22
23
def mma(atom, acc, a, b, **atom_state):
    """Issue ``atom`` once on per-lane operand vectors; return the new acc.

    ``atom_state`` is forwarded to ``fx.gemm`` (e.g. ``scale_a``/``scale_b``
    for scaled MFMA).
    """
    c = _fragment(acc)
    fx.gemm(atom, c, _fragment(a), _fragment(b), c, **atom_state)
    return fx.Vector(c.load())