perf(modules): b12x MoE wrapper 를 geometry 별로 공유 — 층마다 541 MiB 를 중복으로 들고 있었다

[choiceoh]

Summary

층마다 B12xMoEWrapper 를 새로 만듭니다. 각 wrapper 는 geometry 로 크기가 정해지는 graph-stable 스크래치를 듭니다. 이 배치에서 실측:

wrapper 1개 : allocated +541.1 MiB
MoE 층 43개 : 22.72 GiB (랭크마다)
공유 시 절감: 22.19 GiB

GLM-5.3-Flash 의 43 개 MoE 층은 geometry 가 전부 같습니다 — 즉 GMU 예산의 4분의 1이 중복 버퍼였습니다.

엔진이 보고한 수치와 맞습니다

가중치 45.8 GiB + wrapper 22.7 GiB = 68.5 GiB
GMU 0.73 × 119.69          = 87.4 GiB 예산
→ KV ≈ 16 GiB                (엔진 로그: 16.52 GiB)

오늘 메모리로 겪은 것들 — 워커 OOM, GMU 를 올려도 내려도 죽던 것, KV 가 16 GiB 를 못 넘던 것 — 의 상당 부분이 여기서 나옵니다.

안전성

MoE 층은 같은 스트림에서 차례로 돕니다. wrapper 는 호출마다 자기 버퍼에 쓰고 호출 밖으로 남기는 것이 없습니다. 값은 약참조라 마지막 층이 놓으면 같이 사라집니다.

출처

vllm-project/vllm#48698 과 같은 발상을 이 이미지가 싣는 파일에 맞춰 쓴 것입니다. (#53081 은 wrapper 대신 workspace 를 공유하지만 FlashInfer 의 shared_static_workspace(flashinfer-ai/flashinfer#4603)가 필요하고, 이 빌드 0.6.18.dev20260819 에는 없습니다.)

Verification

• tests/test_logic.py — all OK (666 checks)
• compose-overlays.sh glm53 — 13 overlays from 11 modules
• base preimage sha 고정 (86dfe87e…)
• 부팅 실측은 다음 단계 — KV 가 16 → 38 GiB 로 늘어야 하고, 늘지 않으면 이 수치 해석이 틀린 것입니다

🤖 Generated with Claude Code

---

Note

Medium Risk
Shared mutable GPU scratch assumes strictly sequential MoE execution on one stream; wrong overlap or future parallelism could cause cross-layer corruption. The patch calls logger.info_once without an obvious import in the shown file, which could fail on first wrapper creation if not provided by the base module.

Overview
Adds a b12x_shared_workspace overlay that patches flashinfer_b12x_moe.py so B12xMoEWrapper is created once per MoE geometry (experts, top_k, hidden/intermediate sizes, activation/SwiGLU params) instead of once per layer. Layers with identical geometry—e.g. 43 GLM-5.3-Flash MoE blocks—reuse the same wrapper via a thread-locked WeakValueDictionary, with each layer still holding a strong ref in _ensure_wrapper.

The intent is to drop ~22 GiB per rank of duplicate graph-stable scratch (~541 MiB per wrapper) so more GMU budget is available for KV cache. Wiring: new overlay manifest.tsv, duplicate under build/glm53/, and b12x_shared_workspace prepended to MODULES in profiles/glm53.env.

^{Reviewed by Cursor Bugbot for commit 772e0ba. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[cursor]

Bugbot couldn't run - usage limit reached

Bugbot 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_a92f1f9e-805f-4873-be86-685517ec316d)

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-08-30T21:23:38.833929Z	f117dea	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[cursor]

Bugbot couldn't run - usage limit reached

Bugbot 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_111d8cd4-16cb-4cc8-94b2-cc94072cf13d)

[cursor]

Stale comment

Not approved: Cursor Bugbot was present but skipped (usage limit), so the required automated-review signal did not complete successfully. Human review is needed; a reviewer was assigned.

  

_{Sent by Cursor Approval Agent: Pull Request Router and Approver}

[cursor]

Not approved: Cursor Bugbot was present but skipped (usage limit), so the required automated-review signal did not complete successfully. Human review is needed; a reviewer was requested.

  

_{Sent by Cursor Approval Agent: Pull Request Router and Approver}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f117dead1e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Define the logger before creating the shared wrapper

With the glm53 profile enabled, the first MoE apply() always takes this cache-miss path; after allocating and caching the wrapper, it calls logger.info_once, but this module never defines or imports logger. That raises NameError before _shared_wrapper() can return, so the first model execution fails instead of realizing the memory-saving shared wrapper.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Regenerate the committed glm53 manifest

When a consumer uses the committed build/glm53 output without rerunning the composer, this newly enabled module is not applied: build/glm53/manifest.tsv still contains only the previous 12 rows and has no entry for flashinfer_b12x_moe.py. Deployment and launcher mounting enumerate files exclusively from that manifest, so the newly committed build file is ignored; regenerate the build manifest so the versioned composed output matches this MODULES list.

Useful? React with 👍 / 👎.