perf(megakernel): MK-KDA 2라운드 -- gates bf16 mma GEMM, p1+p2 병합, split-K 상한, conv 언롤: 스팬 199 → 176 µs - #230
Conversation
…split-K 상한, conv 언롤, o_proj L2 예열 스팬 199 → SPAN_FINAL µs(배수 진단, acc=3), 벤치 320 → BENCH_FINAL µs. - p1 gates: 워프당 내적(행 재사용, 5단 셔플)을 cp.async 32행×128 타일 (8 KB, 동적 smem 이중 버퍼) + mma.sync m16n8k16 bf16 GEMM 으로: 14.5 → 6.2 (하한 4.5, 가중치 1 MB). - p1+p2 병합: conv 는 p0 의 qkv 만 읽으므로 gates 뒤 배리어 제거(대기 3.4 → 0). - split-K 상한: o_proj(k=2048) 스윕 r=3 40.8 / r=2 33.4 / r=1 35.5 / r=4 43.8 → 슬라이스 ≥ 8 k-블록(비용 모델의 "라운드 수"는 파이프라인 채움 비용을 모른다). p5 39.7 → 35.1; GEMM 형상 불변(n=1024 첫 발사 32.7 → 30.7). - conv: 런타임 인덱스 람다(st/kept 가 로컬 메모리)를 스펙 창 8 로 언롤한 선택식으로: 6.6 → 4.7. 호스트가 mql ≤ 8 을 검사. - delta 동안 노는 16 블록이 o_proj 팩 4.5 MB 를 prefetch.global.L2 로 예열: p5 PREFETCH_RESULT. - 프로브 노브 VLLM_GLM53_MK_KSR_IN/OUT(split-K 강제, 진단 전용). 수치는 모든 단계에서 동일 급(out 5~6e-4, conv_state 0, rec_state 1.6e-6), VERDICT PASS. MEASUREMENTS.md 10차, README 단계 예산. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
🟡 Changes recommended
It introduces launch-breaking RUNBOOK commands (profile keys passed via EXTRA_ENV), contains a README/code inconsistency about L2 prefetch, and needs a stricter KDA mql contract to avoid incorrect conv_state writes under narrower query lengths.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR updates the GLM-5.3 megakernel KDA implementation to reduce per-layer decode latency (notably by reworking the gates path and tightening KDA split-K/conv behavior), and concurrently removes previously-added experimental overlay modules and associated probes/docs that are no longer shipped in the glm53 profile.
Changes:
- MK-KDA perf work: gates rewritten as cp.async + bf16 tensor-core MMA GEMM, phase barrier removed between gates/conv, split-K capped to keep ≥8 k-blocks per slice, conv unrolled to the spec window with an
mqlguard and split-K probe knobs. - Contract updates:
tests/test_logic.pyupdated to assert the new KDA kernel invariants/text signatures. - Cleanup/revert: remove
glm53_prep_fused,glm53_async_dflash, andglm53_indexer_gate_splitkoverlays plus related probes/docs, and drop them fromprofiles/glm53.envand the composedbuild/glm53/manifest.tsv.
File summaries
| File | Description |
|---|---|
| tools/trace_step_timeline.py | Removed trace analysis helper script. |
| tools/trace_step_composition.py | Removed trace analysis helper script. |
| tests/test_logic.py | Updates overlay symbol contract checking and adds/updates KDA megakernel contract assertions. |
| STEP_KERNEL_MAP.md | Removes previously-added supplemental trace decomposition sections. |
| RUNBOOK_KERNEL_CAMPAIGN2.md | Updates megakernel experiment commands (shadow/arm) and removes EXP-7/8/9 docs. |
| profiles/glm53.env | Drops removed experimental modules and their VLLM_* knobs from the profile. |
| probes/run_prep_fused_check.sh | Removed prep-fused probe wrapper. |
| probes/prep_fused_check.py | Removed prep-fused probe implementation. |
| probes/moe_gate_tile_sweep.py | Removed MoE gate tile sweep probe. |
| probes/indexer_gate_check.py | Removed indexer head-gate split-K probe. |
| overlay/modules/glm53_prep_fused/requires | Removed prep-fused module dependency list. |
| overlay/modules/glm53_prep_fused/README.md | Removed prep-fused module documentation. |
| overlay/modules/glm53_prep_fused/manifest.tsv | Removed prep-fused module manifest. |
| overlay/modules/glm53_model_wiring/glm5next_model.py | Removes prep-fused import/install hook from model wiring. |
| overlay/modules/glm53_model_wiring/glm53_prefill_fastpath.py | Removes indexer head-gate helper indirection; uses stock torch.mm directly. |
| overlay/modules/glm53_megakernel/README.md | Updates MK-KDA phase budget narrative and adds split-K probe knob mention. |
| overlay/modules/glm53_megakernel/glm53_megakernel.cu | Core MK-KDA changes: gates MMA GEMM, barrier removal, conv unroll, split-K cap, alignment + mql checks, ksr probe knobs. |
| overlay/modules/glm53_indexer_gate_splitk/requires | Removed indexer-gate split-K module dependency list. |
| overlay/modules/glm53_indexer_gate_splitk/README.md | Removed indexer-gate split-K module documentation. |
| overlay/modules/glm53_indexer_gate_splitk/manifest.tsv | Removed indexer-gate split-K module manifest. |
| overlay/modules/glm53_indexer_gate_splitk/glm5next_attention.py | Removed indexer-gate split-K attention overlay file. |
| overlay/modules/glm53_indexer_gate_splitk/glm53_indexer_gate.py | Removed indexer-gate split-K kernel file. |
| overlay/modules/glm53_async_dflash/README.md | Removed async-dflash module documentation. |
| overlay/modules/glm53_async_dflash/manifest.tsv | Removed async-dflash module manifest. |
| MEASUREMENTS.md | Adds MK-KDA “10차” measurement section; removes large EXP-7/8/9 writeups. |
| build/glm53/manifest.tsv | Removes composed rows for the deleted experimental overlays. |
| build/glm53/glm5next_model.py | Mirrors removal of prep-fused install hook in composed build output. |
| build/glm53/glm5next_attention.py | Removed composed attention overlay file for indexer-gate split-K. |
| build/glm53/glm53_prep_fused.py | Removed composed prep-fused implementation. |
| build/glm53/glm53_prefill_fastpath.py | Mirrors removal of indexer head-gate helper indirection in composed build output. |
| build/glm53/glm53_megakernel.cu | Mirrors MK-KDA kernel changes in composed build output. |
| build/glm53/glm53_indexer_gate.py | Removed composed indexer-gate split-K kernel file. |
Review details
- Files reviewed: 33/36 changed files
- Comments generated: 4
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| EXTRA_ENV="VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_KDA_SHADOW=1" \ | ||
| bash launchers/start-glm53-nvfp4-tp4.sh | ||
| # bench-tp4 1회 내내 [megakernel] kda shadow 로그에 DRIFT 0 확인 | ||
|
|
||
| # 4. 브래킷 — 세그먼트별 개별 암(MHC와 GEMM은 별도 부팅으로 분리) | ||
| VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_MHC=1 \ | ||
| EXTRA_ENV="VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_MHC=1" \ | ||
| bash launchers/start-glm53-nvfp4-tp4.sh # cand A | ||
| VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_GEMM=1 \ | ||
| EXTRA_ENV="VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_GEMM=1" \ | ||
| bash launchers/start-glm53-nvfp4-tp4.sh # cand B (MHC 합침은 그 다음) |
| a.mql = (int)ints[2]; | ||
| TORCH_CHECK(a.mql <= KDA_NQ_MAX, | ||
| "kda: max_query_len over the unrolled conv window (KDA_NQ_MAX)"); |
| owners = {} # dotted module path -> (module dir, source path) | ||
| for manifest in sorted(glob.glob( | ||
| os.path.join(REPO, "overlay", "modules", "*", "manifest.tsv"))): | ||
| moddir = os.path.dirname(manifest) | ||
| for raw in open(manifest, encoding="utf-8"): | ||
| line = raw.rstrip("\n") | ||
| if not line or line.startswith("#"): | ||
| continue | ||
| source, target = line.split("\t")[:2] | ||
| if not source.endswith(".py"): | ||
| continue | ||
| # Targets are absolute for image-bound overlays and relative to the | ||
| # package root for portable ones, so anchoring on "/vllm/" silently | ||
| # skipped every portable module -- including moe_gate_sm121, the one | ||
| # this check exists for. | ||
| if target.startswith("vllm/"): | ||
| rel = target | ||
| elif "/vllm/" in target: | ||
| rel = "vllm/" + target.split("/vllm/", 1)[1] | ||
| else: | ||
| continue | ||
| dotted = rel[:-3].replace("/", ".") | ||
| owners[dotted] = os.path.join(moddir, source) |
| cp.async + `mma.sync m16n8k16 bf16` GEMM over 32-row weight tiles; the | ||
| conv is unrolled to the 8-token spec window (the host refuses a wider | ||
| `max_query_len`); the delta rule runs two blocks per head (rows split, S | ||
| register-resident, per-token state stores staged through smem) while the | ||
| 16 head-less blocks warm L2 with the o_proj pack (`prefetch.global.L2`); | ||
| p4 emits the o_proj's fp8 A tiles itself so p5 starts without a prologue; |
fix(repo): main 복구 -- #230 이 되돌린 다른 세션의 작업 33 파일 복원 + KDA 2라운드 재적용
…라운드는 유지 #230 의 커밋(7a287e0)은 0bd06b7 에서 갈라졌으면서 glm53_prep_fused·glm53_async_dflash· glm53_indexer_gate_splitk 모듈, probes/tools 5개, glm5next_model.py 의 prep-fused 훅, 프로필 MODULES/노브, 테스트 419줄, 원장·RUNBOOK·커널맵 항목을 삭제했다(#224 와 같은 패턴: 오래된 작업 트리로 add -A). 유지한 것 = #230 이 실제로 작성한 것: 메가커널 .cu/README(KDA 2라운드), 테스트의 KDA 계약 2 hunk, 원장 10차 항목. 나머지는 이 브랜치 (0bd06b7 + 리뷰 처리분) 그대로. 41,247 checks OK. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
이 PR 의 커밋 7a287e0 은 0bd06b7 에서 갈라졌는데 |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 7a287e0. Configure here.
| #pragma unroll | ||
| for (int j = 0; j < NQ_MAX; ++j) v = (q == j) ? xin[j] : v; | ||
| } | ||
| a.conv_state[sbase + i] = v; |
There was a problem hiding this comment.
Conv state write drops older history
Medium Severity
The unrolled conv state write only keeps three prior values (kept[0..2]) and maps every i >= 2 onto kept[2]. keep is conv_width - nq_tok (10 minus the request length), so any spec-verify sequence shorter than 7 tokens writes the wrong prefix into conv_state. Later steps then convolve from a corrupted window. Host code only checks mql <= 8, and _kda_eligible still admits those shorter cu_seqlens spans.
Reviewed by Cursor Bugbot for commit 7a287e0. Configure here.
요약
KDA 커널 스팬(L2 배수, acc=3) 199 → 176 µs, 벤치 MK-KDA 320 → 280 µs(stock 640
650). 수치는 모든 단계에서 동일 급(out 56e-4, conv_state 0, rec_state 1.6e-6), VERDICT PASS. GEMM 형상은 불변(n=1024 첫 발사 32.7→30.7).mma.sync m16n8k16 bf16GEMM 으로prefetch.global.L2최종 176 = in_proj 76 + gates 6 + conv 4 + delta 34 + norm 0.5 + o_proj 35 + 배리어 ~17. in_proj·o_proj 는 W4 스트림의 DRAM 속도 그 자체(~190 GB/s 유효)라 커널 안에서 더 줄일 것이 거의 없습니다.
VLLM_GLM53_MK_KSR_IN/OUT(split-K 강제, 진단 전용). 계약 테스트: mma GEMM, 배리어 4개, split-K 상한, conv 언롤·mql 검사.🤖 Generated with Claude Code
Note
Medium Risk
Megakernel KDA changes alter synchronization, conv state handling under speculative decode, and GEMM split-K heuristics—high impact if drift slips past shadow/contract tests. Removing build-time async-dflash and indexer split-K paths is lower risk if serving no longer depended on those artifacts.
Overview
MK-KDA in
glm53_megakernel.cuis the main change: kernel span 199 → 176 µs and bench 320 → ~280 µs, with the same numerical gates (out/rec_state, VERDICT PASS).The gates phase is rewritten from per-warp dot products to a bf16 tensor-core GEMM (
mma.sync m16n8k16) with cp.async-staged 32-row weight tiles and stricter 16 B alignment checks. Phase 1 and conv (phase 2) no longer sit behind an extra grid barrier because both only read p0 qkv. Short conv uses a compile-time spec window (KDA_NQ_MAX = 8) with unrolled token/history selection instead of runtime-indexed local arrays, and the host rejectsmqlabove that limit.mk_choose_ksrcaps split-K when k-slices are shorter than ~8 k-blocks (favoring r=2 on o_proj), with optionalVLLM_GLM53_MK_KSR_IN/OUToverrides for sweeps.Docs record the 10th megakernel round in MEASUREMENTS.md and trim campaign material: RUNBOOK megakernel examples use
EXTRA_ENV="...", and long EXP-7/8/9 (prep fused, dflash async, indexer split-K) sections are removed from the runbook and kernel map supplements.Removed from
build/glm53/: the vendoredglm53_config_vllm.pyfork (dflash async allowlist) andglm53_indexer_gate.pysplit-K helper;glm53_prefill_fastpath.pygoes back to stocktorch.mmfor the fp32 head-gate.Reviewed by Cursor Bugbot for commit 7a287e0. Bugbot is set up for automated code reviews on this repo. Configure here.