metadata
tags:
- kernel
license: apache-2.0
my_softmax_function
Fused softmax-attention kernel written in the CUTLASS Python DSL
(cutlass.cute), packaged for the Hugging Face kernels Hub.
Computes, in a single online pass (max-shifted for numerical stability):
out = softmax(scale * Q @ Kᵀ) @ V
Usage
import torch
from kernels import get_kernel
k = get_kernel("monteiro-t/my_softmax_function") # replace with your repo id
Q = torch.randn(128, 64, device="cuda", dtype=torch.float32)
K = torch.randn(256, 64, device="cuda", dtype=torch.float32)
V = torch.randn(256, 64, device="cuda", dtype=torch.float32)
out = k.attention(Q, K, V) # scale defaults to 1/sqrt(d)
out = k.attention(Q, K, V, scale=0.5)
API
| Function | Signature | Notes |
|---|---|---|
attention(Q, K, V, scale=None) |
(M,d),(N,d),(N,d) -> (M,d) |
Convenience wrapper; scale defaults to 1/sqrt(d). |
softmax_attention(Q, K, V, scale) |
registered torch.ops op |
scale required; torch.compile-friendly (has a fake/meta impl). |
Supported shapes / dtypes / hardware
| Property | Support |
|---|---|
| Dtype | float32 only |
| Layout | 2-D, contiguous (inputs are forced contiguous) |
| Shapes | Q: (M, d), K: (N, d), V: (N, d) → out: (M, d) |
| Batch / heads | Not supported (single 2-D attention; no leading batch/head dims) |
| Masking | None (dense attention only) |
| Backward pass | None (inference / forward only) |
| Device | CUDA only |
Limitations
Naive one-row-per-thread implementation intended as a reference/demo kernel. No batching, masking, dropout, or gradients. Not performance-tuned.