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())