README: re-measured perf on torch 2.13.0; exact chunk ~1.3x; chunked-wrapper dg caveat quantified; compiled decode step; bwd_batched from ~16 items

Browse files

Files changed (2) hide show

README.md +63 -37
build/torch-neuron/__init__.py +4 -4

README.md CHANGED Viewed

@@ -60,11 +60,14 @@ Verify with `python -c "from kernels.backends import _backend; print(_backend().
   scalar-mean decay approximation for the intra-chunk term — **see the warning below**.
 - **`kda_chunk_step_exact(q, k, v, beta, g, state_in)`** — numerically exact
   per-channel prefill (sub-chunk + WY reformulation). Use this when the model's gate
-  decay is non-trivial (see warning below). ~1.05× the latency of `kda_chunk_step`.
 - **`kda_chunk_step_exact_multihead(q, k, v, beta, g, state_in)`** — head-interleaved
   exact prefill over `NV` heads (`[NV, C, dk]` shapes).
  - **`kda_decode_batch(q, k, v, g, beta, state_in)`** — batched multi-(request, head)
    decode; advances all `B*nv` items one token in a single call. Shapes `[B, nv, dk]`.
  - **`kda_chunk_step_exact_bwd(q, k, v, beta, g, state_in, d_output, dS_final)`** —
    exact chunked BACKWARD (gradient of `kda_chunk_step_exact`). Returns
    `dq, dk, dv, dg, dbeta, dstate_in`, matching fla-core autograd at **cos_sim 1.0
@@ -82,20 +85,23 @@ Verify with `python -c "from kernels.backends import _backend; print(_backend().
   state carry-over across chunks.
 - **`kda_chunked_fused(q, k, v, g, beta, initial_state=None)`** → `(output, final_state)`.
   Same numerics as `kda_chunked`, but processes all chunks in one NKI launch per
-  direction. Faster on multi-chunk sequences. **Recommended for training.**
  - Raw kernels also exported: `kda_recurrent_fwd_v2`, `kda_chunk_step_v2`,
    `kda_recurrent_bwd`, `kda_chunk_bwd`, `kda_fused_chunked_fwd`, `kda_fused_chunked_bwd`.
  - **`kda_recurrent_bwd_batched(q, k, v, g, beta, state_stack, d_output, dS_final)`**
    — batched recurrent backward over all `B*nv` (request, head) items in one launch
-   (flat work list + ping-pong state). ~1.1–1.24× over a per-(b,h) loop; all six
    gradients bit-match the single-item `kda_recurrent_bwd`.
 ## Requirements
 - **Hardware**: AWS Trainium (tested on trn2.3xlarge).
-- **SDK / runtime**: PyTorch Native (`device="neuron"`), torch-neuronx 2.11+, PyTorch 2.11+.
 - **NKI** ≥ 0.4.0.
-- **`kernels`** ≥ 0.15.2 (to load via `get_kernel`).
 - **`transformers`** with `KernelConfig` support, if using the `KernelConfig` path.
 ## Usage
