perf(cake_ssd_combined): preprocess, exact-scan and chunk-parallel prefill improvements (SM100 / SM103)

CLAUDE.md

@@ -725,6 +725,7 @@ Used by `flashinfer.trace` / `fi_trace`.
 | `FLASHINFER_CUDNN_PREFILL_SHAPE_OVERRIDE` | `1` | `flashinfer/cudnn/prefill.py` (ragged token-indptr path) | `0` makes the cuDNN ragged prefill graph declare the exact `(batch, max_seq_len)` per call instead of building one override-enabled graph per length class (declared max_len 128 or 65536+, batch 4096+) and passing the real shapes at execute time via cuDNN's shape override. With matching FE 1.31+ Python/native support, BF16 SM100 MLA prefixes (D192/V128, equal 4–64 heads, B1–4, Q64–1024, KV2K–32K at least 4x Q) use bounded power-of-two batch/Q/KV classes and a packed-Q capacity instead; this preserves reuse within each class while exposing useful plan-time occupancy and workspace bounds. Crossing a class builds one cached plan. Override exists because cuDNN finalizes one execution plan per declared shape (55-70 ms, ~1 s with a kernel compile) and serving changes the shape every step; needs cuDNN 9.22+ and cudnn-frontend 1.29+, otherwise the exact declaration is used. Disable only to bisect a suspected override regression. |
 | `FLASHINFER_ALLOW_CUDNN_PLAN_BUILD_IN_CAPTURE` | `0` | `flashinfer/gemm/gemm_base.py` (read once at import) | `1` downgrades the refusal to build a cuDNN execution plan during CUDA graph capture from a `CudnnPlanBuildInCaptureError` to a warning. Set it only if you have verified your cuDNN version tolerates plan build under capture; the supported alternative is to run each shape once eagerly outside the capture region, which populates the per-shape plan cache the captured call then reuses. |
 | `FLASHINFER_DECODE_AUTO_CUDNN` | unset | `flashinfer/decode.py` (`_auto_decode_prefers_cudnn`) | `BatchDecodeWithPagedKVCacheWrapper(backend="auto")`: unset applies the measured rule (SM100/SM103 with cudnn-frontend 1.30+, fp16/bf16 head_dim 128, GQA group dividing 128, `2 <= q_len_per_req <= 4` with `q_len_per_req * group` in [32, 128] rows per CTA and `batch_size * num_kv_heads >= 64` CTAs, no RoPE / soft-cap / window resolve to `cudnn`, whose default SM100 engine runs those multi-token rows on its FROST decode tile at 0.35-0.95x fa2; single-token decode stays on fa2 except head_dim-256 GQA models with group 4 / 8 / 16 and `batch_size * num_kv_heads` in [96, 256], where the d256 decode tile measures at 0.5-0.95x fa2, from 64 under `use_cuda_graph=True` with a caller-owned table on cudnn-frontend 1.31+, which the wrapper tells the graph is replayed (`is_cuda_graph_replay_expected`) so the tile splits KV; multi-token rows without tensor cores, which have no fa2 kernel, take cudnn whenever its decode path can run them); `0` never resolves to cudnn; `1` applies the rule on an older frontend too (benchmarking; the backend engine is ~8x slower there). Under CUDA graphs `auto` takes cudnn only with a caller-owned `block_tables` (the auto-built table cannot grow once captured) or for multi-token rows without tensor cores, and the resolution is frozen after the first plan. A caller-owned table narrower than 2048 tokens keeps fa2 on cudnn-frontend 1.30 (its placement ranks the backend's prefill-class engine first below that; 1.31+ lifts the limit). `auto` takes cudnn only where the frontend's FROST engines can run (its own CuTe-DSL check: `nvidia-cutlass-dsl` >= 4.7.0, above the 4.6.2 the dependency metadata admits; the `1` override and the no-fa2 route are not gated), and callers passing `fixed_split_size` / `disable_split_kv` (fa2's batch-invariance controls) keep fa2 wherever fa2 can serve the rows. `wrapper.resolved_backend` reports the choice. |
+| `FLASHINFER_CAKE_SSD_CHUNK_PARALLEL` | `auto` | `flashinfer/mamba/cake_ssd_combined.py` (`SSDCombined(backend="cake")`, read on every call) | Chooses between the two exact Mamba2 SSD prefill program families of the Cake backend: `auto` applies the calibrated rule (`chunk_parallel_selected`: the chunk-parallel family only when the call has fewer `(sequence, head)` work items than SMs, its per-tile S/H workspace fits in 256 MiB and its predicted main-kernel time beats the exact scan's by the selection margin), `always` forces the chunk-parallel family, `never` forces the exact scan. Both families run the same exact-scan arithmetic. Any other value raises `ValueError`. |
 | `FLASHINFER_NVFP4_4OVER6` | unset | `flashinfer/quantization/nvfp4_quantization_utils.py`, `csrc/nv_internal/cpp/common/envUtils.cpp` | **Legacy.** `1` enables the NVFP4 "4over6" block-scale search. Standard NVFP4 gives each 16-element block the single E4M3 scale implied by its absolute max (`amax / 6`, because 6 is the E2M1 max). 4over6 additionally tries the 1.5x tighter `amax / 4`, quantizes the block both ways, dequantizes both, and keeps whichever reconstructs the block with less error — a little more quantizer work for lower quantization error. Requires fp16/bf16 input and E4M3 (not UE8M0) block scales. Consulted **only** when the public `nvfp4_4over6=` parameter is omitted (its default); a `FutureWarning` is emitted only when this variable actually enables 4over6 (unset or `0` stays silent). An explicit `None` (off) or `NVFP44Over6Config(...)` ignores this variable and the three below, with no per-field merge. Prefer the parameter. |
 | `FLASHINFER_NVFP4_4OVER6_ERR_MODE` | `MAE` | `flashinfer/quantization/nvfp4_quantization_utils.py`, `csrc/nv_internal/cpp/common/envUtils.cpp` | **Legacy.** Metric used to choose between the two 4over6 candidates: `MAE` (sum of absolute errors) or `MSE` (sum of squared errors). The `amax / 4` candidate buys resolution for the bulk of the block by clipping its largest magnitudes, and MSE penalizes those few large residuals much harder than MAE does, so MSE selects it less often — prefer MSE when a handful of outlier activations dominate downstream accuracy. Read only when `FLASHINFER_NVFP4_4OVER6=1` and `nvfp4_4over6=` is omitted. Parameter spelling: `NVFP44Over6Config(err_mode="MSE")`. |
 | `FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH` | unset | `flashinfer/quantization/nvfp4_quantization_utils.py`, `csrc/nv_internal/cpp/common/envUtils.cpp` | **Legacy.** `1` evaluates the 4over6 candidate error in fp16 (a packed `f16x2` dequant of each candidate) instead of exactly in fp32. Cheaper, and on near-ties it can pick the other candidate, so it is a distinct recipe rather than a pure speed knob. Same gating as `FLASHINFER_NVFP4_4OVER6_ERR_MODE`. Parameter spelling: `NVFP44Over6Config(err_use_fast_math=True)`. |