File size: 1,761 Bytes
6f4b39f | 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 | ---
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.
|