@@ -151,6 +157,10 @@ loss.backward()   # gradients flow through the NKI backward
 # q.grad, k.grad, v.grad, g.grad, beta.grad now populated
 ```
 Kernels operate per (batch, head); loop `B*H` in the caller.
 ## Input contract
@@ -200,13 +210,15 @@ cosine similarity vs the fla-core reference:
 **If your model has non-trivial gate decay, use `kda_chunk_step_exact`** (or
 `kda_chunk_step_exact_multihead`), which is numerically exact (cos_sim ≥ 0.9999999
-across all gate regimes) at ~1.05× the latency. The recurrent kernels
 (`kda_recurrent_fwd`, `kda_recurrent`) are exact in all regimes.
 Because `kda_chunked` / `kda_chunked_fused` differentiate the approximate forward,
-their `dg` gradient is the exact gradient of the *approximate* forward — self-consistent
-for training with these kernels, but not equal to the exact-per-channel `dg` unless the
-approximation is accurate (i.e. small gate decay).
 ## Parity
@@ -238,43 +250,57 @@ use `kda_chunk_step_exact_bwd` to fix it).
 ## Performance
-Measured on trn2.3xlarge, LNC=2, single logical core, single (batch, head) invocation.
-### Prefill (chunked)
-| Metric | Value |
-|--------|-------|
-| Wall-clock per chunk (C=128) | 87 μs |
-| Per-token effective | 0.68 μs |
-| Achieved TFLOPS | 1.89 |
-| MFU (BF16 peak 158 TFLOPS/LNC=2) | 1.19% |
-| MBU (empirical 218 GB/s/LNC=2) | 3.11% |
-### Decode (recurrent)
-| Metric | Value |
-|--------|-------|
-| Wall-clock per call (S=128) | 527 μs |
-| Per-token (amortized) | 4.1 μs |
-The recurrent kernel is overhead-dominated at small sequence lengths. For real decode
-throughput, batch tokens (or requests via `kda_decode_batch`): per-token wall-clock
-drops from ~70 μs at S=1 to ~6 μs at S=128, and `kda_decode_batch` amortizes launch
-overhead across a whole serving batch.
-### Training (fused vs Python chunk loop, fwd+bwd)
-| S | Chunks | `kda_chunked` (loop) | `kda_chunked_fused` |
-|---|--------|----------------------|---------------------|
-| 256 | 2 | 695 μs | 336 μs |
-| 512 | 4 | 1320 μs | 338 μs |
-| 1024 | 8 | 2588 μs | 529 μs |
-The fused path pays per-launch overhead once instead of per chunk, so its advantage
-grows with sequence length. Prefer `kda_chunked_fused` for training.
-MFU/MBU denominators are per LNC=2 core on trn2 (NeuronCore-v3), from the AWS
-Trainium2 architecture guide and empirical measurements.
 ## Not provided

   scalar-mean decay approximation for the intra-chunk term — **see the warning below**.
 - **`kda_chunk_step_exact(q, k, v, beta, g, state_in)`** — numerically exact
   per-channel prefill (sub-chunk + WY reformulation). Use this when the model's gate
+  decay is non-trivial (see warning below). ~1.3× the latency of `kda_chunk_step`.
 - **`kda_chunk_step_exact_multihead(q, k, v, beta, g, state_in)`** — head-interleaved
   exact prefill over `NV` heads (`[NV, C, dk]` shapes).
  - **`kda_decode_batch(q, k, v, g, beta, state_in)`** — batched multi-(request, head)
    decode; advances all `B*nv` items one token in a single call. Shapes `[B, nv, dk]`.
+   9.4-19.5× faster than one launch per item. **But** on the tested build a `torch.compile`d
+   plain-PyTorch batched step (three `einsum`s, below) is 1.8-6.6× faster than this kernel at
+   B=1-32 × 32 heads; prefer that when your decode step is a static batched function.
  - **`kda_chunk_step_exact_bwd(q, k, v, beta, g, state_in, d_output, dS_final)`** —
    exact chunked BACKWARD (gradient of `kda_chunk_step_exact`). Returns
    `dq, dk, dv, dg, dbeta, dstate_in`, matching fla-core autograd at **cos_sim 1.0
   state carry-over across chunks.
 - **`kda_chunked_fused(q, k, v, g, beta, initial_state=None)`** → `(output, final_state)`.
   Same numerics as `kda_chunked`, but processes all chunks in one NKI launch per
+  direction. Faster on multi-chunk sequences. **Its `dg` is not the exact gradient** (cos
+  0.60-0.73 vs exact, even at g=0.01; see the warning below). For exact training use
+  `kda_recurrent`, or `kda_chunk_step_exact` + `kda_chunk_step_exact_bwd`.
  - Raw kernels also exported: `kda_recurrent_fwd_v2`, `kda_chunk_step_v2`,
    `kda_recurrent_bwd`, `kda_chunk_bwd`, `kda_fused_chunked_fwd`, `kda_fused_chunked_bwd`.
  - **`kda_recurrent_bwd_batched(q, k, v, g, beta, state_stack, d_output, dS_final)`**
    — batched recurrent backward over all `B*nv` (request, head) items in one launch
+   (flat work list + ping-pong state). Use from ~16 items: 0.80× at 4 items, 0.84× at 8,
+   1.40× at 16, 1.33× at 32 vs a per-(b,h) loop (S=8). All six
    gradients bit-match the single-item `kda_recurrent_bwd`.
 ## Requirements
 - **Hardware**: AWS Trainium (tested on trn2.3xlarge).
+- **SDK / runtime**: PyTorch Native (`device="neuron"`). Tested: torch 2.13.0 / torch-neuronx 2.13.3 / neuronx-cc 2.0.404056 / NKI 0.7.0b1.
 - **NKI** ≥ 0.4.0.
+- **`kernels`** ≥ 0.17 (to load via `get_kernel`).
 - **`transformers`** with `KernelConfig` support, if using the `KernelConfig` path.
 ## Usage
 # q.grad, k.grad, v.grad, g.grad, beta.grad now populated
 ```
