README.md

# transformer-layout-primitives

Generic BF16 layout, RoPE, and text-state primitives for transformer pipelines.

Hub repo: `flashrt/transformer-layout-primitives`

## Public API

- `fill_neginf_bf16(dst) -> dst`
- `add_bias_bf16_(data, bias) -> data`
- `repeat_interleave_heads_bf16(src, repeat) -> bf16`
- `text_gather_bf16(src, batch, seq) -> bf16`
- `text_scatter_bf16(dst, src, batch, seq) -> dst`
- `rope_rotate_half_bf16_(x, cos, sin) -> x`
- `qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, eps=1e-6) -> qk`
- `qk_pair_rmsnorm_rope_bf16(q, k, q_weight, k_weight, cos, sin, eps=1e-6) -> (q, k)`
- `gather_rows_bf16(src, row_indices, out=None) -> bf16`
- `scatter_rows_bf16(src, row_indices, rows, out=None) -> bf16`

The package is intentionally model-neutral. It exposes Tensor APIs for common
transformer integration gaps: head repeat for GQA/MQA, first/last token gather
and scatter, bias add, RoPE rotate-half, and fused Q/K RMSNorm+RoPE. The pair
API handles different Q and KV head counts in one launch.

## Example

```python
from kernels import get_kernel
import torch

ops = get_kernel("flashrt/transformer-layout-primitives", version=1)

q = torch.randn((128, 32, 128), device="cuda", dtype=torch.bfloat16)
weight = torch.ones((128,), device="cuda", dtype=torch.bfloat16)
cos = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
sin = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
ops.qk_rmsnorm_rope_bf16_(q, weight, cos, sin)

k = torch.randn((128, 8, 128), device="cuda", dtype=torch.bfloat16)
q, k = ops.qk_pair_rmsnorm_rope_bf16(
    q, k, weight, weight, cos, sin
)
```

## Shape contract

- All tensors are contiguous CUDA BF16 tensors.
- `repeat_interleave_heads_bf16`: `src` is `(seq, heads, head_dim)`.
- `text_gather_bf16`: `src` is flattened `(batch * seq, dim)` and returns
  first and last token rows as `(2 * batch, dim)`.
- `text_scatter_bf16`: writes `(2 * batch, dim)` rows back to first and last
  positions in `(batch * seq, dim)`.
- RoPE functions use rotate-half layout with `cos/sin` shaped `(seq, head_dim)`
  or `(rows, head_dim)`.
- `qk_pair_rmsnorm_rope_bf16` accepts Q `(rows, q_heads, head_dim)` and K
  `(rows, kv_heads, head_dim)`. `head_dim` must be even and in `[8, 256]`.
  Q and K may have different head counts.
- `gather_rows_bf16` and `scatter_rows_bf16` use contiguous CUDA int64 row
  indices. Scatter indices must be unique.
- Indexed-row validation includes the Cosmos3-Edge production layout
  `128 -> 60` rows with hidden size `2048`, including exact CUDA Graph replay.

## Validation

Correctness is tested against PyTorch BF16/FP32 reference formulas with exact
checks for pure layout operations and strict BF16 tolerances for math ops. The
pair API must additionally match the two staged single-tensor kernels exactly.

See `benchmarks/RESULTS.md` for current local RTX 5090 source benchmark data.