| --- |
| tags: |
| - kernel |
| license: apache-2.0 |
| --- |
| |
| # my_softmax_function |
|
|
| Fused **softmax-attention** kernel written in the [CUTLASS Python DSL](https://github.com/NVIDIA/cutlass) |
| (`cutlass.cute`), packaged for the Hugging Face [`kernels`](https://github.com/huggingface/kernels) Hub. |
|
|
| Computes, in a single online pass (max-shifted for numerical stability): |
|
|
| ``` |
| out = softmax(scale * Q @ Kᵀ) @ V |
| ``` |
|
|
| ## Usage |
|
|
| ```python |
| 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. |
|
|