+> `g.grad` from `kda_chunked_fused` is the gradient of the approximate forward (cos 0.60-0.73
+> vs exact on the tested build). If the gate is trained, use `kda_recurrent` instead (same call
+> signature) or the exact chunk kernels.
 Kernels operate per (batch, head); loop `B*H` in the caller.
 ## Input contract
 **If your model has non-trivial gate decay, use `kda_chunk_step_exact`** (or
 `kda_chunk_step_exact_multihead`), which is numerically exact (cos_sim ≥ 0.9999999
+across all gate regimes) at ~1.3× the latency. The recurrent kernels
 (`kda_recurrent_fwd`, `kda_recurrent`) are exact in all regimes.
 Because `kda_chunked` / `kda_chunked_fused` differentiate the approximate forward,
+their `dg` gradient is the gradient of the *approximate* forward, not the exact
+per-channel `dg`. Measured on the tested build, it is **cos 0.60-0.73 vs the exact
+gradient even at g=0.01**, so do not rely on it for training a model with per-channel
+gates. Use `kda_recurrent`, or `kda_chunk_step_exact` + `kda_chunk_step_exact_bwd`
+(all six gradients cos ≥ 0.9999997 in every gate regime).
 ## Parity
 ## Performance
+Measured on trn2.3xlarge, LNC=2, single logical core, 32 heads × 128 (dk = dv = 128).
+Build: torch 2.13.0 / torch-neuronx 2.13.3 / neuronx-cc 2.0.404056 / NKI 0.7.0b1. Median device time with `torch.neuron.synchronize()` before each timer stop.
+Measured on a pre-release PyTorch Native build; re-run on the release with `bash run.sh` from the companion benchmark repo.
+### Prefill (one 128-token chunk, single (batch, head))
+| Kernel | Wall-clock per chunk |
+|--------|---------------------:|
+| `kda_chunk_step` (approximate) | 248 μs |
+| `kda_chunk_step_exact` | 318 μs (~1.3×) |
+`kda_chunk_step_exact_multihead`, per head: 304 / 180 / 123 / 97 μs at NV = 1 / 2 / 4 / 8,
+vs 340 / 204 / 156 / 148 μs for a per-head loop (1.53× at NV=8).
+### Decode
+| Kernel | Wall-clock per call |
+|--------|--------------------:|
+| `kda_recurrent_fwd`, S=1 | 297 μs |
+| `kda_recurrent_fwd`, S=128 | 1,078 μs |
+One decode step over B requests × 32 heads (μs per step):
+| B | per-item launches | `kda_decode_batch` | compiled PyTorch step |
+|--:|------------------:|-------------------:|----------------------:|
+| 1 | 3,649 | 387 | **213** |
+| 8 | 29,157 | 1,581 | **372** |
+| 32 | 118,191 | 6,051 | **911** |
+The compiled step is:
+```python
+def decode_step(q, k, v, g, beta, S):          # [B, nv, D], state [B, nv, D, D]
+    S = S * g.exp().unsqueeze(-1)
+    kv = torch.einsum("bhk,bhkv->bhv", k, S)
+    S = S + torch.einsum("bhk,bhv->bhkv", beta * k, v - kv)
+    return torch.einsum("bhk,bhkv->bhv", q, S), S
+step = torch.compile(decode_step, backend="neuron", dynamic=False)
+```
+### Training (fwd+bwd)
+| S | `kda_chunked_fused` |
+|---|--------------------:|
+| 256 | 1.41 ms |
+| 512 | 1.89 ms |
+| 1024 | 2.72 ms |
+`kda_chunked` (Python chunk loop) at S=1024: 7.22 ms. `kda_recurrent` at S=128: 4.24 ms.
+See the warning above before using the chunked wrappers for training.
 ## Not provided

