Skip to content

vllm.cute_utils

Functions:

  • simple_tma_copy –

    A simple helper that wraps group_modes() and tma_partition()

simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None, mcast_mask=None)

A simple helper that wraps group_modes() and tma_partition() NOTE: this should be called WITHOUT cute.elect_one()

Source code in vllm/cute_utils/__init__.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None, mcast_mask=None):
    """A simple helper that wraps group_modes() and tma_partition()
    NOTE: this should be called WITHOUT cute.elect_one()
    """
    g2s_ops = (
        cpasync.CopyBulkTensorTileG2SOp,
        cpasync.CopyBulkTensorTileG2SMulticastOp,
    )
    if isinstance(atom.op, g2s_ops):
        gmem = src
        smem = dst
    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
        smem = src
        gmem = dst
    else:
        raise ValueError

    s_part, g_part = cpasync.tma_partition(
        atom,
        0,
        cute.make_layout(1),
        cute.group_modes(smem, 0),
        cute.group_modes(gmem, 0),
    )

    if isinstance(atom.op, g2s_ops):
        kwargs = {}
        if mcast_mask is not None:
            kwargs["mcast_mask"] = mcast_mask
        cute.copy(
            atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy, **kwargs
        )
    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
        cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
    else:
        raise ValueError