Triton Sigmoid Attention

Fused sigmoid attention kernels for NVIDIA GPUs built with Triton.

Overview

Two implementations:

• Dense - Same-length sequences (fastest)
• Padded - Variable-length with padding masks

Features torch.compile support and causal masking. Requires Ampere architecture or newer.

Installation

pip install triton-sigmoid

Development install from source

git clone https://github.com/MSDLLCpapers/triton-sigmoid.git
cd triton-sigmoid
uv sync --extra dev
source .venv/bin/activate

Quick Start

Dense Attention (Same-Length Sequences)

import torch
from triton_sigmoid import sigmoid_attention

batch, seq_len, n_heads, head_dim = 2, 1024, 8, 64
q = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)
k = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)
v = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)

output = sigmoid_attention(q, k, v, is_causal=False)
output_causal = sigmoid_attention(q, k, v, is_causal=True)

Padded Attention (Variable-Length Sequences)

import torch
from triton_sigmoid import sigmoid_attention_padded

batch, seq_len, n_heads, head_dim = 2, 1024, 8, 64
q = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)
k = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)
v = torch.randn(batch, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)

seq_lens_k = torch.tensor([800, 950], device='cuda', dtype=torch.int32)
seq_lens_q = torch.tensor([800, 950], device='cuda', dtype=torch.int32)

output = sigmoid_attention_padded(q, k, v, seq_lens_k=seq_lens_k, seq_lens_q=seq_lens_q)

compiled_fn = torch.compile(sigmoid_attention_padded)
output = compiled_fn(q, k, v, seq_lens_k=seq_lens_k, seq_lens_q=seq_lens_q, is_causal=True)

Performance

See benchmarks/README.md for details.

Documentation

• API Reference - Detailed function signatures and parameters
• Algorithm Details - Implementation details and sigmoid variants
• Benchmarking Guide - Performance benchmarks and how to run them
• Testing Guide - Running tests and test structure

License

MIT License - see LICENSE file for details.

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

Citation

If you use this work in your research, please cite our paper:

@misc{sadashivaiah2026bettermodelsfastertraining,
      title={Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models}, 
      author={Vijay Sadashivaiah and Georgios Dasoulas and Judith Mueller and Soumya Ghosh},
      year={2026},
      eprint={2604.27124},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2604.27124}, 
}

Acknowledgments

• Built with Triton
• Inspired by Flash Attention