build/torch-neuron/__init__.py CHANGED Viewed

@@ -37,8 +37,8 @@ Functional warning:
     regimes. See the README for details.
 Requirements:
-    - PyTorch Native (device="neuron"), torch-neuronx 2.11+, PyTorch 2.11+
-    - NKI >= 0.4.0
 License: Apache-2.0. Inference/training runtime kernel package (not a model).
 The fla-core algorithm reference is MIT-licensed and compatible.
@@ -66,7 +66,7 @@ from .kda_fused_chunked_bwd import kda_fused_chunked_bwd
 # amortizes launch overhead across a serving batch.
 # - kda_chunk_step_exact: numerically exact per-channel prefill
 # (cos_sim >= 0.9999999 in all gate regimes) vs the scalar-mean approximation
-# in kda_chunk_step. ~1.05x the latency of the approximation.
 # - kda_chunk_step_exact_multihead: head-interleaved exact prefill.
 from .nki_kda_decode_batch import kda_decode_batch
 from .nki_kda_chunked_exact import kda_chunk_step_exact
@@ -84,7 +84,7 @@ from .nki_kda_chunked_exact_bwd import kda_chunk_step_exact_bwd
 # Batched recurrent BACKWARD: advances all B*nv (request, head) items' reverse-token
 # backward in one launch (flat work list + ping-pong state), mirroring the batched
-# forward decode. ~1.1-1.24x over the per-(b,h) loop; all 6 gradients bit-match the
 # single-item kernel.
 from .nki_kda_recurrent_bwd_batched import kda_recurrent_bwd_batched

     regimes. See the README for details.
 Requirements:
+    - PyTorch Native (device="neuron"); tested on torch 2.13.0 / torch-neuronx 2.13.3
+    - NKI >= 0.4.0 (tested 0.7.0b1)
 License: Apache-2.0. Inference/training runtime kernel package (not a model).
 The fla-core algorithm reference is MIT-licensed and compatible.
 # amortizes launch overhead across a serving batch.
 # - kda_chunk_step_exact: numerically exact per-channel prefill
 # (cos_sim >= 0.9999999 in all gate regimes) vs the scalar-mean approximation
+# in kda_chunk_step. ~1.3x the latency of the approximation.
 # - kda_chunk_step_exact_multihead: head-interleaved exact prefill.
 from .nki_kda_decode_batch import kda_decode_batch
 from .nki_kda_chunked_exact import kda_chunk_step_exact
 # Batched recurrent BACKWARD: advances all B*nv (request, head) items' reverse-token
 # backward in one launch (flat work list + ping-pong state), mirroring the batched
+# forward decode. 1.3-1.4x over the per-(b,h) loop from ~16 items (slower below 8); all 6 gradients bit-match the
 # single-item kernel.
 from .nki_kda_recurrent_bwd_batched import kda_recurrent_bwd_batched