[ROCm] Retune MI355 selective_state_update float16 config on the unified effective_batch grid

[vanshbhatia-amd]

Summary

Re-tunes the AMD Instinct MI355X selective_state_update (Mamba SSU decode kernel) cache_dtype=float16 launch config onto the same effective_batch grid already used by the MI300X (#47945) and MI350X (#48159) configs, so all AMD Instinct SSU configs are sampled on a uniform set of points.

The original MI355 config (#47767) was keyed on a narrower, deployment-specific grid (10 … 5120). This PR regenerates it on the standard grid (128 … 262144), matching MI300X/MI350X. Keyed on effective_batch = batch * nheads, so a config transfers across (model, TP) combos sharing (headdim=64, dstate=128, cache_dtype). bf16 shares this file (the loader canonicalizes bf16 → float16, since the kernel only sees the state bit-width).

This is purely a config re-tune for uniformity across the AMD Instinct line — no code changes.

How it was generated

Produced with the in-tree tuning script on an MI355X:

python -m benchmarks.kernels.benchmark_selective_state_update \
  --dstate 128 --dtype float16 --mamba-ssm-cache-dtype float16 \
  --save-configs --compare --validate

The device name and dtype key in the filename are derived automatically, so the file drops straight into configs/selective_state_update/ and is picked up at runtime with no code changes.

Test plan

• --validate: all tuned effective_batch points pass against the CPU reference (12/12).
• --compare: consistent speedup over the generic heuristic across the swept grid.
• Loaded at runtime: get_ssm_configs() resolves device_name=AMD_Instinct_MI355_OAM,cache_dtype=float16.json (log: Using SSM config from .../cache_dtype=float16.json for selective_state_update); removing it logs Using default Mamba SSU config. Performance might be sub-optimal!.
• End-to-end correctness verified with lm_eval lambada_openai (see Results).

The config only selects Triton launch geometry (BLOCK_SIZE_M, num_warps), not the kernel math, so model outputs are unaffected.

Results

Perf — selective_state_update kernel, tuned vs the current heuristic (M=4, w=4), relative speedup from benchmark_selective_state_update.py --compare:

effective_batch	fp16 speedup
128	1.21×
256	1.38×
1024	1.71×
2048	1.71×
4096	1.93×
8192	2.08×
16384	1.87×
32768	1.97×
65536	1.90×
131072	2.12×
196608	2.28×
262144	2.19×

~1.2× at small batch (latency-bound), ~1.9–2.3× across the decode range.

Correctness — end-to-end, unchanged. lm_eval lambada_openai (500 examples), config present vs removed, AntonV/mamba2-2.7b-hf (headdim=64, dstate=128), fp16 activations, float16 SSM state cache:

config	acc	perplexity
tuned (this PR)	0.706	4.2667
heuristic	0.706	4.2668

Identical acc and matching perplexity — expected, since the config only changes the Triton launch geometry, not the math.

Companion PR retunes the float32 variant on the same grid.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.