Uploaded using `kernel-builder`.

README.md

@@ -17,20 +17,21 @@ tags:
 
 Fused **CuteDSL** kernels for the loss functions that dominate post-training
 workloads: PPO-family policy losses (BNPO, GRPO) and reverse-KL
-self-distillation. Each kernel ships a **single-launch fused forward +
-backward** path that returns `(loss, grad_logprobs)` directly — no
-`torch.autograd.Function` wrapper, no extra `grad_output * dpolicy` backward
+self-distillation. 
+
+Each kernel ships a **single-launch fused forward +
+backward** path that returns `(loss, grad_logprobs)` directly. No `torch.autograd.Function` wrapper, no extra `grad_output * dpolicy` backward
 kernel, and no host-side syncs in the hot path.
 
 Background and benchmarks: see the
 [release post](https://geometric.so/blog/2026/05/08/hf-kernel-hub).
 
 - **Backend**: CUDA (NVIDIA CUTLASS DSL).
-- **Min GPU**: SM80 (Ampere) — required by `nvidia-cutlass-dsl`. Tested on A100 (SM80) and H100 (SM90). Works on SM86 (RTX 3090, A40), SM89 (RTX 4090, L40S), SM90a (H100 SXM), and SM100 (Blackwell B200/GB200).
+- **Min GPU**: SM80 (Ampere) - required by `nvidia-cutlass-dsl`. Tested on H100 (SM90). Should work on SM80 (Ampere), SM86 (RTX 3090, A40), SM89 (RTX 4090, L40S), SM90a (H100 SXM), and SM100 (Blackwell B200/GB200).
 - **Min CUDA**: 12.8.
 - **Dtypes**: `float32`, `float16`, `bfloat16`.
 - **Dynamic shapes**: a single compile handles arbitrary batch size and
-  sequence length — no recompiles when shapes change between calls (common
+  sequence length, no recompiles when shapes change between calls (common
   in post-training rollouts).
 
 ## Kernels
@@ -45,15 +46,15 @@ Background and benchmarks: see the
 
 Each kernel family exposes three entry points with the same underlying CuteDSL kernel:
 
-- **`<name>(...)`** — fused fwd+bwd, returns `(loss, grad)` from one `@cute.jit`
+- **`<name>(...)`** - fused fwd+bwd, returns `(loss, grad)` from one `@cute.jit`
   dispatch. Lowest-overhead path; the caller chains the gradient into the upstream
   model with `policy_logprobs.backward(grad)`. Use this in custom training loops
   where you control gradient flow.
-- **`<name>_autograd(...)`** — same kernel, registered via
+- **`<name>_autograd(...)`** - same kernel, registered via
   `torch.library.custom_op` + `register_autograd`. `loss.backward()` works
   and composes with `torch.compile(fullgraph=True)`. There is a noticeable
   per-call dispatcher overhead vs. the direct path.
-- **`<name>_fwd(...)`** — forward-only, returns scalar `loss` and skips
+- **`<name>_fwd(...)`** - forward-only, returns scalar `loss` and skips
   the gradient buffer entirely. Use for inference / validation /
   reward-model scoring.
 
@@ -85,7 +86,7 @@ per_token  = −advantages · min(ratio, clipped)
 kl         = exp(ref_logprobs − policy_logprobs) − (ref_logprobs − policy_logprobs) − 1
 ```
 
-The global denominator is computed entirely on-GPU via cross-CTA atomics —
+The global denominator is computed entirely on-GPU via cross-CTA atomics -
 no host-side `mask.sum()` sync. When `beta=0` the KL branch is dead-coded
 at compile time.
 
@@ -110,7 +111,7 @@ ref_logprobs        = torch.randn(bs, seq_len, dtype=torch.bfloat16, device=devi
 advantages          = torch.randn(bs, dtype=torch.bfloat16, device=device)
 completions_mask    = (torch.rand(bs, seq_len, device=device) > 0.2).to(torch.int8)
 
-# 1) Direct (loss, grad) — lowest overhead training path
+# 1) Direct (loss, grad) - lowest overhead training path
 loss, grad = km.bnpo_loss(
     policy_logprobs, old_policy_logprobs, ref_logprobs,
     advantages, completions_mask,
@@ -118,7 +119,7 @@ loss, grad = km.bnpo_loss(
 )
 policy_logprobs.backward(grad)
 
-# 2) Autograd-aware — works with loss.backward() and torch.compile
+# 2) Autograd-aware - works with loss.backward() and torch.compile
 loss = km.bnpo_loss_autograd(
     policy_logprobs.requires_grad_(),
     old_policy_logprobs, ref_logprobs,
@@ -127,7 +128,7 @@ loss = km.bnpo_loss_autograd(
 )
 loss.backward()
 
-# 3) Forward-only — inference / reward scoring, no gradient buffer
+# 3) Forward-only - inference / reward scoring, no gradient buffer
 loss = km.bnpo_loss_fwd(
     policy_logprobs, old_policy_logprobs, ref_logprobs,
     advantages, completions_mask,
@@ -140,7 +141,7 @@ loss = km.bnpo_loss_fwd(
 ## GRPO Loss
 
 **Group Relative Policy Optimization** implements TRL's default
-**per-response normalization** variant — each response is normalized by its
+**per-response normalization** variant - each response is normalized by its
 own valid-token count before averaging across the batch:
 
 ```
@@ -150,12 +151,12 @@ loss = mean_r( ((per_token_loss + β·kl) · mask).sum(-1) / max(mask.sum(-1), 1
 `per_token_loss` and `kl` are the same clipped-ratio and KL expressions as BNPO.
 `completions_mask` is **required** because the per-response denominator is
 mask-derived. The kernel uses one CTA per row so the per-row mask sum is
-reduced inside the block — no cross-CTA atomics on the scaling pass.
+reduced inside the block - no cross-CTA atomics on the scaling pass.
 
 **Inputs**:
 - `policy_logprobs`, `old_policy_logprobs`, `ref_logprobs`: `(bs, seq_len)`, fp32/fp16/bf16
 - `advantages`: `(bs,)`
-- `completions_mask`: `(bs, seq_len)`, bool or int8 — **required**
+- `completions_mask`: `(bs, seq_len)`, bool or int8 - **required**
 
 **Returns**: `(loss, grad_policy_logprobs)` from `grpo_loss`; scalar `loss` from `grpo_loss_fwd`.
 
@@ -173,7 +174,7 @@ ref_logprobs        = torch.randn(bs, seq_len, dtype=torch.bfloat16, device=devi
 advantages          = torch.randn(bs, dtype=torch.bfloat16, device=device)
 completions_mask    = (torch.rand(bs, seq_len, device=device) > 0.2).to(torch.int8)
 
-# 1) Direct (loss, grad) — lowest overhead training path
+# 1) Direct (loss, grad) - lowest overhead training path
 loss, grad = km.grpo_loss(
     policy_logprobs, old_policy_logprobs, ref_logprobs,
     advantages, completions_mask,
@@ -181,7 +182,7 @@ loss, grad = km.grpo_loss(
 )
 policy_logprobs.backward(grad)
 
-# 2) Autograd-aware — works with loss.backward() and torch.compile
+# 2) Autograd-aware - works with loss.backward() and torch.compile
 loss = km.grpo_loss_autograd(
     policy_logprobs.requires_grad_(),
     old_policy_logprobs, ref_logprobs,
@@ -190,7 +191,7 @@ loss = km.grpo_loss_autograd(
 )
 loss.backward()
 
-# 3) Forward-only — inference / reward scoring, no gradient buffer
+# 3) Forward-only - inference / reward scoring, no gradient buffer
 loss = km.grpo_loss_fwd(
     policy_logprobs, old_policy_logprobs, ref_logprobs,
     advantages, completions_mask,
@@ -222,7 +223,7 @@ grad_student_v = scale · p_v · (log p_v − log q_v − kl_per_row)
 where `scale = mask[r] · inv_n_valid`.
 
 **Inputs**:
-- `student_logits`, `teacher_logits`: `(*, V)` — arbitrary leading dims (typically `(bs, seq_len, vocab)`); both must share shape and dtype
+- `student_logits`, `teacher_logits`: `(*, V)` - arbitrary leading dims (typically `(bs, seq_len, vocab)`); both must share shape and dtype
 - `completions_mask`: shape matching `student_logits.shape[:-1]`
 
 > ⚠️ **Fully-masked batches**: `inv_n_valid = 1 / mask.sum()` is not clamped, so a batch where every token is masked produces inf/NaN. Guard upstream if that case is reachable.
@@ -242,17 +243,17 @@ student_logits  = torch.randn(bs, seq_len, vocab, dtype=torch.bfloat16, device=d
 teacher_logits  = torch.randn(bs, seq_len, vocab, dtype=torch.bfloat16, device=device)
 completions_mask = (torch.rand(bs, seq_len, device=device) > 0.2)
 
-# 1) Direct (loss, grad) — lowest overhead training path
+# 1) Direct (loss, grad) - lowest overhead training path
 loss, grad = km.reverse_kl(student_logits, teacher_logits, completions_mask)
 student_logits.backward(grad)
 
-# 2) Autograd-aware — works with loss.backward() and torch.compile
+# 2) Autograd-aware - works with loss.backward() and torch.compile
 loss = km.reverse_kl_autograd(
     student_logits.requires_grad_(), teacher_logits, completions_mask
 )
 loss.backward()
 
-# 3) Forward-only — inference / KL monitoring, no gradient buffer
+# 3) Forward-only - inference / KL monitoring, no gradient buffer
 loss = km.reverse_kl_fwd(student_logits, teacher_logits, completions_mask)
 ```
 
@@ -260,37 +261,22 @@ loss = km.reverse_kl_fwd(student_logits, teacher_logits, completions_mask)
 
 ## Performance
 
-Numbers below are geometric-mean speedups from our in-house benchmark
-(`triton.testing.do_bench`, fresh subprocess per shape). Baselines are
-**eager PyTorch** and **`torch.compile(mode="max-autotune-no-cudagraphs",
-fullgraph=True)`** with `torch._dynamo.config.trace_autograd_ops = True` so
-the compiled baseline is a real Inductor-fused fwd+bwd graph.
+All numbers are geometric-mean speedups over H100 SXM (SM90a). Full methodology
+and per-shape plots in the [release post](https://geometric.so/blog/2026/05/08/hf-kernel-hub).
 
-| Kernel | β | vs eager | vs `torch.compile` |
-| --- | --- | --- | --- |
-| `bnpo_loss_fwd` | 0   | 1.6× | 1.3× |
-| `bnpo_loss`     | 0   | 1.5× | 1.2× |
-| `bnpo_loss_fwd` | ≠ 0 | 1.4× | 1.1× |
-| `bnpo_loss`     | ≠ 0 | 1.3× | 1.0× |
-| `grpo_loss_fwd` | ≠ 0 | 1.5× | 1.2× |
-| `grpo_loss`     | ≠ 0 | 1.4× | 1.1× |
-| `reverse_kl_fwd`|     | 1.3× | 1.1× |
-| `reverse_kl`    |     | 1.2× | 1.0× |
-
-Profiled on H100 SXM (SM90a). BNPO and GRPO benchmarked separately for `β = 0`
-(KL term dead-coded at compile time) and `β ≠ 0`. Shapes:
-
-- **BNPO / GRPO**: `(16, 1024)`, `(32, 2048)`, `(64, 4096)`, `(128, 8192)`,
-  `(128, 8193)` — the last entry exercises the predicated tail-tile path.
-- **Reverse KL** (vocab = 248320, matching Qwen3.5): `(1, 64)`, `(2, 128)`,
-  `(4, 256)`, `(8, 512)`, `(16, 1024)`, `(8, 1981)`.
-
-Reproduce locally:
-
-```bash
-make bench-kernel KERNEL=grpo_loss   # or bnpo_loss, reverse_kl
-```
+### `kernels` CLI benchmark
+
+Timed with `time.perf_counter` + `cuda.synchronize()`, single iteration per
+shape, mean over 100 iterations. Baseline runs once after warmup.
 
+| Kernel | vs eager | vs `torch.compile` |
+| --- | --- | --- |
+| `grpo_loss_fwd` | 5.60×  | 2.52× |
+| `grpo_loss`     | 19.92× | 2.28× |
+| `bnpo_loss_fwd` | 5.58×  | 2.54× |
+| `bnpo_loss`     | 17.46× | 2.15× |
+| `reverse_kl_fwd`| 6.95×  | 2.44× |
+| `reverse_kl`    | 7.06×  | 2.59× |
 ---
 
 ## Benchmark animations
@@ -299,40 +285,40 @@ make bench-kernel KERNEL=grpo_loss   # or bnpo_loss, reverse_kl
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/bnpo_loss_eager/bnpo_loss_eager_dark_animation.svg">
-  <img src="benchmark_results/bnpo_loss_eager/bnpo_loss_eager_light_animation.svg" alt="BNPO loss latency vs eager PyTorch">
+  <img width="90%" src="benchmark_results/bnpo_loss_eager/bnpo_loss_eager_light_animation.svg" alt="BNPO loss latency vs eager PyTorch">
 </picture>
 
 ### BNPO Loss vs torch.compile
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/bnpo_loss_compiled/bnpo_loss_compiled_dark_animation.svg">
-  <img src="benchmark_results/bnpo_loss_compiled/bnpo_loss_compiled_light_animation.svg" alt="BNPO loss latency vs torch.compile">
+  <img width="90%" src="benchmark_results/bnpo_loss_compiled/bnpo_loss_compiled_light_animation.svg" alt="BNPO loss latency vs torch.compile">
 </picture>
 
 ### GRPO Loss vs eager PyTorch
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/grpo_loss_eager/grpo_loss_eager_dark_animation.svg">
-  <img src="benchmark_results/grpo_loss_eager/grpo_loss_eager_light_animation.svg" alt="GRPO loss latency vs eager PyTorch">
+  <img width="90%" src="benchmark_results/grpo_loss_eager/grpo_loss_eager_light_animation.svg" alt="GRPO loss latency vs eager PyTorch">
 </picture>
 
 ### GRPO Loss vs torch.compile
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/grpo_loss_compiled/grpo_loss_compiled_dark_animation.svg">
-  <img src="benchmark_results/grpo_loss_compiled/grpo_loss_compiled_light_animation.svg" alt="GRPO loss latency vs torch.compile">
+  <img width="90%" src="benchmark_results/grpo_loss_compiled/grpo_loss_compiled_light_animation.svg" alt="GRPO loss latency vs torch.compile">
 </picture>
 
 ### Reverse KL vs eager PyTorch
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/reverse_kl_eager/reverse_kl_eager_dark_animation.svg">
-  <img src="benchmark_results/reverse_kl_eager/reverse_kl_eager_light_animation.svg" alt="Reverse KL latency vs eager PyTorch">
+  <img width="90%" src="benchmark_results/reverse_kl_eager/reverse_kl_eager_light_animation.svg" alt="Reverse KL latency vs eager PyTorch">
 </picture>
 
 ### Reverse KL vs torch.compile
 
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="benchmark_results/reverse_kl_compiled/reverse_kl_compiled_dark_animation.svg">
-  <img src="benchmark_results/reverse_kl_compiled/reverse_kl_compiled_light_animation.svg" alt="Reverse KL latency vs torch.compile">
+  <img width="90%" src="benchmark_results/reverse_kl_compiled/reverse_kl_compiled_light_animation.svg" alt="Reverse KL latency vs torch.compile">
 </picture>