Skip to content

[Perf] [AMD] Kimi-K3: fuse attention-residual aggregation with TP collectives - #46656

Open
aryatsc wants to merge 9 commits into
sgl-project:mainfrom
modal-projects:devin/pr-attnres
Open

aryatsc wants to merge 9 commits into
sgl-project:mainfrom
modal-projects:devin/pr-attnres

Conversation

@aryatsc

@aryatsc aryatsc commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

K3's attention-residual aggregation ran as a separate all-reduce plus aggregation kernels per layer.

Modifications

Accuracy Tests

  • The existing test_attn_res.py and test_attn_res_mix_large_tokens.py pass on MI355. This PR adds no new test.

Full K3 stack on MI355X TP8 (DFlash, prod config) vs B300, McNemar:

  • GPQA-D 89.4 vs 87.9 (p=0.55)
  • AIME25 96.7 vs 93.3 (p=1.0)
  • MATH500 92.2 vs 92.2 (p=1.0)
  • MMLU-Pro 89.6 vs 90.0 (p=0.77)
  • GSM8K 97.2%

Speed Tests and Profiling

Throughput change from turning this feature off on the full stack: -2.7/-1.1/-3.4% on the agentic replay at C=1/8/32, and -6.3/-2.2% at 90k prefix at C=1/8 (mean -3.1%).

Dependencies

Checklist


CI States

Latest PR Test (Base): ❌ Run #42125653195
Latest PR Test (Extra): ❌ Run #42125652668
Latest PR Test (AMD ROCm 10): ❌ Run #42125652934

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

aryatsc and others added 3 commits October 26, 2026 17:46
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	python/sglang/srt/environ.py
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	python/sglang/kernels/ops/moe/moe_sorting_small.py
@github-actions

This comment has been minimized.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	python/sglang/kernels/ops/moe/__init__.py
@github-actions

This comment has been minimized.

…erCore.run

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

This branch has not been deployed

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant