MK GEMM: the shared expert's pair on a barrier-free local-quant kernel (EXP-22, knob off) + two review rounds - #307
Conversation
…k GEMM launches (VLLM_GLM53_MK_LOCALQ) The 30-45 us class of the decode step -- the shared expert's gate_up [1024 x 4096] and down [4096 x 512], 86 launches a step, 2.4 ms of it exposed on the critical path (28차) -- moves 1-2 MB each and pays ~25 us of fixed cost: a grid-wide A quant (one k-block per block into g_mk_aq), the publishing barrier and its skew (the x loads queue behind the hoisted W fill), the sxs round trip, and 16 idle blocks held in all of it. When a launch has at most one unit per block, the standalone kernel now takes an LQ instantiation of the phase: every block quantizes the A k-blocks of ITS unit straight into smem from x in L2, as it stages them (load a whole iteration ahead like the tile copy it replaces), with quant_store's arithmetic -- so the mma reads the same bytes and the output is bitwise the global path's. No barrier, no global tiles, idle blocks exit at once and hand their SMs to the PDL successor. The host picks it (mk_units <= grid), the phase re-derives the condition, KDA's inlined phases stay the <false> instantiation. VLLM_GLM53_MK_LOCALQ=0 is the kill switch. Bench: --gemm-sweep (local x split, bitwise column, replay) and --stamps (phase stamps of a MK_PHASE_TS build) over the small shapes; gemm_plan / set_probe pybind for in-process sweeps; a concurrent probe that captures the pair under the served b12x MoE call on a forked stream, the way the step runs it. Numbers land in MEASUREMENTS.md once measured. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ernel keeps its 80-register binary One kernel holding both instantiations allocated registers for the union (ptxas 80 -> 128), so the global path's code was no longer the binary it was measured as. mk_gemm_lq_kernel carries the <true> instantiation (127 registers, no spill); the host launches it when the plan says localq and the two kernels resolve to the same resident grid. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… and probes Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… quant off the sync section, background-gated; the review's fixes; ledger 29차 29차 (srv2): the first form lost standalone (m=8 n=1024 24.7 -> 28.7 us, m=32 +12 us at the same split; the global prologue it skips is 5.5 us on the stamps, not the 4차 cold-launch 8-10) and under the routed MoE kernel hid the pair only when the MoE was issued first (44 -> 17 us exposed; the pair-first order of serving's aux stream stayed at 42: 48 blocks take every SM). v2: lq_quant before the pipeline wait over every row (the per-row break serialised the shuffle chains), the lq kernel launched on as many blocks as it has units (VLLM_GLM53_MK_LQ_GRID caps it), the fp8-dense hook marks the shared expert's linears background and the knob is 0 (default, declared in the profile) / 1 = those / 2 = every one-unit-per-block launch. Review (10 angles, verified): one host plan for the launch and the bench pybind (padded n, the lq-grid check inside); MKGemmCtx::localq gone; set_probe -1 = back to the env; env values range-checked; the boot self-test exact-gates BOTH kernels and requires them bitwise equal; the stamp reader keeps blocks that ran a unit (global-path idle blocks stamp slot 4); sweep: no duplicate split rows, --gemm-shapes feeds it, the stock table is kept, split rows gated at the summation-noise class; the concurrent probe imports the served MoE fixture, prints plans per knob setting and amortises the graph-launch gap; idle blocks of the local path leave before the PDL wait; mk_split_ok shared by phase and host; stale MK_SPLIT comments; RUNBOOK order item placed after item 12; bracket.py snapshots the new knobs. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…29차 v2 numbers The v2 form clamped every row like the prologue and reduced all four, so the m=8 launch did four times the work of its one live row per warp (pair alone 35.2 -> 43.0 us). A block-uniform row bound keeps the unrolled loops convergent. 29차: v2 on 32 blocks takes the pair's exposure under the routed MoE from 47.4 to 31.8 us a layer in serving's issue order (x43 = -0.67 ms/step), 8 blocks is too few (the pair outlives the MoE tail). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…, never a 48-block barrier on a 32-block launch), the global path's order restored, one A quantizer for the lane, the fewer-blocks control Review round 2 (10 angles): the lq kernel carried the whole barrier path and "a drift degrades to the global path" was a hang -- the launch is sized to the units, the barrier waits for c.grid arrivals; now if constexpr and __trap(). v2 had moved the first k-block's A store past the expansion on the GLOBAL path too, so "mk_gemm_kernel's binary is the measured one" was false; main's order is back. The three A quantizers (gemm prologue, local path, kda p4) share mk_warp_amax (redux.max on the bits), mk_pow2_scale (exponent arithmetic, exact by construction), mk_pow2_rcp and mk_pack4 (paired cvt + prmt): the same bytes by construction, the global kernel's SASS changes (3,360 -> 3,240 lines, 80 registers) and its control rows are re-measured, not assumed. Local path: x two k-blocks ahead in a two-array register ring (a runtime index went to local memory), the quant ahead of the mma, dead rows zeroed, the live-row bound in one spelling. The boot self-test and probe_exact share one fixture builder, check the plan really took the local kernel (otherwise the bitwise check compared the global kernel with itself), restore the knobs in a finally, and log the extension's plan. Launch grids are the bench's only (set_probe; the env surface is gone; the bracket snapshot lists what serving can carry). The concurrent probe builds the served MoE from the MoE probe's builders, projects by the model's 42 MoE layers (not 43) and labels it a projection, and carries the control row that separates "fewer blocks" from "no barrier": the global kernel on 32 blocks, on its own ticket counter, plus the 48-block local row of the first form. Sweep: rows deduplicated by plan, the control shape's forced splits skipped, marks by the tolerance judged, stamps cleared before the launch they read. Docs: the retracted ceiling removed from the RUNBOOK bullets, "wash"/"hidden" replaced by the measured numbers, the v1 -> v2 moe-first gap recorded as unexplained. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…mk-kernel-improvements-a8b6ea -- both lanes kept, my EXP-17 renumbered EXP-22 The other session's non-persistent v2 lane (mk_gemm2_kernel, VLLM_GLM53_MK_GEMM2, 30차) and this branch's local-quant kernel (mk_gemm_lq_kernel, VLLM_GLM53_MK_LOCALQ, 29차) share the diagnosis -- a resident barrier kernel serialises with the MoE kernel on the aux stream -- and differ in prescription. Both stay, both default off: mk_run_gemm dispatches the v2 lane first, then the plan; the bench keeps the --gemm2 columns beside the local column and --gemm-sweep; probe_exact runs the v1 pair of kernels on the shared fixture and the v2 lanes after it; the ledger orders 30차 above 29차; the v2 kernel's scalar e4m3 conversion is back for its own pack. test_logic 44,646 OK; nvcc gemm 80 / lq 122 / gemm2 124 registers. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: serverGenReqId_72e1ebb3-4bad-45dc-8a5b-aeb7f62a1e4e) |
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
The new set_probe pybind API only enforces upper bounds (silently accepting out-of-contract negative values), and should validate documented lower bounds across the source and generated CUDA copies.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR introduces EXP-22 for GLM53 MK-GEMM: an opt-in barrier-free local-quant kernel path intended to reduce aux-stream exposure for the shared expert’s small GEMM pair when co-scheduled with routed MoE. The change is gated behind a new profile-declared knob (VLLM_GLM53_MK_LOCALQ, default off) and reinforced with an exact dual-kernel self-test plus new/updated probes and documentation.
Changes:
- Add a new
mk_gemm_lq_kernel(compile-timemk_gemm_phase_t<true>) and host planning (mk_gemm_plan_for) to select global vs local-quant kernel based on units/grid andbgmarking. - Plumb a
bgflag from fp8-dense call sites into the GEMM launcher, and add bench/probe hooks (gemm_plan,set_probe) plus a new concurrent probe. - Update docs/runbook/measurements and extend
tests/test_logic.pyassertions to cover the new path, helper refactor, and knob wiring.
File summaries
| File | Description |
|---|---|
| tests/test_logic.py | Updates invariant checks for bg/localq plumbing, plan selection, helper refactor, and probe/tooling assertions. |
| RUNBOOK_KERNEL_CAMPAIGN2.md | Documents EXP-22 rationale, knob semantics, gating steps, and measurement/probe commands. |
| profiles/glm53.env | Declares VLLM_GLM53_MK_LOCALQ=0 so launchers forward the knob (default off). |
| probes/moe_decode_stream_probe.py | Refactors served wrapper + weight-set construction into reusable helpers for other probes. |
| probes/mk_gemm_concurrent_probe.py | Adds a new probe to measure shared-expert GEMM pair exposure under routed MoE with forked streams. |
| probes/megakernel_glm53_bench.py | Adds standalone localq-vs-global sweep (incl. stamps) and refactors exact probe fixture reuse. |
| overlay/modules/glm53_megakernel/README.md | Adds detailed README section for local-quant path design, selection rules, and probe interpretation. |
| overlay/modules/glm53_megakernel/glm53_megakernel.py | Adds bg plumbing, shared exact fixture utilities, and dual-kernel exact self-test logic. |
| overlay/modules/glm53_megakernel/glm53_megakernel.cu | Implements shared quant helpers, split gate refactor, local-quant kernel + plan selection + bench hooks. |
| overlay/modules/glm53_fp8_dense/glm53_fp8_dense.py | Marks shared-expert projections as background and passes bg into MK GEMM hook. |
| MEASUREMENTS.md | Records EXP-22 measurement results and review outcomes (standalone vs concurrent exposure). |
| build/glm53/glm53_megakernel.py | Mirrors source changes into built artifact. |
| build/glm53/glm53_megakernel.cu | Mirrors source CUDA changes into built artifact. |
| build/glm53/glm53_fp8_dense.py | Mirrors fp8-dense bg plumbing into built artifact. |
| build/dsv4/glm53_megakernel.py | Mirrors source changes into built artifact for dsv4 build. |
| build/dsv4/glm53_megakernel.cu | Mirrors source CUDA changes into built artifact for dsv4 build. |
| bench/bracket.py | Adds VLLM_GLM53_MK_LOCALQ to bracket env snapshot list. |
Review details
- Files reviewed: 17/17 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| TORCH_CHECK(ksr <= KBLK_MAX && localq <= 2 && lq_grid <= MK_GRID_CAP && | ||
| bg_grid <= MK_GRID_CAP, | ||
| "set_probe: a knob above its range"); |
| TORCH_CHECK(ksr <= KBLK_MAX && localq <= 2 && lq_grid <= MK_GRID_CAP && | ||
| bg_grid <= MK_GRID_CAP, | ||
| "set_probe: a knob above its range"); |
| TORCH_CHECK(ksr <= KBLK_MAX && localq <= 2 && lq_grid <= MK_GRID_CAP && | ||
| bg_grid <= MK_GRID_CAP, | ||
| "set_probe: a knob above its range"); |
…ernel-improvement-5a65a8 kernel: MKGemmCtx keeps both the SMLP pair-activation fields and bar_id; bench: probe_smlp beside the new probes, smlp segment/knob kept; tests: PDL kernel count 7 (gemm, gemm-lq, gemm2, kda, mhc, smlp, mla), ksr calls 7. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…31차로 PR #304 는 #301(672430e)에서 갈라져 나왔고, 그 뒤 main 이 #300·#302·#303·#305· #306·#307·#308(#309)로 앞서 나가면서 충돌 상태(dirty)가 됐다. 코드 파일은 겹치지 않았고(KDA 원패스·kpool 은 #304 단독), 겹친 것은 문서·프로필·테스트 넷이다. **RUNBOOK 충돌 2건** (1) 새 EXP 절의 삽입 위치: main 이 EXP-16 뒤에 EXP-22(로컬 양자화)·EXP-21(v2 레인)을 넣었고 #304 는 같은 자리에 EXP-20(미세 융합 묶음 2)을 넣었다. 경쟁이 아니라 가산이므로 셋 다 두되 기존 배열(최신 먼저)을 따라 22 → 21 → 20 순으로 둔다. (2) "순서와 근거" 절: main 은 평평한 번호 목록을 지우고 자기가 새로 만든 상태 표로 보내는 안내문으로 바꿨고(`채택·기각·닫힘은 위 상태 표에 있다`), #304 는 그 목록에 18번(EXP-20)을 덧붙였다. **main 의 재구성을 취하되**, 사라질 뻔한 #304 의 항목을 main 의 "다음 부팅 창에서" 목록에 8번으로 옮겨 실험이 목록과 함께 증발하지 않게 했다. **원장 번호 충돌 (자동 병합이 못 보는 것)** #304 는 자기 측정을 `★29차` 로 적었는데, 그 사이 main 에 다른 29차(메가커널 로컬 양자화, #307)와 30차(비상주 v2 레인, #305)가 먼저 들어왔다. 텍스트로는 충돌하지 않아 병합 결과에 **서로 다른 29차 둘**이 남았다. 28차 항목이 세운 선례("원장 번호 27 은 PR #290 이 쓰고 있어 28 로 적는다")대로 다음 빈 번호인 **31차** 로 옮겼고, 왜 옮겼는지 절 머리에 한 줄 남겼다. 측정값 자체는 손대지 않았다. 같은 번호를 가리키던 참조 6곳을 함께 고쳤다: MEASUREMENTS 절 제목, RUNBOOK EXP-20 표의 기각 축, `glm53_kda_onepass/README.md`, `moe_gate_sm121/README.md`, `profiles/glm53.env`, `tests/test_logic.py`. main 이 자기 29차를 가리키는 참조 (`VLLM_GLM53_MK_LOCALQ` 주석)는 그대로 두었다. **자동 병합분 검증** - `profiles/glm53.env`: main 의 `VLLM_GLM53_MK_LOCALQ` 와 #304 의 세 노브 (`KDA_DUAL_GEMM`·`KDA_ONEPASS`·`KPOOL_UPDATE_DIRECT_POS`) 전부 기본 0 으로 잔존. `MODULES=` 목록은 main 의 항목 하나도 잃지 않고 `glm53_kda_onepass` 만 늘었다. - overlay ↔ build 사본 3쌍(`glm53_kda_onepass.py`, `glm5next_kda.py`, `sparse_attn_indexer_kpool.py`) 모두 동일. - `tests/test_logic.py`: 양쪽 테스트 공존, `micro-fusion bundle 2 contracts` 통과. 검증: `tests/test_logic.py` all OK (6106 checks; 이 호스트엔 torch 부재로 일부 SKIP). 충돌 마커 0, 변경된 .py 전부 파싱 OK, `run_micro_fusion_check.sh` bash -n OK. GPU 검증(프로브 VERDICT·트레이스 물리확인)은 하지 않았다 — 이 환경에 GPU 가 없다. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JTghtTjyfT23vuhZwHzBSM
**1. 원장 번호 충돌 — #311 의 29차를 32차로** main 에 서로 다른 29차 두 개가 남아 있었다: `메가커널 29차`(로컬 양자화 커널, PR #307, 5d52c16 으로 먼저 안착)와 `29차`(레버 2~7 브래킷 체인, PR #311, 65b2362 로 나중). 두 계열은 한 수열이다 — 28차 항목이 "원장 번호 27 은 PR #290 이 쓰고 있어 28 로 적는다"고 적으며 메가커널 27차를 피해 간 것이 그 증거다. 그 선례대로 **나중에 들어온 #311 쪽을 다음 빈 번호 32차로** 옮겼다(30·31 사용 중). 절 제목과 본문의 `item 2/5/6` 참조를 포함해 50곳을 고쳤고, 어느 쪽을 가리키는지 한 줄씩 분류해 **메가커널 29차를 가리키는 13곳은 그대로 두었다**(RUNBOOK EXP-22 5곳, `glm53.env` 2곳, megakernel README·py 4곳, 절 제목, #304 재번호 주석). 왜 옮겼는지는 31차와 같은 형식으로 절 머리에 남겼다. 측정값은 손대지 않았다. **2. RUNBOOK 상태 표 — 빠진 세 줄과 낡은 한 줄** 표가 EXP-20·21·22 를 아예 담고 있지 않았다. 각 절과 원장에 적힌 값만으로 채웠다: - EXP-20(미세 융합 묶음 2) — 브래킷 대기, 노브 3개 기본 0, 오프라인 게이트 PASS, 런치 −249/1,548 · −0.25 ms/스텝(C=1, 31차) - EXP-21(MK-GEMM v2) — 프로브 먼저, `MK_GEMM2` 프로필 미선언, exact PASS 이후 노출 프로브·스탬프 대기(30차 §4~5) - EXP-22(로컬 양자화) — 프로브 먼저 · 단독 부팅 금지, 노출 47.4 → 31.8 µs/층 (투영 −0.66 ms/스텝, 메가커널 29차) 그리고 **EXP-7 행이 낡아 있었다**: #311 이 `PREP_FUSED` 를 기본값 1 로 올리고 16.39 → 17.59 step/s(+7.3%)를 실측했는데 표는 아직 "부팅 대기 · `PREP_FUSED=0`" 이었다. 프로필 실값으로 확인해 "채택 · 기본값" 으로 고쳤다. 이제 표는 EXP 1~22 를 빠짐없이, 중복 없이 담는다. 검증: `tests/test_logic.py` all OK (6167 checks), 변경된 .py 파싱 OK, `bash -n` OK, overlay ↔ build 사본 동기 유지. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JTghtTjyfT23vuhZwHzBSM
MK GEMM: the shared expert's pair on a barrier-free local-quant kernel (EXP-22, knob off) + two review rounds
Target. The decode step's 30-45 us class (28차 trace reading): the shared expert's gate_up
[1024 x 4096]and down[4096 x 512], two launches in each of the 42 MoE layers, 3.5 ms, 2.4 ms of it not covered by another stream. They move 1-2 MB each (5-12 us at the lane's rate). Serving forks them onto the aux stream beside the routed MoE kernel, and a persistent 48-block barrier kernel there gets its SMs as MoE blocks retire and holds every resident block for the last one -- the same diagnosis 30차 (#305) reached from the armed trace.Change.
mk_gemm_lq_kernel(mk_gemm_phase_t<true>, a compile-time instantiation with none of the barrier path): every block quantizes the A k-blocks of its unit straight into smem from x in L2 (loaded two k-blocks ahead in a register ring, reduced ahead of the mma), launched on as many blocks as it has units, no grid barrier; a host/kernel drift on the unit rule traps. The lane's three A quantizers (gemm prologue, this path, KDA p4) share one set of helpers (mk_warp_amax= redux.max,mk_pow2_scale= exponent arithmetic,mk_pow2_rcp,mk_pack4= paired cvt + prmt): the same bytes by construction, and the boot self-test runs the exact e2m1 fixture through BOTH kernels, checks the plan really took the local one, and requires bitwise equality.mk_gemm_kernelstays at 80 registers (its prologue's quantizer changed with the helpers, so its control rows are re-measured, not assumed).Knob.
VLLM_GLM53_MK_LOCALQ= 0 (default, declared inprofiles/glm53.envso the launcher forwards it) / 1 = the launches the fp8-dense hook marks background (mlp.shared_experts.*) / 2 = every one-unit-per-block launch (bench). Launch grids and the fewer-blocks control are the bench's only (set_probe). Serving does not change on merge.Measured (29차, srv2). Standalone the local kernel is slower than the global one (the prologue it skips is 5.5 us on the stamps; first form +4 us at m=8 / +12 at m=32, the 32-block form +8 us on the pair; the ring/row-bound form of HEAD is unmeasured). Under the routed MoE kernel (
probes/mk_gemm_concurrent_probe.py, one graph with the aux fork, U=40): the global kernel's pair is exposed whole in either issue order (47.4 us a layer); the local kernel on 32 blocks 36.2 (MoE first) / 31.8 (pair first, serving's order) -- x42 layers a projection of about -0.66 ms/step, not a step number. Pending in the next idle window: the control row (the global kernel on 32 blocks, own ticket counter) that separates "fewer blocks" from "no barrier", the 48-block local row under the same method, and the row-bound form's standalone delta. This branch and the v2 lane of #305 are two prescriptions for one diagnosis; they are measured side by side, then the EXP-6 bracket decides.Reviews. Two rounds (10 angles each, verified): default-on without profile wiring, a vacuous both-kernel gate, a drift that would have hung the aux-stream launch on a 48-block barrier, the global path's order silently changed by v2, stale/overstated ceilings, sweep/stamp/probe bugs -- all fixed; findings and outcomes are in the session's review reports and MEASUREMENTS.md 29차.
Merge note. origin/main's MK-GEMM v2 lane (#305) landed in the same files; both lanes are kept, both default off, and this branch's EXP number is 22 (17 is #303's TileLang passes). test_logic 44,646 checks OK; nvcc gemm 80 / lq 122 / gemm2 124 registers, no spills.
🤖 Generated with Claude Code
Note
Medium Risk
Touches hot-path CUDA GEMM launch selection and changes global-kernel SASS via shared quant helpers; mitigated by dual-kernel exact self-test and default-off knob, but wrong
bg/plan wiring could hang aux-stream launches or alter numerics.Overview
Adds EXP-22: a separate
mk_gemm_lq_kernelthat quantizes activations per-block from L2 (no grid-wideg_mk_aq+ barrier), launched with one block per unit whenVLLM_GLM53_MK_LOCALQ=1and the caller marks the launch background (bg).Serving / Python:
gemm_w4a8andrun_gemmtakebg; fp8-dense sets_mk_bgonmlp.shared_experts.*projections. Hostmk_gemm_plan_forpicks global vs local kernel, grid size, and optional bench control (set_probe/gemm_plan).Kernel refactor:
mk_gemm_phase_t<LQ>shares quant helpers (mk_warp_amax,mk_pow2_scale,mk_pack4) across prologue, local path, and KDA; globalmk_gemm_kernelbinary changes slightly but stays a separate 80-register entry point.Validation: Boot self-test runs
exact_fixturethrough both kernels and requires bitwise match;VLLM_GLM53_MK_LOCALQadded to bracket env keys. Docs/runbook record 29차 benches (standalone slower; MoE concurrent probe projects ~−0.66 ms/step at pair-first, knob default off).Reviewed by Cursor Bugbot for commit 0f8978e. Bugbot is set up for automated code reviews on this repo. Configure here.