CausalMask (Attention)

Mask query-key dot products for causal attention.

Attention(
    dim=768,
    plugins=[
        CausalMask()
    ],
)