fix(mla): pass merge_attn_states' prefill_tokens_with_context at runtime

[gbyu-amd]

What

merge_attn_states_kernel declared prefill_tokens_with_context as tl.constexpr. Every MLA caller (attention_mla.py, the three chunked-prefill merges) leaves it at its num_tokens default, so the kernel was compiled once per distinct batch token count: a JIT compile on the serving path, which stalls every TP rank.

The value is only compared against token_idx (tl.program_id), a branch that was already a runtime branch. This PR makes it a plain int with do_not_specialize, so the generated IR differs only in reading the bound from a kernel argument instead of a constant.

Why it matters

Kimi-K3 AgentX (aiperf, 3600 s window) on MI355X, TP8, with a fresh triton cache:

	c1	c14
in-window triton compiles	291 / h	712 in the first ~20 min
of which merge_attn_states_kernel	274	711

Requests whose prefill overlapped one of these compiles were ~0.8–0.9 s slower to first token than the same request (paired by conversation/turn) without one. At c14, 445 of 520 requests were hit and TTFT p50 doubled.

Because the key is the token count, a warm cache never fully covers it: the observed values ran from 2 to 8192.

Validation

• Compile count, 60 random batch sizes × OUTPUT_LSE on/off: 120 compiled kernels before, 6 after (OUTPUT_LSE × the num_tokens ==1 / %16 / other classes).
• Bitwise identical to the old kernel over 2304 cases: T ∈ {1..4099}, H ∈ {1,5,12,16}, D ∈ {128,192}, bf16 and fp8 (e4m3fn) output, ±inf LSE edge rows, prefill_tokens_with_context ∈ {None, 0, 1, T//2, T}.
• tests/test_merge_attn_states.py: 22 passed.
• do_bench, interleaved A/B × 5 rounds, H=12 D=128: OUTPUT_LSE=False within +0–2%, True within −1%..+4% (≤ 1.4 µs/call).
• Serving, same c1 workload: 0 in-window compiles of this kernel (was 274/h). TTFT p50 is back at the no-JIT level: 735 ms, vs 1418 ms before on the same requests.

The vLLM plugin paths call vLLM's own merge_attn_states, where this argument is already a runtime int, so they are unaffected.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2378 --add-label <label>