[Kernel] Add opt-in autotuning for Softmax (#1022)

* [Kernel] Add opt-in autotuning for Softmax

Softmax hard-coded BLOCK_THREADS=256 for every shape and dtype. A forced search on
gfx950 selects 512 for wide bf16/f16 rows and smaller blocks for narrow ones, so 256
is a compatibility default rather than a universal optimum.

Make Softmax the second direct-JIT autotune adopter after RMSNorm. The kernel factory
takes BLOCK_THREADS as a static input, softmax_direct exposes it as a Constexpr, and
softmax_autotuned serves the searched winner, then a matching artifact, then the 256
default. Ordinary calls never benchmark; FLYDSL_AUTOTUNE=1 is the only search path.

Two narrow additions to the shared autotuner were required. validate_hook runs a
numerical check once per candidate, outside the timed repetitions but under the same
stream, compile hints and reset/restore policy, so a candidate that launches cleanly
but computes the wrong answer cannot win. An all-rejected search now chains the last
failure, which keeps a numerical rejection distinguishable from a compile failure.

The candidate gate compares against the row's own scale. The previous absolute
tolerance was sized for O(1) values while softmax elements are O(1/N), so at N=8192
the bf16 bound exceeded the signal by two orders of magnitude and an all-zero output
would have passed. tests/kernels/test_softmax.py adopts the same criterion so the
default is held to the standard it imposes on candidates.

tuning_schema is a declared key axis, not kernel input: the scratch winner cache does
not fingerprint kernel source, so bumping it is the only way to invalidate stale
winners after a change that can move the result.

* test(softmax): isolate autotune cache state

* fix(autotune): reject candidates with unwritten outputs

Wrap validation around an untimed candidate launch so output poisoning catches no-op or partial-store kernels without affecting benchmark rankings.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(softmax): address autotune correctness gaps

Synchronize reused reduction scratch before the next block reduction, reject streams from a different device, and pass fully bound call arguments to candidate validation.

* fix(ci): pin compatible nanobind for LLVM builds

---------

Co-authored-by: zhimding <zhimding@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

docs/autotune_guide.md

@@ -46,6 +46,57 @@ Artifacts intentionally do not include a compiler or kernel-source fingerprint.
 Treat them as reviewed deployment inputs, and retune after a compiler, kernel,
 compile-hint, or search-space change that can affect the winner.
 
+The scratch winner cache has the same blind spot: it fingerprints the device,
+toolchain, environment and compile hints, but not the adopter's kernel source or
+search space. An adopter that needs stale scratch winners invalidated should
+declare an integer schema parameter on its entry point and list it in `key`, then
+bump it with any change that can move the winner. Softmax does this with
+`tuning_schema`.
+
+## Candidate correctness gate
+
+`validate_hook(sig_args)` returns a context manager around one untimed candidate
+launch. Code before `yield` can poison outputs; code after `yield` validates the
+result, so skipped or partial stores cannot inherit a previous candidate's data.
+The launch uses the same stream, compile hints, reset/restore policy and arguments
+as timing, but validation work never affects ranking. `sig_args` maps every kernel
+parameter name to its value, including positional tensors. Softmax uses this hook
+to fill its output with NaNs before launch and check numerics afterward.
+
+Raising from the hook rejects that candidate. If every candidate is rejected the
+search raises `RuntimeError("All autotune configs failed")` with the last failure
+chained, so a numerical rejection stays distinguishable from a compile failure.
+Use it wherever a candidate could launch successfully and still compute the wrong
+answer, and hold every candidate to the same tolerance as the default.
+
+## Device timing contract
+
+The shared `do_bench` timer queues a GPU-side backlog before batched event
+windows. This is required for sub-100 µs kernels: a fresh event pair on an empty
+stream can time the host enqueue gap instead of the kernel. Each window averages
+several launches, and the reported value is the median across windows. The
+callable must enqueue asynchronous work on the current stream and must not
+synchronize internally.
+
+For Softmax results within 2% of the measured minimum, selection prefers the
+compatibility default, then a config without an explicit occupancy override,
+then the candidate packing more rows per block. This prevents event granularity
+from turning equivalent 6--10 µs candidates into unstable deployment artifacts;
+an improvement outside the band still wins normally. Softmax uses 10 warmup and
+100 measured launches, split into five backlogged event windows; the larger
+sample stabilized bandwidth-scale rows that moved by more than the tie band with
+the generic 25-launch default.
+
+## Adopters
+
+| Kernel | Module | `artifact_name` | Tuned axes |
+|---|---|---|---|
+| RMSNorm | `kernels/norm/rmsnorm_autotune.py` | `rmsnorm` | `BLOCK_THREADS`, `waves_per_eu` |
+| Softmax forward | `kernels/norm/softmax_autotune.py` | `softmax_fwd` | full-row threads, `waves_per_eu`, threads/rows per block for short rows |
+
+Softmax backward is not an adopter yet; its existing kernel and dispatch are
+unchanged by `softmax_fwd` artifacts.
+
 ## Failure behavior
 
 FlyDSL ignores missing, unreadable, mismatched, or structurally invalid

docs/prebuilt_kernels_guide.md

@@ -8,7 +8,7 @@ This guide covers the available FlyDSL kernels — normalization, softmax, GEMM,
 |---|---|---|---|---|
 | **LayerNorm** | `build_layernorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Two-pass vectorized normalization |
 | **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16; optional fp32 weight | LDS-cached 3-pass pipeline |
-| **Softmax** | `build_softmax_module(M, N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Online softmax, adaptive block size |
+| **Softmax** | `build_softmax_module(M, N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Register-buffered softmax, opt-in autotuning |
 | **Softmax backward** | `build_softmax_bwd_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | fp32 dot reduction, native-dtype register buffering |
 | **GEMM** | `compile_preshuffle_gemm(...)` | `@flyc.kernel` | fp8, int8, fp16, bf16 | Preshuffle B, ping-pong LDS, MFMA 16x16 |
 | **FlashAttention** | `build_flash_attn_func_module(...)` | `@flyc.kernel` | bf16, f16 (any arch); fp8 e4m3fn (gfx950, D=128, dense) | Dual-wave SWP fwd, GQA/MQA, causal, descale ABI |
@@ -117,9 +117,63 @@ executor = build_softmax_module(M=32768, N=8192, dtype_str="bf16")
 **Configuration:**
 | Parameter | Value | Description |
 |---|---|---|
-| `BLOCK_SIZE` | `min(256, next_power_of_2(N))`, min 32 | Adaptive block size |
-| `VEC_WIDTH` | 8 | Vector load/store width |
-| `WARP_SIZE` | 64 | AMD wavefront size |
+| `BLOCK_THREADS` | 256 by default; 64/128/256/512 for full-row candidates | Total threads per block |
+| `THREADS_PER_ROW` | Defaults to `BLOCK_THREADS`; 8/16/32/64 for short-row candidates | Reduction subgroup assigned to one row |
+| `ROWS_PER_BLOCK` | 1 by default; derived from `BLOCK_THREADS / THREADS_PER_ROW` | Independent rows packed into one block |
+| `vec_width` | `128 // elem_bits` (8 for f16/bf16, 4 for f32) | Derived from the 128-bit transaction contract |
+| `WARP_SIZE` | 64 on CDNA, 32 on RDNA | Wavefront size, resolved from the target arch |
+
+`THREADS_PER_ROW` also selects the data-movement path: with
+`tile_cols = THREADS_PER_ROW * vec_width`, a row takes the vectorized fast path when
+`N % tile_cols == 0` and the scalar generic path otherwise.
+
+**Opt-in autotuning** (`kernels/norm/softmax_autotune.py`):
+```python
+from kernels.norm.softmax_autotune import softmax_autotuned
+
+softmax_autotuned(x, y)              # serves the tuned or default config, never searches
+```
+Ordinary calls follow the searched-winner cache → offline artifact → compatibility default
+(`BLOCK_THREADS=256`) ordering and never benchmark. `FLYDSL_AUTOTUNE=1` forces a search over
+a bounded, shape-aware space:
+
+- full-row `BLOCK_THREADS ∈ {64,128,256,512}` ×
+  `waves_per_eu ∈ {none,1,2,4}`;
+- Quack-style short-row packing that decouples `THREADS_PER_ROW` from total
+  block threads and processes several rows per block.
+
+The search-space rationale was checked against AITER
+`536118aaf94047b0b559e0730749352659419b34`, SGLang
+`955704544c60e920672aa434cefa2ce78c0ceb4c`, and Tri Dao's Quack
+`60d88082272a256fa9b3b2ab631c82cfa78337c6`. Quack's portable ideas are the
+row-width-dependent reduction subgroup, a separate 128/256-thread CTA size,
+multiple rows per CTA, and an online/non-online algorithm choice; its
+multi-CTA cluster reduction is NVIDIA-specific. AITER's standalone Triton
+kernel is a fixed two-pass online/reload implementation (`.cg`, eight warps,
+two stages, `waves_per_eu=2`), not an autotuned space. The pinned SGLang tree
+has attention-local and top-k softmax implementations but no directly comparable
+standalone row-wise kernel; attention tile/stage choices are therefore not
+imported here.
+
+Every candidate is numerically validated before ranking and uses the shared
+GPU-backlog and batched-event timer. Input cache policy is deliberately not a
+search axis: on gfx950, non-temporal loads changed rank between repeated use of
+one address and rotation across fresh addresses. Cache residency is absent from
+the shape-only artifact identity, so persisting either result would encode an
+unstated workload assumption. The tested three-pass reload algorithm is also
+excluded because it lost to the register-buffered compatibility path; a future
+algorithm axis should implement a true online pair reduction before entering
+the default search.
+
+Within a 2% timing tie, the selector favors the compatibility default, then no
+explicit `waves_per_eu`, then more rows per block. Larger measured improvements
+still win; the tie rule only avoids persisting noise-level differences between
+6–10 µs candidates. Softmax uses 10 warmup and 100 measured launches, divided
+into five GPU-backlogged event windows, so bandwidth-scale candidates are also
+ranked from a stable sample.
+
+Artifacts use the name `softmax_fwd` and cover forward only. Softmax backward
+has no autotune adopter in this change. See [`autotune_guide.md`](autotune_guide.md).
 
 **Algorithm (6 stages):**
 1. **Load data**: Vectorized global loads into register buffer with validity masks
@@ -130,11 +184,12 @@ executor = build_softmax_module(M=32768, N=8192, dtype_str="bf16")
 6. **Normalize + store**: Divide by sum, convert to output dtype, vectorized store
 
 **Kernel signature:**
-```
-GPU_MODULE_NAME = f"softmax_{dtype_str}"
+```python
+build_softmax_module(M, N, dtype_str="f32", BLOCK_THREADS=256)  # M is vestigial
+launch_softmax(A, C, m_in, stream=...)                          # returned launcher
 
-@kernel
-softmax_kernel(self, A, C, m_in)
+# Direct-JIT entry point used by the autotuner
+softmax_direct(A, C, m_in, N, dtype_str, BLOCK_THREADS, tuning_schema, stream=...)
 ```
 
 ### 2.2 Softmax backward (`kernels/norm/softmax_bwd_kernel.py`)
@@ -421,6 +476,7 @@ What operation do you need?
 | `kernels/norm/rmsnorm_kernel.py` | RMSNorm (layout API) |
 | `kernels/norm/softmax_kernel.py` | Softmax (layout API) |
 | `kernels/norm/softmax_bwd_kernel.py` | Softmax backward (layout API) |
+| `kernels/norm/softmax_autotune.py` | Softmax opt-in autotune adopter |
 | `kernels/attention/fused_rope_cache_kernel.py` | Fused RoPE + KV cache |
 | `kernels/comm/custom_all_reduce.py` | Multi-GPU all-reduce |
 | `kernels/gemm/rdna_f16_gemm.py` | RDNA FP16 GEMM |
@@ -447,6 +503,7 @@ What operation do you need?
 | `tests/kernels/test_rmsnorm.py` | RMSNorm |
 | `tests/kernels/test_softmax.py` | Softmax |
 | `tests/kernels/test_softmax_bwd.py` | Softmax backward |
+| `tests/kernels/test_softmax_autotune.py` | Softmax autotune selection and candidate correctness |
 | `tests/kernels/test_fused_rope_cache.py` | Fused RoPE + KV cache |
 | `tests/kernels/test_allreduce.py` | Multi-GPU all-reduce |
 | `tests/kernels/test_rdna_gemm.py` | RDNA GEMM |