[kacperkadziolka]

fla's state-scan kernel (chunk_gated_delta_rule_fwd_kernel_h_blockdim64) runs in every DeltaNet layer. Triton autotunes it: on the first call it times 12 launch configs and caches the fastest. The cache key has no sequence count and no length, so the first call decides the config for every later pass. In a Kev server that first call is the CudaGraphs probe at load, one sequence of 16 tokens. On Kev-4B (H100) the probe picks BV 32 (BV is the value-block width each program handles). A batch of 16-32 question rows runs the scan 1.3-1.6x slower with BV 32 than with the best config.

This PR gives the scan two autotune caches: one for passes with 4 or more sequences (question rows, batched new states), one for everything else. A small wrapper, in the same style as _FixedNB, swaps the autotuner's cache before each launch:

n = kwargs["k"].shape[0] if kwargs["cu_seqlens"] is None else len(kwargs["cu_seqlens"]) - 1
self.tuner.cache = self.caches.setdefault(n >= WIDE_SCAN, {})

It also turns off Triton's on-disk autotune cache for this kernel. That cache doesn't know about the split and would give both caches the first pick. The cost is one extra autotune per process, on the first wide pass.

Every config produces bit-identical output, so this changes speed, not answers. I checked that on all 12 configs × 69 shapes, and served probabilities were bit-identical in every A/B block below.

Why 4 sequences. I timed all 12 configs on 69 scan shapes on an H100, for 16, 32 and 48 value heads (scripts/scan_config_map.py, runs/scan-config-map-h100-1). Wide passes (16-32 rows, or 4-16 states of 1k-4k tokens) all want BV 64 with 2-3 stages, and the probe's pick is 1.1-1.8x slower there. Short passes are within noise for any config. Single long states already get the best config from the probe. So wide passes are the only place a second cache helps.

Results. Kev-4B (jaredpalmer/kev-4b@6cfce5c2), H100, 16 requests × 6 questions, GPU time per batch (scripts/scan_width_ab.py, runs/scan-width-ab-4b-h100-1). The run alternates main and this PR seven times on one loaded model, each block starting from empty autotune caches and fresh graphs. I set the bar before the run: this PR must be faster than both neighbouring main blocks every time, by more than twice the main-to-main drift.

workload	main	this PR	change	drift
96 question rows, cached states	105.3-106.4 ms	101.1-101.5 ms	-4.4 ms (-4.2%)	1.1 ms
same requests, new states	109.9-110.1 ms	105.6-105.7 ms	-4.3 ms (-3.9%)	0.2 ms
one new 16k / 61k-token state	584-592 / 2110-2123 ms	586-589 / 2115-2124 ms	none	8 / 13 ms

Both short workloads clear the bar; long states are unchanged. On Kev-27B (jaredpalmer/kev-27b@af0e6d55, an earlier version of this change, runs/scan-class-ab-27b-h100-1) the gain is under 1%. With 48 value heads its probe already picks BV 64.

Not covered. Only H100 was measured, so whether 4 is the right split on L4 or L40S is unknown; autotuning still picks the config on each GPU. Kev-0.8B and Kev-9B weren't measured in the model. If question rows later move to their own copy of this kernel (the indexed-state idea in #107 item 4), most of what's left here is the state-pass gain. Adding the sequence count to fla's autotune key upstream would make this wrapper unnecessary; until then it is pinned to fla 0.5.2, like _FixedNB.

Related to #107 (item 4). Doesn't close it.

Test plan

• New test_state_scan_tuned_by_width (CPU, fla stubbed): split at 4 sequences, varlen counted from cu_seqlens, disk cache off
• CI's fast suites, rebased on main: 595 passed, 26 skipped
• test_cuda_graphs_match_eager on an H100 (modal_app.py::gpu_tests): passed
• Kev-4B A/B above: faster on rows and new states, long states unchanged, bit-identical in all 7 blocks

🤖 Generated with Claude Code