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.