Add Cosmos3-Edge BF16 production gates
Browse files
README.md
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# transformer-layout-primitives
|
| 2 |
+
|
| 3 |
+
Generic BF16 layout, RoPE, and text-state primitives for transformer pipelines.
|
| 4 |
+
|
| 5 |
+
Hub repo: `flashrt/transformer-layout-primitives`
|
| 6 |
+
|
| 7 |
+
## Public API
|
| 8 |
+
|
| 9 |
+
- `fill_neginf_bf16(dst) -> dst`
|
| 10 |
+
- `add_bias_bf16_(data, bias) -> data`
|
| 11 |
+
- `repeat_interleave_heads_bf16(src, repeat) -> bf16`
|
| 12 |
+
- `text_gather_bf16(src, batch, seq) -> bf16`
|
| 13 |
+
- `text_scatter_bf16(dst, src, batch, seq) -> dst`
|
| 14 |
+
- `rope_rotate_half_bf16_(x, cos, sin) -> x`
|
| 15 |
+
- `qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, eps=1e-6) -> qk`
|
| 16 |
+
- `qk_pair_rmsnorm_rope_bf16(q, k, q_weight, k_weight, cos, sin, eps=1e-6) -> (q, k)`
|
| 17 |
+
- `gather_rows_bf16(src, row_indices, out=None) -> bf16`
|
| 18 |
+
- `scatter_rows_bf16(src, row_indices, rows, out=None) -> bf16`
|
| 19 |
+
|
| 20 |
+
The package is intentionally model-neutral. It exposes Tensor APIs for common
|
| 21 |
+
transformer integration gaps: head repeat for GQA/MQA, first/last token gather
|
| 22 |
+
and scatter, bias add, RoPE rotate-half, and fused Q/K RMSNorm+RoPE. The pair
|
| 23 |
+
API handles different Q and KV head counts in one launch.
|
| 24 |
+
|
| 25 |
+
## Example
|
| 26 |
+
|
| 27 |
+
```python
|
| 28 |
+
from kernels import get_kernel
|
| 29 |
+
import torch
|
| 30 |
+
|
| 31 |
+
ops = get_kernel("flashrt/transformer-layout-primitives", version=1)
|
| 32 |
+
|
| 33 |
+
q = torch.randn((128, 32, 128), device="cuda", dtype=torch.bfloat16)
|
| 34 |
+
weight = torch.ones((128,), device="cuda", dtype=torch.bfloat16)
|
| 35 |
+
cos = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
|
| 36 |
+
sin = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
|
| 37 |
+
ops.qk_rmsnorm_rope_bf16_(q, weight, cos, sin)
|
| 38 |
+
|
| 39 |
+
k = torch.randn((128, 8, 128), device="cuda", dtype=torch.bfloat16)
|
| 40 |
+
q, k = ops.qk_pair_rmsnorm_rope_bf16(
|
| 41 |
+
q, k, weight, weight, cos, sin
|
| 42 |
+
)
|
| 43 |
+
```
|
| 44 |
+
|
| 45 |
+
## Shape contract
|
| 46 |
+
|
| 47 |
+
- All tensors are contiguous CUDA BF16 tensors.
|
| 48 |
+
- `repeat_interleave_heads_bf16`: `src` is `(seq, heads, head_dim)`.
|
| 49 |
+
- `text_gather_bf16`: `src` is flattened `(batch * seq, dim)` and returns
|
| 50 |
+
first and last token rows as `(2 * batch, dim)`.
|
| 51 |
+
- `text_scatter_bf16`: writes `(2 * batch, dim)` rows back to first and last
|
| 52 |
+
positions in `(batch * seq, dim)`.
|
| 53 |
+
- RoPE functions use rotate-half layout with `cos/sin` shaped `(seq, head_dim)`
|
| 54 |
+
or `(rows, head_dim)`.
|
| 55 |
+
- `qk_pair_rmsnorm_rope_bf16` accepts Q `(rows, q_heads, head_dim)` and K
|
| 56 |
+
`(rows, kv_heads, head_dim)`. `head_dim` must be even and in `[8, 256]`.
|
| 57 |
+
Q and K may have different head counts.
|
| 58 |
+
- `gather_rows_bf16` and `scatter_rows_bf16` use contiguous CUDA int64 row
|
| 59 |
+
indices. Scatter indices must be unique.
|
| 60 |
+
- Indexed-row validation includes the Cosmos3-Edge production layout
|
| 61 |
+
`128 -> 60` rows with hidden size `2048`, including exact CUDA Graph replay.
|
| 62 |
+
|
| 63 |
+
## Validation
|
| 64 |
+
|
| 65 |
+
Correctness is tested against PyTorch BF16/FP32 reference formulas with exact
|
| 66 |
+
checks for pure layout operations and strict BF16 tolerances for math ops. The
|
| 67 |
+
pair API must additionally match the two staged single-tensor kernels exactly.
|
| 68 |
+
|
| 69 |
+
See `benchmarks/RESULTS.md` for current local RTX 5090 source benchmark data.
|