[perf][gemm]: tile short-token FFN weight gradients

benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py

@@ -0,0 +1,82 @@
+import json
+import sys
+
+import torch
+
+from rl_engine.kernels.ops.pytorch.ffn.ffn import qwen3_ffn
+
+
+H, I = 4096, 12288
+WARMUP = 5
+FW_ITERS = 20
+FB_ITERS = 10
+
+
+def timed(fn, warmup, iters):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(iters):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end))
+    values = torch.tensor(samples, dtype=torch.float64)
+    return {
+        "median_ms": float(values.median()),
+        "mean_ms": float(values.mean()),
+        "min_ms": float(values.min()),
+        "max_ms": float(values.max()),
+        "samples_ms": samples,
+    }
+
+
+def main():
+    torch.cuda.set_device(0)
+    torch.manual_seed(2026)
+    result = {
+        "device": torch.cuda.get_device_name(0),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "shape": {"hidden": H, "intermediate": I, "dtype": "bfloat16"},
+        "warmup": WARMUP,
+        "forward_iters": FW_ITERS,
+        "forward_backward_iters": FB_ITERS,
+        "rows": [],
+    }
+    for tokens in map(int, sys.argv[1:]):
+        x = torch.randn(tokens, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        gate = torch.randn(I, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        up = torch.randn(I, H, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        down = torch.randn(H, I, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        dout = torch.randn(tokens, H, device="cuda", dtype=torch.bfloat16)
+        for deterministic in (True, False):
+            def forward():
+                return qwen3_ffn(x, gate, up, down, deterministic=deterministic)
+
+            def forward_backward():
+                y = forward()
+                torch.autograd.grad(y, (x, gate, up, down), dout)
+
+            row = {
+                "tokens": tokens,
+                "mode": "det" if deterministic else "prod",
+                "forward": timed(forward, WARMUP, FW_ITERS),
+                "forward_backward": timed(forward_backward, WARMUP, FB_ITERS),
+            }
+            result["rows"].append(row)
+            print(
+                f"tokens={tokens} mode={row['mode']} "
+                f"forward_median={row['forward']['median_ms']:.4f} "
+                f"fwd_bwd_median={row['forward_backward']['median_ms']:.4f}",
+                flush=True,
+            )
+    print(json.dumps(result))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/results/qwen_ffn_h100_trace/kernel_summary.md

@@ -0,0 +1,28 @@
+# Qwen FFN H100 kernel profile
+
+These charts are aggregated from the Nsight Systems `cuda_gpu_kern_sum` report
+for the same `tokens=1,8`, `hidden=4096`, `intermediate=12288` profiling run.
+The benchmark correctness checks passed in both worktrees.
+
+| Profile | Kernel | Total time | Launches | Average |
+|---|---|---:|---:|---:|
+| Before | `det_gemm_sm90_kernel` | 392.262584 ms | 643 | 610.051 us |
+| Before | `det_gemm_naive<bf16,true>` | 157.657338 ms | 90 | 1,751.748 us |
+| Before | `det_gemm_naive<bf16,false>` | 116.315558 ms | 84 | 1,384.709 us |
+| After | `det_gemm_sm90_kernel` | 391.452104 ms | 643 | 608.790 us |
+| After | `det_gemm_naive<bf16,false>` | 116.338858 ms | 84 | 1,384.986 us |
+| After | `det_gemm_db_small_k<bf16>` | 12.006408 ms | 90 | 133.405 us |
+
+The optimized kernel replaces the 90 `det_gemm_naive<bf16,true>` launches:
+
+- 157.657338 ms → 12.006408 ms;
+- 13.13× kernel-time speedup;
+- 92.38% reduction for the replaced kernel;
+- 145.651 ms saved in the selected deterministic GEMM kernel aggregate.
+
+The unchanged SM90 and `det_gemm_naive<bf16,false>` rows provide a useful
+control: the observed reduction is localized to the intended short-token dW
+path rather than a general profiling artifact.
+
+Raw `.nsys-rep` files remain on the H100 profiling node because they are large;
+the SVGs and this reproducible kernel summary are the review artifacts.