Skip to content

(WIP) perf(comm): extend fused AttnRes through M16 - #1142

Draft
panditsa wants to merge 1 commit into
lightseekorg:sanket/wip-k3-iris-producer-directfrom
panditsa:sanket/wip-k3-iris-attnres
Draft

(WIP) perf(comm): extend fused AttnRes through M16#1142
panditsa wants to merge 1 commit into
lightseekorg:sanket/wip-k3-iris-producer-directfrom
panditsa:sanket/wip-k3-iris-attnres

Conversation

@panditsa

@panditsa panditsa commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • extend the fused Iris attention all-reduce plus AttnRes epilogue from M=1 to M=1..16
  • use per-token double-buffered staging with separate ready and consumed epochs
  • keep the public communication wrapper shape-driven

Performance

Isolated fused world-size-8 CUDA-graph latency at head a1a0f8e7:

M Fused all-reduce + AttnRes
1 13.69 µs
2 14.52 µs
4 16.23 µs
8 17.38 µs
16 21.79 µs

Exact parent → head 4K/1K TP8/EP8 M=2, median of three, is flat: 93.98→93.95 tok/s (-0.04%) and 20.534→20.565 ms TPOT (+0.15%). This is expected at this stack boundary because the Kimi runtime still gates batched AttnRes fusion to M=1. #1146 removes that runtime gate; this PR provides and validates the batched kernel capability.

Tests

  • world-size-8 numerical comparison against the split all-reduce + AttnRes reference passes for M=1/2/4/8/16
  • each M is exercised repeatedly in eager mode and through CUDA-graph capture/replay
  • all 12 measured 4K/1K A/B requests completed with valid lengths

Stack

Use per-token double-buffered Iris staging and readiness epochs so the fused attention all-reduce and AttnRes epilogue supports decode batches from M=1 through M=16.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@jhinpan

jhinpan commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Exact #1132 validation of #1141 + this PR + #1146 on current main (9196fa27):

  • forced 131072/1024 TP8/EP8 baseline: 36.80 ms TPOT, 19.7647 tok/s
  • stack, 3 measured runs: 36.25/36.25/36.25 ms; 19.9702/19.9672/19.9734 tok/s
  • mean change: -1.49% TPOT, +1.04% throughput
  • normal 4K/1K DSpark: 155.078 tok/s vs current-main 152.949 (+1.39%), so the 95% protection gate passes
  • forced outputs remain token-identical to ordinary greedy for 4K/1K repeated text and exact 131K/128 natural text

This is a real but small E2E win: the forced round ratio moves only from 2.157x to about 2.125x, not near the 1.5x target. The full orchestration is required; #1141 standalone could not register with SMG.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants