Skip to content

[Config] [Kimi K3 fix] Drop gfx942/gfx950 opus rows from the BF16 tuned-GEMM tables - #4915

Open
hyukjlee wants to merge 1 commit into
ROCm:mainfrom
hyukjlee:hjlee/drop-opus-rows-gfx942-gfx950
Open

hyukjlee wants to merge 1 commit into
ROCm:mainfrom
hyukjlee:hjlee/drop-opus-rows-gfx942-gfx950

Conversation

@hyukjlee

Copy link
Copy Markdown

Motivation

The a16w16 split-K opus kernels keep their partial-sum buffer in a per-stream hipMalloc
registry on gfx942/gfx950. Any HIP graph capture that is the first thing to touch that registry
aborts:

[AITER] opus_gemm.cu:487 splitk workspace not initialized for the current CUDA stream.
Call aiter.opus_gemm_workspace_init() inside with torch.cuda.stream(s): ...

Frameworks that capture cudagraphs (vLLM, SGLang) never make that call — vLLM contains no
opus_gemm references at all — so every tuned row selecting libtype=opus on these two archs
is a latent engine crash
for them. The row itself is fine in eager mode; it is only unusable
under capture.

gfx1250 is unaffected: #4246 gave it a torch.empty workspace (_get_opus_workspace) that is
HIP-graph-capture aware, and its docstring states "No eager pre-warm is required." Its opus
rows are therefore kept.

Technical Details

Removes 259 rows where libtype == "opus" and gfx ∈ {gfx942, gfx950}:

arch opus rows workspace path action
gfx1250 127 torch.empty, capture-safe (#4246) kept
gfx942 183 legacy per-stream registry (kids 10200–10300) removed
gfx950 76 legacy per-stream registry (kids 200–300, 1200–1300) removed
  • 12 files touched, 259 deletions, 0 additions. No row is edited; whole lines are removed so
    every surviving byte is untouched.
  • Removed shapes fall through get_GEMM_A16W16_config to its default branch (asm / skinny /
    torch), all capture-safe.
  • Because the tuned table is indexed by ("gfx", "cu_num", ...), removing a row only affects the
    arch it names.

This is a mitigation, not the root fix. Porting #4246's torch.empty workspace to the
gfx942/gfx950 split-K kids would remove the per-stream registry entirely and let these rows come
back. That is the change I would rather see; this PR unblocks capture-mode users in the meantime.

Framework-side workarounds were tried first and are not sufficient — calling
opus_gemm_workspace_init() on the capture stream from vLLM
(vllm-project/vllm#53238, now closed) lets the PIECEWISE ladder finish but
the abort simply relocates to FULL capture, because the registry is per-stream and the framework
cannot reliably pre-register every stream that will capture.

Test Plan

Kimi-K3 MXFP4 TP8 (bf16 GEMM path), 8x MI355X (gfx950), ROCm 7.2.3, aiter v0.1.20, vLLM main
with cudagraphs enabled, --max-num-seqs 20 (concurrency 8, 48-entry PIECEWISE ladder).

  1. Start the server with the stock tuned tables → observe the abort.
  2. Start it with the opus rows removed → observe capture completing and the server serving.
  3. Run gsm8k (5-shot, lm_eval, real verification sampling) to confirm numerics are unaffected.
  4. Repeat at concurrency 1 to characterise the failure.

Step 2 was executed by pointing AITER_CONFIG_GEMM_BF16 at a table generated by the same rule.
On gfx950 that table is equivalent to this diff by construction: the lookup is arch-keyed, so the
gfx942/gfx1250 rows are not selectable on this hardware.

Test Result

stock tables opus rows removed
opus_gemm.cu:487 aborts 24 0
PIECEWISE capture dies at 58% 100%
FULL capture not reached 100%
engine Worker proc died unexpectedlyEngine core initialization failed starts, serves
gsm8k (128 samples) n/a — never starts 1.00 strict / 1.00 flexible

Concurrency dependence: at concurrency 1 the capture ladder is small enough that no shape
routes to opus, and the stock tables pass. The failure only appears once the ladder is wide enough
to reach the affected shapes (observed at M:20/19/18, N:6288). Low-concurrency smoke tests do not
catch this.

The a16w16 split-K opus kernels keep their partial-sum buffer in a per-stream
hipMalloc registry on gfx942/gfx950. Any HIP graph capture that is the first
thing to touch it aborts:

  [AITER] opus_gemm.cu:487 splitk workspace not initialized for the current
  CUDA stream. Call aiter.opus_gemm_workspace_init() ...

Frameworks that capture cudagraphs (vLLM, SGLang) never make that call, so
every tuned row selecting libtype=opus on these two archs is a latent engine
crash for them. gfx1250 is unaffected -- ROCm#4246 gave it a torch.empty workspace
that is capture-aware -- so its 127 opus rows are kept.

Removes 259 rows (183 gfx942 + 76 gfx950); those shapes fall through
get_GEMM_A16W16_config to asm/skinny/torch, all capture-safe. Deletions only,
no row is edited.

Reproduced on Kimi-K3 MXFP4 TP8 / 8x MI355X (gfx950) under vLLM: capture dies
at 58% of the PIECEWISE ladder on M:20/19/18 N:6288, and after warming the
workspace on the capture stream it merely relocates to FULL capture -- the
per-stream registry cannot be reliably pre-registered from the framework side.
Note the failure is concurrency-dependent: at concurrency 1 the ladder never
reaches an opus shape, so low-concurrency smoke tests pass.

This is a mitigation, not the root fix. Porting ROCm#4246's torch.empty workspace
to the gfx942/gfx950 split-K kids would let these rows come back.

Signed-off-by: hyukjlee <hyukjlee@amd.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@hyukjlee
hyukjlee requested a review from a team August 21, 2026 14:31
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every PR:

  • ✅ Pre-checks (submodule verification, code formatting)
  • ✅ Aiter op tests (gfx942 + gfx950)
  • ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label Tests
ci:gfx1250-ffm-triton Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4915 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

@github-actions github-actions Bot changed the title [Kimi K3 fix] Drop gfx942/gfx950 opus rows from the BF16 tuned-GEMM tables [Config] [Kimi K3 fix] Drop gfx942/gfx950 opus rows from the BF16 tuned-GEMM tables Aug 21, 2026
@zufayu
zufayu requested a review from yzhou103 August 24, 2026 01:43
@yzhou103

Copy link
Copy Markdown
Contributor

We'll fix this at the root — porting the workspace to torch.empty (capture-safe) for gfx942/gfx950 split-K kids, same as the prewarm path already does for the capture stream. That avoids the hipMalloc-inside-capture abort without sacrificing any tuned performance. This PR deletes 259 opus rows, but only 70 of them are split-K (the ones that actually touch the per-stream hipMalloc registry). The other 189 rows are non-splitk kids (flatmm/persistent/mono_tile) that never touch the workspace and are capture-safe — removing them drops tuned performance for no reason.

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