## Summary
Fixes autotune failures under TileLang 0.1.11 for three kernel families
that were not covered by PR #1619.
### Root cause
TileLang 0.1.11 changed its autotuner: when
`autotuned_kernel_fn(**initial_kwargs)` is called with the tunable JIT
parameters already bound, it records those params in the cache key and —
on a cache hit — treats them as "already tuned," skipping the
benchmarking sweep entirely and returning `config=None`.
PR #1619 added `_call_autotuned_kernel()` to seed required JIT params
(fixing the `TypeError: missing N required positional arguments`), but
this inadvertently triggered the new 0.1.11 skip behavior, causing:
- `SSDStatePassingFwdKernel` / `SSDChunkScanFwdKernel` /
`SSDDecodeKernel`: `Best config: None`
- `SoftmaxKernel` / `LogSumExpKernel`: `TypeError: '<' not supported
between instances of 'NoneType' and 'float'`
### Fix
Pass `do_not_specialize=tunable_params` to TileLang's `autotune()`
decorator. This tells 0.1.11 to exclude those parameter names from the
cache key, so the seeded initial values no longer trigger the skip path
— the full benchmarking sweep runs as expected.
## Changes
- **`tileops/kernels/kernel_base.py`** (`Kernel.autotune()`): extract
tunable param names from the JIT signature and pass as
`do_not_specialize`. Fixes `SSDStatePassingFwdKernel`,
`SSDChunkScanFwdKernel`, `SSDDecodeKernel`, and any future kernel that
uses the base `autotune()`.
- **`tileops/kernels/reduction/softmax.py`**
(`SoftmaxKernel.autotune()`): same pattern for the custom per-tile_n
autotune loop.
- **`tileops/kernels/reduction/logsumexp.py`**
(`LogSumExpKernel.autotune()`): same pattern for the custom per-tile_n
autotune loop.
## Validation
Tested against real TileLang 0.1.11 (`0.1.11+cuda.gitcd37ed5f`) in
`flashmlaenv`:
**Mamba kernels** — all run full sweeps and return valid configs:
- `SSDStatePassingFwdKernel` → `{'block_d': 64, 'threads': 128,
'vectorize': False}`
- `SSDChunkScanFwdKernel` → `{'block_l': 64, 'block_p': 64, 'block_n':
64, 'block_s': 64, 'threads': 128}`
- `SSDDecodeKernel` → `{'block_p': 4, 'block_n': 32, 'threads': 128}`
**Softmax/logsumexp** — `tests/ops/test_softmax.py`: **135 passed**
(previously 3 failed with `TypeError`)
**Autotune binding unit tests** —
`tests/kernels/test_kernel_autotune_binding.py`: **3 passed**
**deepseek_dsa_decode** —
`tests/ops/attention/test_deepseek_dsa_decode.py`: **1 passed**
**Ruff** — all PR #1619 validation files pass cleanly.
## Summary
Fixes#1613.
TileLang 0.1.11 validates/binds autotuned JIT signatures before
candidate configs are applied. Several TileOPs kernels expose tunable
JIT parameters such as `block_m`, `block_n`, and `bdim` as required
arguments, so calling the autotuned wrapper with no kwargs can fail
before autotuning starts.
This seeds the autotuned wrapper call with bindable config values while
preserving the existing autotune search space.
## Changes
- Add shared `Kernel._call_autotuned_kernel()` /
`_autotune_initial_kwargs()` helpers.
- Seed base `Kernel.autotune()` with `default_config`, filtered to the
JIT signature.
- Support common JIT parameter aliases such as `threads_arg` and
`npt_arg`.
- Update custom autotune overrides that previously called autotuned
wrappers with no kwargs:
- `SoftmaxKernel`
- `LogSumExpKernel`
- `SparseMlaKernel`
- `FP8LightingIndexerKernel`
- DeltaNet / GatedDeltaNet sub-kernel autotune paths
- Add unit coverage for required tunable binding, signature filtering,
and alias mapping.
## Validation
- `python3 -m pytest -q tests/kernels/test_kernel_autotune_binding.py`
- `python3 -m ruff check tileops/kernels/kernel_base.py
tileops/kernels/reduction/softmax.py
tileops/kernels/reduction/logsumexp.py
tileops/kernels/attention/deepseek_dsa_decode.py
tileops/kernels/fp8_lighting_indexer.py
tileops/kernels/deltanet/deltanet_fwd.py
tileops/kernels/deltanet/deltanet_bwd.py
tileops/kernels/gated_deltanet/gated_deltanet_fwd.py
tests/kernels/test_kernel_autotune_binding.py`
- `python3 -m compileall -q ...` on changed Python files
- commit hooks / pre-commit checks
Note: local environment has TileLang 0.1.9, so the regression test mocks
the TileLang 0.1.11-style binding requirement instead of running the
exact 0.1.11 autotuner.
## Summary
Two fixes to make `test_ssd_state_passing_fwd_bench` produce stable,
apples-to-apples numbers.
## Change 1 — `benchmarks/ops/bench_mamba.py`
### out_dtype=torch.float32 on Triton baseline
TileOPs always outputs `float32`. Without this fix, the Triton baseline
defaults to `out_dtype=states.dtype` (fp16), writing half as many bytes.
This makes TileOPs appear artificially fast on write-bandwidth-bound
shapes.
### Triton autotuner pre-warm
`@triton.autotune` runs all 6 `BLOCK_SIZE` candidates on the first call.
In a fresh nightly process, this fires inside `bench_kernel`'s first
CUPTI active step, polluting the sum with 5 extra kernel launches. The
result is inflated and non-reproducible.
**Nightly evidence:** N108 showed 10/20 Triton configs with anomalous
values (e.g. 0.1878ms, 0.1985ms = 6 candidates × ~0.03ms each). N109
happened to have a warm cache and showed correct 0.013ms values for the
same configs.
Fix: call `mamba_fwd(); torch.cuda.synchronize()` once before
`bm.profile(mamba_fwd)`.
## Change 2 — `tileops/kernels/kernel_base.py`
### autotune warmup 10→25, rep 10→50
With `rep=10` the autotuner measured each candidate with 10 CUPTI
samples. For 2–20µs kernels this produces ~50% measurement noise, making
the winning config non-deterministic across runs.
**Nightly evidence:** longctx shapes in `SSDStatePassingFwdOp` flipped
between configs across consecutive nightly runs, causing 3–10x latency
swings (e.g. 0.2093ms vs 0.0753ms for `longctx-370m-32k`). With `rep=50`
the same config is selected consistently.
This affects **all kernels using `tune=True`**, not just
ssd_state_passing.
## Testing
- 20/20 `test_ssd_state_passing_fwd_bench` pass
- 5 consecutive runs: 20/20 stable within ±2% (verified on H200 at 1830
MHz)
- Nightly Triton baseline: 19/20 stable (1 remaining case is GPU clock
noise at ~3µs floor)
Cao Ying