[Feature] Integrate the AITER MK1 persistent decoder

[ssharma4-amd]

Dependency

Depends on ROCm/aiter#4975.

This PR should remain draft until the AITER MK1 package is available to upstream ATOM CI.

Summary

Adds an opt-in GPT-OSS-120B persistent decode path backed by aiter.MK1 on AMD Instinct MI355X (gfx950). The integration is BF16-only for the initial upstream version.

Design

• ATOM owns serving policy, scheduling, page reservation, request state, and its existing KV cache.
• AITER owns the compiled persistent checkpoint, native binaries, tensor binding, and quantum execution.
• The megakernel binds directly to the existing ATOM KV cache; no duplicate serving cache is allocated.
• Persistent weights are loaded only when persistent decoding is enabled.
• Prelaunch rejection preserves ordinary ATOM fallback. A failure after native execution is terminal because cache state may have changed.

CLI

The integration adds only:

• --persistent-decoder {off,auto,required}
• --persistent-decoder-checkpoint <local-directory>

The persistent path currently requires the existing BF16 KV-cache configuration. There are no provider-backend, persistent KV-dtype, maximum-quantum, debug, verification, or statistics flags.

Supported envelope

• GPT-OSS-120B revision b5c939de
• AMD Instinct MI355X / gfx950
• BF16 KV cache
• TP1 / PP1 / PCP1 / DCP1
• Batch-one greedy decode
• Maximum resident quantum of eight tokens
• Prefix caching and chunked prefill disabled

Validation

Final source tree validation: AMD MI355X

Cross-validation uses the squashed AITER commit 4b492bb46cbb81deac584d5d68bd4a73e30bf9ae from ROCm/aiter#4975. Its manifest-free native ABI packaging was independently validated by Fleet jobs job-e51aa84a and job-687ce457.

• 13 focused policy and package tests passed.
• Persistent-only checkpoint validated with 471 tensors.
• ATOM chat produced coherent output.
• Serving benchmark completed 60/60 requests.
• Mean TPOT was 1.967 ms for the final BF16-only ATOM tree.
• Mean TPOT was 1.962 ms when cross-tested with the squashed AITER branch.
• Confirmed aiter.persistent_decoder is absent and all integration imports use aiter.MK1.
• Confirmed no libanvil.so dependency.

GSM8K evaluation

Both runs used a BF16 KV cache, all 1,319 GSM8K samples, three-shot prompting, and a 2,048-token generation limit.

Configuration	Samples	Strict exact match	Flexible extraction
Pristine ATOM BF16 baseline	1,319	35.18%	88.17%
ATOM + AITER MK1 persistent decoder	1,319	9.55%	90.52%

The persistent run completed 379,525 native persistent calls, emitted 380,635 tokens, and recorded zero prelaunch or backend failures. Flexible answer extraction improved by 2.35 percentage points. The lower strict-format score indicates a response-format difference; this evaluation does not establish exact token-level equivalence between ordinary and persistent decoding.

Test coverage

tests/test_persistent_decoder_policy.py covers scheduler admission and rejection rules, quantum sizing, BF16 cache geometry, the minimal CLI surface, and BF16 defaults. The test module is CI-only and is not part of the installed ATOM runtime package.

Known validation boundary

These runs validate packaging, loading, native execution, chat coherence, serving performance, and GSM8K answer extraction. Exact standard-versus-persistent numerical equivalence remains a separate correctness gate and is not claimed by this PR.

[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 2011 --add-label <label>