8 days ago
-
K4 (`mlp`) and K1/K3 (`res_layernorm_fold`): -4.8% end-to-end encoder, bitwise identical
Reworks the fused MLP kernel and the residual + LayerNorm kernel. K2, the weight folds and the
layer are unchanged apart from one line in `layers.py` that packs the folded fc1 bias.
### Results (trn2.3xlarge, LNC=2, all outputs `torch.equal` to v1.1)
**Full encoder** (`torch.compile(model.encoder, backend="neuron")`, 24 layers, one 16-frame clip):
| | v1.1 | this PR | change |
|---|---|---|---|
| Wall time, ms/clip (3 rounds, 20 iters each) | 26.632 / 26.635 / 26.639 | 25.356 / 25.357 / 25.347 | **-1.28 ms (-4.8%)** |
| Device time (one profiled forward) | 26.83 ms | 25.50 ms | -5.0% |
| `mfu_estimated_percent` | 39.5% | 41.7% | +2.2 points |
**Standalone at the encoder shape** (2 interleaved rounds each):
| Kernel | v1.1 | this PR | change | efficiency |
|---|---|---|---|---|
| K4 `mlp` ([2048, 1024] -> 4096 -> 1024) | 295.3 us | 248.6 us | -15.8% | MFU 70.3% -> 82.8% |
| K1 `res_layernorm_fold`, 1 branch | 48.0 us | 41.9 us | -12.7% | HBM 73% -> 84% of 716 GB/s |
| K3 `res_layer…
17 days ago
-
v1: V-JEPA2 ViT-L encoder NKI kernels for Trainium2 (4.35x, 44% MFU)
