TorchEmbed is now integrated into DeepSpeed and ships as a dependency. The fused Triton RoPE kernel replaces the rotary embedding path in DeepSpeed’s transformer engine, delivering 4–7× speedup over pure PyTorch with no code changes required from users.

Why RoPE is a bottleneck

Rotary Position Embedding (RoPE) is now the dominant positional encoding scheme across modern LLMs — Llama, Mistral, Qwen, and most production models use it. At scale, the naive PyTorch implementation becomes a memory-bandwidth bottleneck: it materializes intermediate tensors from chunk and cat operations on every forward pass. For a sequence of length 4096, that means allocating and reading back several hundred MB of intermediate state per layer.

torch.compile helps but doesn’t fully eliminate the problem. Even with fusion, compiled code still allocates intermediate tensors that persist across the operation. The fused Triton kernel in TorchEmbed sidesteps this entirely — it reads and writes each element exactly once.

The integration

DeepSpeed’s transformer engine automatically uses TorchEmbed’s fused kernel when Triton is available. Installing TorchEmbed with the Triton extra is all that’s needed:

pip install torchembed[triton]

DeepSpeed picks up the kernel at import time. No config changes, no monkey-patching. If Triton is unavailable (e.g. CPU-only environments), TorchEmbed falls back to the standard PyTorch path transparently.

Using TorchEmbed directly

If you want to use the fused kernel outside of DeepSpeed, the API is minimal:

from torchembed.positional import RotaryEmbedding

rope = RotaryEmbedding(dim=64, use_fused=True).to("cuda")
q, k = rope(q.cuda(), k.cuda())

Setting use_fused=True routes through the Triton kernel. The interface is otherwise identical to a standard PyTorch module — it drops into any attention implementation without modifications.

Benchmark results

Measured on NVIDIA GB10 hardware at float16 across representative tensor shapes:

Shape (B, H, S, D, dim)	PyTorch	torch.compile	TorchEmbed (Triton)	Speedup
(1, 32, 2048, 128, 128)	1.364 ms	0.620 ms	0.305 ms	4.47×
(1, 32, 4096, 128, 128)	2.946 ms	1.284 ms	0.579 ms	5.09×
(1, 32, 8192, 128, 128)	5.909 ms	2.476 ms	1.224 ms	4.83×
(2, 32, 2048, 128, 128)	2.949 ms	1.293 ms	0.630 ms	4.68×
(1, 32, 2048, 256, 128)	2.869 ms	1.323 ms	0.646 ms	4.44×

Peak speedup is 5.09× at sequence length 4096 — the sweet spot for most production inference workloads. Speedups hold across batch sizes and head dimensions.

Other embedding types

TorchEmbed is not limited to RoPE. The library also covers patch embeddings for Vision Transformers and cyclic temporal encodings for time-series models:

# Patch embeddings for ViT
from torchembed.patch import PatchEmbedding

patch_emb = PatchEmbedding(image_size=224, patch_size=16, embed_dim=768)
tokens = patch_emb(images)  # (B, 3, 224, 224) → (B, 196, 768)

# Cyclic temporal encoding
from torchembed.temporal import CyclicEmbedding

hour_enc = CyclicEmbedding(period=24)
features = hour_enc(hour_tensor)

Get started

TorchEmbed requires Python ≥ 3.9 and PyTorch ≥ 2.0. No other dependencies beyond PyTorch itself.

pip install torchembed        # CPU / fallback path
pip install torchembed[triton] # with fused Triton kernels

Source and benchmarking scripts are on GitHub.