[rebklee]

Summary

_apply_attn_res_impl copies block_residual to a contiguous tensor before the fused AttnRes
kernel. The kernel already reads the block residual by its token and block strides, so only the
hidden dimension has to be unit-stride; this PR copies only when it is not. Outputs are
bit-identical. Production builds its block residual with torch.cat, which is contiguous, so it
never copied and nothing changes there. Mono decode, which passes slices of one block bank, saves
108 us of GPU time per target forward (0.8%) at the agentic recipe's C1 launch.

Motivation

Mono decode keeps the closed blocks in one bank, blocks[T, 8, H], and passes a seam the blocks
before it as blocks[:, :n]. With more than one row and n < 8 that slice has strides
(8 H, H, 1), so the call copied it first: 22 copies per multi-row step (the attention seams of
MLA layers 3, 7, ..., 83, and layer 0's MLP seam; layers 87, 91 and 92 take the whole bank,
which is contiguous), 9.7 MB per step at C1.

Changes

• br = block_residual if block_residual.stride(-1) == 1 else block_residual.contiguous().
• The module docstring and the two comments say the block residual is read by stride.

Performance

8x MI355X, the agentic recipe's C1 launch (mono decode on, 8-row verify steps), main against main
with this PR. GPU time of the 22 seams that copied, per target forward in rank 0's profile:

the 22 seams, a target forward	main	this PR	change
block-residual copies	116.2 us	none	-116.2 us
_attn_res_fused_kernel	226.7 us	234.4 us	+7.7 us
total	342.9 us	234.4 us	-108.5 us

The AttnRes calls at these seams take 0.35 us longer each (+3.4%). A microbenchmark of the seam
traces that to the call losing the copy in front of it (it now starts right after K2, on input
no kernel has just read), not to the strided read. Median ITL, 14.75 ms on main and 14.65 ms with
this PR, is consistent with the net, but one server run per arm cannot resolve 0.1 ms.

Method

• rocm/atom-dev:nightly_202611011508-rocm10.0, ATOM main 606aceb11 mounted over the image's,
  with and without this commit. Every arm preloads /usr/lib/x86_64-linux-gnu/libgomp.so.1,
  without which main's loader (feat(mono): shared DSA indexer mono launch; Kimi-K3 mono to 64 rows, plain decode and ReplaySSM #2871) cannot resolve omp_set_num_threads in this image.
• recipes/Agentic-Kimi-K3.md's C1 launch: TP8, DCP1, DSpark with 7 speculative tokens, synthetic
  acceptance length 3.84, mono on with main's preconditions (fp32 state, five projections
  excluded from ptpc_fp8).
• atom.benchmarks.benchmark_serving, random dataset, 8192 in / 1024 out (range ratio 0.8),
  --ignore-eos, 5 prompts at concurrency 1, seed 0: a warmup wave, then two timed waves. With
  speculative decoding the median ITL is the step period.
• Profile: 3 s of torch profiler in a third wave. The target forward is the largest graph replay;
  the 22 seams are its first 22 _attn_res_fused_kernel calls. Its last 4 read the whole bank
  and are within 0.1 us of main's. Each row is the median over replays of a sum over a forward.
• Seam microbenchmark, one MI355X: CUDA graphs of the 21 MLA seams at C1's widths, each behind a
  264 MB weight-streaming GEMM in K2's place, reproduce the slowdown (+0.52 us a call). A tiny
  kernel after the GEMM plus a just-read slice recover all of it; without the GEMM there is none.

Limitations

• A block residual whose hidden dimension is strided is still copied.
• The kernel's offsets stay i32: with the bank's 8-block stride, t * stride_br_t overflows past
  37,449 rows. Mono decode steps are at most 64 rows; production's torch.cat bank has the same
  limit.
• perf(kimi-k3): stream attn-res candidate rows through LDS with TDM on gfx1250 #2814 (gfx1250 TDM AttnRes) merges cleanly, but its TDM check does not test stride(0), which
  a mono slice now sets to 8 H. That is aligned and its kernel takes the stride as an argument,
  so nothing breaks, but the check should cover it; and on gfx1250 the in-place test below would
  see the TDM launch. Happy to adapt whichever lands second.

Testing

gfx950 (MI355X), rocm/atom-dev:nightly_202611011508-rocm10.0.

• tests/model_ops/test_kimi_k3_attn_res_strided.py (new, GPU): 34 passed.
  • Every bank width 1..8 at 1, 6, 16, 48 and 130 tokens in all six output variants, and slices
    past block 0 and past token 0, are byte-identical to slice.contiguous(), scales included.
  • The kernel is handed the slice itself. Restoring the unconditional .contiguous() fails this
    test and none of the other 33.
  • A strided hidden dimension is still copied, and matches.
• Kimi-K3 and mono tests: 244 passed on GPU, none skipped; 251 on CPU, as on main (both with
  a tests package in the image's site-packages, which shadows ATOM's, masked).
• black and ruff clean.

python3 -m pytest -q tests/model_ops/test_kimi_k3_attn_res_strided.py

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

[valar-cop]

LGTM

ci green.

[HaonanWang98]

LGTM

[valar-cop]

ci triage for db6ae8f:

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/41769882692

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/41769882570

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/41769882784