File size: 2,036 Bytes
490d9fe fcfe221 490d9fe fcfe221 490d9fe fcfe221 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """Public torch custom-op wrapper around the CUTLASS-DSL softmax-attention kernel.
This is the API a Hugging Face `kernels` consumer sees: a normal, registered
`torch.ops` op that works with autograd tracing / `torch.compile`, not the raw
`@cute.kernel`.
"""
import math
import torch
from cutlass.cute.runtime import from_dlpack
from ._ops import add_op_namespace_prefix
from .attention_v3 import solve
@torch.library.custom_op(add_op_namespace_prefix("softmax_attention"), mutates_args=())
def softmax_attention(
Q: torch.Tensor, # (M, d)
K: torch.Tensor, # (N, d)
V: torch.Tensor, # (N, d)
scale: float,
) -> torch.Tensor: # (M, d)
"""Fused, max-shifted softmax attention: softmax(scale * Q @ K^T) @ V.
Q, K, V must be contiguous, 2-D, float32, and on the same CUDA device.
"""
if not (Q.is_cuda and K.is_cuda and V.is_cuda):
raise ValueError("Q, K, V must be CUDA tensors")
if not (Q.dim() == K.dim() == V.dim() == 2):
raise ValueError("Q, K, V must be 2-D (M,d)/(N,d)/(N,d)")
M, d = Q.shape
N = K.shape[0]
if K.shape[1] != d or V.shape[1] != d or V.shape[0] != N:
raise ValueError("shape mismatch between Q/K/V")
Q = Q.contiguous()
K = K.contiguous()
V = V.contiguous()
output = torch.empty((M, d), dtype=torch.float32, device=Q.device)
solve(
from_dlpack(Q),
from_dlpack(K),
from_dlpack(V),
from_dlpack(output),
M,
N,
d,
float(scale),
)
return output
@softmax_attention.register_fake
def _(Q, K, V, scale):
# Shape/dtype/device metadata only — no compute. Lets torch.compile trace.
M, d = Q.shape
return Q.new_empty((M, d))
def attention(Q, K, V, scale=None):
"""Convenience entry point with a default 1/sqrt(d) scale."""
if scale is None:
scale = 1.0 / math.sqrt(Q.shape[-1])
# Call the registered op directly (its namespace is build-unique).
return softmax_attention(Q, K, V, scale)
|