IHA-CUDA

KO | EN

gitlite — search

#python #hacktoberfest #typescript #go #machine-learning #rust #java #nodejs #cli #security #openai #ios #mcp #open-source

(★ 10)

Interleaved Head Attention CUDA kernels for NVIDIA Hopper (SM90) - built on FlashAttention-3

IHA-CUDA 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# IHA-CUDA: Interleaved Head Attention Kernels for Hopper GPUs

GPU kernels for **Interleaved Head Attention (IHA)** ([arXiv:2602.21371](https://arxiv.org/abs/2602.21371)) on NVIDIA Hopper (SM90).
IHA extends standard Multi-Head Attention by mixing information across
attention heads via learned weights before computing attention, then applying
**joint softmax** across P pseudo-heads. The kernels are built on
[FlashAttention-3](https://github.com/Dao-AILab/flash-attention) using the
[CuTe DSL](https://github.com/NVIDIA/cutlass) (cutlass-dsl 4.4.1), with
pre-mixing kernels written in [Triton](https://github.com/triton-lang/triton).

## Kernel Design

![FlashAttention-3 Optimization](https://raw.githubusercontent.com/saisuryadv/IHA-CUDA/master/assets/fa3_optimization.png)

**Left (Naive FA3):** A naive application of FlashAttention-3 leads to many threadblocks (CTAs), one per query token group, iterating over key-value pairs with a sliding window.

**Right (Optimized IHA):** Our modified FA3 CuTe DSL kernel launches fewer CTAs by using a nested loop over both query pseudo-tokens and key-value pairs. A separate Triton kernel handles pseudo-token construction (memory-bound), since naive `torch.einsum` is suboptimal for this operation.

### Pseudo-Token Mixing (Triton Premix Kernels)

Before attention, IHA mixes Q, K, V across heads via learned weights:
`Out[h,p] = Σ_s W[h,p,s] * X[s]` — a linear map from H source heads to
H×P pseudo-head outputs, applied independently at every (batch, position,
feature) coordinate. This is the "premix" step.

We implement this as fused Triton kernels rather than `torch.einsum` because
the operation is memory-bound with a tiny weight tensor (H×P×H scalars) and
large input/output tensors (B×N×H×d). The Triton kernels tile over sequence
positions and unroll the head loops at compile time via `tl.static_range`.

**Forward** (`premix.py`): One kernel per Q/K/V. Each thread block handles
one (batch, pseudo-head, output-head) slice and tiles over N. Supports
optional fused RoPE rotation.

**Backward** (`premix_bwd.py`): Two fused kernels handle all three Q/K/V
in a single launch each (6→2 kernel launches):
- `_premix_dx_3qkv_kernel`: Computes dX = Wᵀ dY (transpose of forward map)
- `_premix_dw_3qkv_kernel`: Computes dW = Σ_{b,n,d} dY · Xᵀ (weight
  gradients via atomic reduction)

The premix is the main source of IHA overhead relative to MHA during training. Its cost
scales as O(H² × P × N) per layer, while attention scales as O(N² × H).
The ratio premix/attention ∝ H×P/N, so increasing sequence length
amortizes the premix overhead.

## End-to-End Training Benchmark

We profile a **2.4B-parameter LLaMA** model (hidden=2560, heads=20, d=128,
28 layers, MHA) by comparing FlashAttention-3 (FA3) with FLOP-matched IHA
across sequence lengths and IHA layer fractions. Both the FA3 baseline and IHA
use the same CuTe DSL attention kernels on SM90; IHA layers additionally
perform pseudo-token mixing via Triton kernels. Both models use gradient
checkpointing, batch size 1, AdamW optimizer, and a `DataLoader` with
`pin_memory` transfers. P=4 pseudo-heads are used for IHA layers. The window
size (WSL) is set to tile-match the number of KV tile visits between IHA and
MHA causal, ensuring equal computational work.

| Frac | N | WSL | Train FA3 | Train IHA | Ratio | Tput FA3 | Tput IHA | Ratio |
|------|-------|------|-----------|-----------|-------|----------|----------|-------|
| 50% | 8K | 1023 | 448 ms | 527 ms | 1.18x | 18,281 | 15,546 | 0.85x |
| 50% | 16K | 2049 | 953 ms | 1,111 ms | 1.17x | 17,193 | 14,745 | 0.86x |
| 50% | 32K | 4225 | 2,328 ms | 2,639 ms | 1.13x | 14,075 | 12,416 | 0.88x |
| 50% | 64K | 8577 | 6,493 ms | 7,133 ms | 1.10x | 10,093 | 9,188 | 0.91x |
| 80% | 8K | 1023 | 449 ms | 574 ms | 1.28x | 18,253 | 14,278 | 0.78x |
| 80% | 16K | 2049 | 951 ms | 1,205 ms | 1.27x | 17,225 | 13,593 | 0.79x |
| 80% | 32K | 4225 | 2,324 ms | 2,818 ms | 1.21x | 14,100 | 11,626 | 0.82x |
| 80% | 64K | 8577 | 6,527 ms | 7,485 ms | 1.15x | 10,040 | 8,756 | 0.87x |

**Frac**: fraction of attention layers replaced by IHA (the rest use FA3).
**N**: sequence length. **WSL**: sliding-window size in IHA layers.
**Train**: per-step latency (forward + backward + optimizer). **Tput**:
throughput in tokens/sec.

The overhead decreases with sequence length as the premix cost is
increasingly amortized by the quadratically scaling attention computation.

### How to reproduce

**Hardware**: NVIDIA GH200 120GB (SM90, Hopper).

**Environment setup**:
```bash
# Python 3.11+ with PyTorch 2.x + CUDA
pip install nvidia-cutlass-dsl    # CuTe DSL for kernel compilation
pip install triton                # Triton for premix kernels

# Clone this repository and the upstream FlashAttention-3
git clone <this-repo> IHA_CUDA
git clone https://github.com/Dao-AILab/flash-attention flash-attention-upstream
```

The IHA kernels (forward, backward) and premix kernels live in this repository
under `flash_attn/cute/`. The upstream FlashAttention-3 repository provides
the base FA3 kernel classes that IHA subclasses. Both paths are resolved
automatically by the benchmark script.

**Run the benchmark**:
```bash
# 50% IHA layers, N=64K, P=4, tile-matched wsl=8577
python bench_llama_iha.py \
    --hidden 2560 --heads 20 --layers 28 \
    --P 4 --B 1 --N 65536 --wsl 8577 \
    --iha_frac 0.5 --n_iter 3

# 80% IHA layers, N=32K, P=4, tile-matched wsl=4225
python bench_llama_iha.py \
    --hidden 2560 --heads 20 --layers 28 \
    --P 4 --B 1 --N 32768 --wsl 4225 \
    --iha_frac 0.8 --n_iter 3
```

The benchmark script (`bench_llama_iha.py`) is self-contained. It builds three
HuggingFace `LlamaForCausalLM` models from the same config: an SDPA baseline,
an FA3 baseline (all layers use CuTe FA3 kernels), and an IHA model (selected
layers use IHA, the rest use FA3). CuTe kernels are compiled on first run,
then the script measures prefill, decode, and training step times.

**Baseline**: The FA3 baseline replaces every `LlamaAttention` layer with
`FA3LlamaAttention`, which calls the upstream FlashAttention-3 CuTe forward
and backward kernels via a `torch.autograd.Function` wrapper. The IHA model
additionally replaces selected layers with `IHALlamaAttention`, which adds
premixing (Triton) before calling the IHA v3 CuTe kernel (a subclass of FA3).

## FLOP Matching

IHA uses P pseudo-heads with a sliding window of size W = WSL + 1. For a fair
comparison, we set WSL so that the total number of KV tile visits in the IHA
kernel matches the MHA causal kernel. The approximate formula is
`WSL = N/(2P) - 1`, but discrete tile boundaries (128-token tiles) cause small
deviations at large N. We precompute exact tile-matched WSL values:

| N | Formula WSL | Tile-matched WSL |
|-------|-------------|------------------|
| 8192 | 1023 | 1023 |
| 16384 | 2047 | 2049 |
| 32768 | 4095 | 4225 |
| 65536 | 8191 | 8577 |

## Repository Structure

```
flash_attn/cute/
  flash_fwd_iha_v3.py       # IHA v3 forward kernel (CuTe DSL, SM90)
  flash_bwd_iha_v3.py       # IHA v3 backward kernel (CuTe DSL, SM90)
  premix.py                 # Triton premix forward: Q/K/V head mixing
  premix_bwd.py             # Triton premix backward: dX + dW gradients
  flash_fwd.py              # FA3 forward base class (upstream)
  flash_bwd_sm90.py         # FA3 backward base class (upstream)
  flash_bwd_preprocess.py   # Backward preprocessing (dPsum, lse)
  flash_bwd_postprocess.py  # Backward postprocessing (dQaccum -> dQ)
  interface.py              # Python entry points
  block_info.py             # Tile range computation (causal/windowed)
  mask.py                   # Attention masking
  softmax.py                # Online softmax

bench_llama_iha.py          # End-to-end LLaMA training benchmark (this table)
run_p8_sweep.py             # Sweep script for multiple configs
test_bwd_iha_v3.py          # Backward correctness tests
test_bwd_iha_v3_bench.py    # Kernel-level forward + backward benchmarks
test_quick.py               # Forward correctness tests
```

## Requirements

- NVIDIA Hopper GPU (SM90, e.g. GH200)
- Python 3.11+
- PyTorch 2.x with CUDA 12+
- `nvidia-cutlass-dsl >= 4.4.1`
- `triton`
- `transformers` (for the LLaMA benchmark)

IHA-CUDA

IHA-CUDA 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

IHA-CUDA

IHA-CUDA 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README