[junhaha666]

aiter removed aiter.ops.triton.attention.mla_decode in ROCm/aiter#6633. With current ATOM and aiter main, importing the Kimi MLA/GDN backend raises ModuleNotFoundError even when Triton MLA is disabled; the Triton decode path also depends on the deleted module.

Use the maintained aiter.ops.triton.attention.mla.mla_decode_fwd API for both plain and shuffled KV caches. View plain token rows as scheduler-sized pages and pass the existing block tables, context lengths, query offsets, and FP8 scales. Remove the obsolete CSR-to-dense conversion and decode scratch buffers from the Triton metadata builder.

Validation in jun_atom2 with aiter main at b3ad6c4617:

• 77 tests passed across the new Triton MLA regression tests and existing attention-family, shuffled-KV, and MTP contract tests.
• GPU comparisons against an FP32 PyTorch reference cover plain/shuffled KV, block sizes 16/64, one-token/four-token decode, and BF16/FP8 inputs. FP8-query comparisons account for the kernel's FP8 probability rounding.
• CUDA graph replay matches eager execution exactly after updating context lengths and block tables.
• Black, Ruff, and git diff --check passed.

Full-model end-to-end inference has not been run.

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