overhauled benchmarking system for cuda source

CONTRIBUTING.md

@@ -37,6 +37,73 @@ If any step fails, the commit will be aborted.
 find csrc -name "*.cu" -o -name "*.cuh" | xargs clang-format -i
 ```
 
+## Benchmark Annotations
+
+Kernel benchmarks are auto-generated from annotations in `.cuh` headers. Add `// @bench` annotations above dispatch functions (not `__global__` kernels) to generate benchmarks automatically.
+
+### Quick example
+
+```cpp
+// @bench rmsnorm
+// @title: RMSNorm
+// @configs: (N,C) = [(128,4096),(512,4096),(1024,4096)]
+// @in: inp float[N * C] random(-1,1), weight floatX[C] bf16(0.75,1.25)
+// @out: out float[N * C]
+// @scalar: eps=1e-5f
+// @bandwidth: (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + (size_t)N * C * sizeof(float)
+// @profile: rmsnorm_forward_profiled N=1024 C=4096 grid=N
+inline void rmsnorm_forward(float* out, const float* inp, const floatX* weight,
+                            int N, int C, float eps, cudaStream_t stream = nullptr) {
+```
+
+### Annotation fields
+
+| Field | Required | Description |
+|-------|----------|-------------|
+| `@bench` | Yes | Marks function for codegen. Optional name: `@bench custom_name` |
+| `@configs` | Yes | Dimension tuples: `(M,N) = [(32,32),(64,64)]` |
+| `@in` | Yes | Input buffers: `name type[size] fill`. Comma-separated for multiple |
+| `@out` | Yes | Output buffers: `name type[size]` |
+| `@title` | No | Section title in benchmark output |
+| `@scalar` | No | Scalar args: `eps=1e-5f` |
+| `@bandwidth` | No | Custom bytes expression (default: sum of buffer sizes) |
+| `@flops` | No | FLOPs expression for GFLOPS metric |
+| `@profile` | No | Profiled variant call with config overrides |
+| `@group` | No | Group multiple `@bench` blocks into one output file |
+| `@extra-buf` | No | Buffers not in the function signature |
+| `@setup` | No | Custom C++ setup code (runs after device allocs) |
+| `@pre-setup` | No | Setup code before device allocs (e.g. workspace size calc) |
+| `@pre-iter` | No | Per-iteration reset: `buf.zero()` |
+| `@runs` | No | Iteration count (default: 100) |
+
+### Fill types
+
+- `random(lo,hi)` — fp32 random fill, optional seed: `random(-1,1,43)`
+- `bf16(lo,hi)` — bf16 random fill
+- `zero` — device-side zero (no host allocation)
+- `ones(val)` — constant fill: `ones(1.0/n)`
+- `computed` — filled by `@setup` code
+
+### Multi-line annotations
+
+Long lines can wrap using `//` continuation:
+
+```cpp
+// @configs: (N,C) =
+// [(128,4096),(256,4096),(512,4096),(1024,4096)]
+// @bandwidth: (size_t)N * C * sizeof(float) +
+// (size_t)C * sizeof(floatX)
+```
+
+### Generating and checking
+
+```bash
+make benchgen        # regenerate all benchmark .cu files
+make benchgen-check  # verify generated files are up to date (for CI)
+```
+
+Generated files go to `src/csrc/benchmarks/bench_<name>.cu` and start with `// AUTO-GENERATED by benchgen.py`. Do not edit them manually.
+
 ## Code Style
 
 - 4 space indents

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all build test test-python bench bench-python clean format lint install install-dev force-install
+.PHONY: all build test test-python bench bench-python clean format lint install install-dev force-install benchgen benchgen-check
 
 CUDA_ARCH ?= "80;86;89;90;100"
 BUILD_DIR = build
@@ -55,12 +55,18 @@ test: build
 test-python: install
 	LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6 pytest src/python/tests -v
 
-bench: build
+bench: benchgen build
 	@for b in $(BUILD_DIR)/bench_*; do echo "Running $$b..."; $$b; done
 
 bench-python: install
 	LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6 python src/python/benchmarks/bench_layer.py --all-configs --backward
 
+benchgen:
+	python3 scripts/benchgen.py --force
+
+benchgen-check:
+	python3 scripts/benchgen.py --check
+
 clean:
 	rm -rf $(BUILD_DIR) $(STAMP_DIR)
 

README.md

@@ -15,6 +15,8 @@ Once the image builds the first time, it will be cached and will not require a r
 
 ### Benchmark
 
+Run benchmark suite
+
 ```bash
 # python bench
 modal run runmodal.py --gpu h100 --mode bench --scope python
@@ -26,6 +28,16 @@ modal run runmodal.py --gpu h100 --mode bench --scope native
 modal run runmodal.py --gpu h100 --mode bench --scope all
 ```
 
+Generate benchmark files (this is automatically run in the above)
+
+```bash
+# generate the benchmark files
+make benchgen
+
+# check status of benchmark files
+make benchgen-check
+```
+
 ### Test
 
 ```bash
@@ -108,29 +120,19 @@ in section A.1):
 
 | Batch | Hidden | n  | Forward | Backward |
 |-------|--------|----|---------|----------|
-| 320   | 1280   | 4  | 13.2x   | 11.1x    |
-| 512   | 1920   | 4  | 9.0x    | 7.7x     |
-| 1280  | 2560   | 4  | 5.1x    | 3.6x     |
-| 2560  | 1280   | 4  | 5.0x    | 3.5x     |
-| 128   | 1280   | 8  | 13.6x   | 11.5x    |
-| 256   | 1280   | 8  | 10.3x   | 9.8x     |
-| 32    | 1280   | 32 | 5.8x    | 2.9x     |
-| 64    | 1280   | 32 | 4.7x    | 2.2x     |
-| 128   | 1280   | 32 | 3.5x    | 1.5x     |
+| 320   | 1280   | 4  | 15.20x  | 10.07x   |
+| 512   | 1920   | 4  | 10.52x  | 9.20x    |
+| 1280  | 2560   | 4  | 5.66x   | 4.34x    |
+| 2560  | 1280   | 4  | 5.66x   | 4.21x    |
 
 **Dynamic H Path** (per-batch H values computed via Equations 7-9 from paper):
 
 | Batch | Hidden | n  | Forward | Backward |
 |-------|--------|----|---------|----------|
-| 320   | 1280   | 4  | 6.7x    | 11.0x    |
-| 512   | 1920   | 4  | 6.8x    | 9.0x     |
-| 1280  | 2560   | 4  | 4.6x    | 5.1x     |
-| 2560  | 1280   | 4  | 4.6x    | 5.0x     |
-| 128   | 1280   | 8  | 6.4x    | 11.2x    |
-| 256   | 1280   | 8  | 6.1x    | 10.6x    |
-| 32    | 1280   | 32 | 1.9x    | 3.1x     |
-| 64    | 1280   | 32 | 1.8x    | 2.5x     |
-| 128   | 1280   | 32 | 1.7x    | 1.9x     |
+| 320   | 1280   | 4  | 7.39x   | 3.35x    |
+| 512   | 1920   | 4  | 7.38x   | 3.47x    |
+| 1280  | 2560   | 4  | 5.33x   | 3.07x    |
+| 2560  | 1280   | 4  | 5.21x   | 3.02x    |
 
 
 ## Format