Granite Mamba2 Inference Optimization

State-space models like Mamba2 replace attention with a recurrent selective-scan, which makes their inference profile very different from a transformer's: long prefill is a chunked parallel scan, while decode is a per-token recurrence over convolution and SSM cache state. Production inference engines squeeze this path with hand-tuned Triton kernels, and the gap between a straightforward eager PyTorch implementation and an optimized one is large.

The workspace is a standalone port of the real Hugging Face GraniteMoeHybridMambaLayer, with weights extracted from the pinned checkpoint ibm-granite/granite-4.0-h-1b-base (layer-0 Mamba weights, final norm, and tied embedding head). Everything runs in bfloat16 on CUDA, which constrains optimization choices such as Triton intrinsics and accumulation precision.

The agent must implement CandidateBlock in /app/src/candidate_impl.py(a stub subclassing the reference is provided) and make the layer's inference path faster without changing its semantics. The performance bar is an optimized implementation built on the same production Triton kernels vendored in /app/vllm_ops/: match it, then beat it.

• Keep the fixed constructor and forward(hidden_states, cache=None, attention_mask=None) signature; the internal cache layout may change as long as the returned cache still exposes conv_state, ssm_state, has_previous_state, and decode-position semantics.
• Match the reference on hidden states, convolution and SSM cache states, and last-token readout logits, across prefill, cached decode, and variable-length padded batches. Speed only counts once parity holds.
• Gains are measured on long prefill, batched variable-length prefill, and per-token decode latency; consistent gains across shapes beat a single outlier.
• torch.compile, Triton, custom CUDA kernels, CUDA streams, and calls into transformers are all allowed. /app/reference_impl.py, /app/task_fixtures.py, and /app/vllm_ops/ must stay untouched; all changes go in the single candidate file.

A ready-to-run workspace at /app: the fixed reference port (reference_impl.py, a clean port of the HF torch_forward path that does not call transformers inside the forward), the extracted checkpoint slice in assets/, the vendored vLLM Triton kernels in vllm_ops/ as building blocks, and fixed utilities in task_fixtures.py including the visible correctness and benchmark workloads.

Local loops are provided: verify_api.py checks parity against both the reference and the pinned transformers implementation, and run_dev_bench.py runs a local latency comparison. The hidden optimized baseline itself is not exposed to candidate code. There is no internet access at run time; the image bakes a ready .venv, so uv run --no-sync python ... works offline, and a sandbox-timer CLI reports the remaining budget.

A clean-room verifier runs in a fresh container on the captured /app: it restages pristine copies of the fixed files (task_fixtures.py, reference_impl.py, vllm_ops/) so only the candidate file carries over, then runs a correctness gauntlet followed by a paired GPU benchmark. Candidate code executes only inside an unprivileged worker process; the trusted parent owns workload generation, comparisons, and timing.

• Correctness is checked against both the eager reference and the pinned transformers implementation on hidden prefill and decode workloads: hidden states, conv and SSM cache states, readout logits, and readout KL divergence. Any failure zeroes the result before performance is considered.
• The timed path is additionally validated at the large benchmark shapes, with tolerances calibrated from the trusted baseline's own drift, and an output digest is compared on every timed call, so returning garbage-fast output only on the timed shapes is a hard correctness failure.
• Speed is a paired ABBA benchmark against a hidden optimized baseline built on the same Triton kernel family, timed exclusively on the trusted parent's wall clock (worker self-reported timings are never scored), with median-of-pairs aggregation and a frozen Triton autotune cache.
• The reward is the fraction of starter runtime eliminated: with G the geometric-mean speedup over the unchanged starter (measured against a live no-op floor), a correct candidate scores 0 when G ≤ 1 and 1 − 1/G otherwise; 2x maps to 0.5, 4x to 0.75.
• Source scans reject references to verifier internals and any delegation to the trusted baseline; build failures and a missing candidate file zero the result.