# ltx-kernels

Custom CUDA/C++ All2All communication kernels for multi-GPU tensor parallelism. Used by the sequence parallel inference path in `ltx-core`.

## Requirements

- CUDA toolkit (nvcc) matching your GPU architecture
- PyTorch with CUDA support
- Linux

## Building

From the repository root:

```bash
uv pip install -e internal/ltx-kernels --no-build-isolation
```

Set `TORCH_CUDA_ARCH_LIST` to target specific architectures (speeds up compilation):

```bash
# H100 only
TORCH_CUDA_ARCH_LIST="9.0" uv pip install -e internal/ltx-kernels --no-build-isolation

# Multiple architectures
TORCH_CUDA_ARCH_LIST="9.0 9.0a 10.0 12.0" uv pip install -e internal/ltx-kernels --no-build-isolation
```

## Testing

Tests require a CUDA GPU:

```bash
uv run pytest internal/ltx-kernels/tests/ -v
```

## Operations

- **send_recv_heads** -- Redistributes attention heads across GPUs (All2All)
- **gather_heads** -- Inverse of send_recv_heads
- **allgather** -- Gathers sequence tokens from all ranks

All operations support BFloat16 and Float8 (e4m3fn) data types.