perf: unify metadata publication and reuse page tables and Triton kernels

[valarLip]

Forward preparation uploads many small metadata buffers and repacks page tables even when the page mapping is unchanged. Changing batch sizes and prefill lengths also creates new Triton specializations for values that only control runtime bounds.

This PR publishes metadata at its first consumer under a shared ownership contract, reuses page-table state across attention backends, and reuses Triton kernels across dynamic lengths. ATOM_H2D_BACKEND=direct remains the default; packed combines eligible metadata into one H2D followed by scatter.

Changes

• Add checked publication with source lifetime, stream/device identity, semantic padding, and failure handling owned by the existing runner slots. A shared producer decorator checks write permission before host data can be overwritten. Packed arena ownership follows actual packed transfers, including mixed direct/packed publication.

• Share versioned page-table preparation across common attention, MHA/MLA, V4, and V4.1. Unchanged mappings avoid page-ID scans and uploads; appends and request reordering update the affected rows. CPU preparation and successful GPU publication have separate revisions. PP/TBO buffers retain independent state, and V4.1 uses the common entry point.

• Integrate runner, attention, draft/GDN, and Engram consumers at their first-use boundaries. Keep uniform sampling filters on the CPU, reuse independent TBO buffers, and preserve graph padding and fixed destination storage.

• Remove exact-length specialization from SWA writes and 13 additional Triton kernels. Preserve model geometry, tiles, and useful bounded alignment variants. Cached KV quantization keeps vectorized accesses and supports inputs exceeding 2^31 elements; only the existing small-input reduction configuration uses two-way loop unrolling.

• Add ATOM:: runner stage annotations when profiling is active, document the publication contract, and organize H2D regressions by owner/consumer responsibility.

• Bind and retain the model-runner TCPStore before spawning workers. All managed ranks connect as clients, removing the gap between probing a free port and binding it later. Single-node runs allocate the port with port=0; multi-node runs share an explicit fixed port. RapidServe retains separate prefill/decode stores, and standalone ModelRunner callers keep their existing rendezvous path.

Validation

Validation used an AMD Instinct MI355X with PyTorch 2.9.1+ROCm7.1.1 and Triton 3.8.0.

• H2D/shared-page-table commits: 892 passed, 6 skipped; skips require additional GPUs or an external V4.1 reference. The core, profiler, and SWA commits were also checked independently after history cleanup.
• Final dynamic-Triton changes: 176 passed, covering numerical output, tail maxima and sentinels, request boundaries, graph replay, and JIT reuse. The 13 targeted kernels reuse one variant across six changing values within a fixed geometry/alignment class.
• TCPStore startup and related regressions: 123 passed, including real ModelRunner initialization and collectives for TP, DP, TP+DP, PP, PCP, DCP and eight-rank DP; standalone TP2 compatibility; port ownership/release; fixed-port conflicts; and remote-node client selection. Multi-node configuration was tested over loopback, not across physical hosts.
• A real 1,048,576-token × 12-head KV input passed full output-byte validation. The previous implementation failed to compile this shape because its K/V constexpr element counts inferred incompatible integer types.
• Formatting and diff checks passed. Ruff passes for the dynamic-Triton changes and introduces no new findings in the startup fix. Existing diagnostics include three Optional annotation warnings and four exception-handling warnings in the manager.

The suites overlap and their counts should not be added. Performance checks use 1M context and operator/metadata microbenchmarks, not end-to-end model throughput. In a controlled shared-buffer comparison, 4096-token KV quantization measured 11.95 → 11.02 μs after the reduction adjustment (11.14 μs before the Triton changes). Some small dynamic shapes retain overhead, including about 0.30 μs for 257-token KV; the accepted shared-page-table path also retains about 2.4 μs in its B256 single-request subsequent-chunk case. Full-model serving, quality, and throughput validation was not rerun.

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