From 9c9861efd80b4b1fbeb1ff7b4edc3c02660d9257 Mon Sep 17 00:00:00 2001 From: choiceoh Date: Fri, 4 Sep 2026 16:59:53 +0900 Subject: [PATCH 1/3] =?UTF-8?q?feat(osar,megakernel):=20AR=20=EB=8C=80?= =?UTF-8?q?=EA=B8=B0=EB=A5=BC=20=EB=8B=A4=EC=9D=8C=20=EC=BB=A4=EB=84=90?= =?UTF-8?q?=EC=9D=98=20L2=20=ED=94=84=EB=A6=AC=ED=8E=98=EC=B9=98=EB=A1=9C,?= =?UTF-8?q?=20=EC=84=9C=EB=B9=99=20PDL=20=EA=B8=B0=EB=B3=B8=EA=B0=92,=20Mo?= =?UTF-8?q?E=20go/no-go=20=ED=94=84=EB=A1=9C=EB=B8=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - tp_oneshot_ar: k_oneshot 이 HintArgs(최대 8 범위)를 받아 워프 1~7 이 prefetch.global.L2 로 걷는다 -- 스레드 0 의 폴링과 t_wait 는 그대로, n==0 은 옛 커널 그대로. oneshot_ar_hint / phase_counters 바인딩. shim 은 MK 발사의 note_consumer 를 "타깃 forward 의 몇 번째 콜렉티브 뒤" 로 파일해 두고 캡처가 발사에 굳힌다; forward 경계는 컴파일 영역 위의 Glm5NextForConditionalGeneration.forward. 노브 VLLM_GLM53_AR_PREFETCH (0 / 1 = 12 MB / N MB, 1..20). - 서빙 PDL: 드라이버는 env 를 읽는데 프로필에도 ab-glm53.sh 에도 없어서 지금까지 무장 부팅이 전부 PDL 없이 돌았다 -> glm53.env 기본 1, cand 팔에 명시. 그래프 캡처 안 체인 검사 probes/mk_pdl_graph_check.py. - probes: run_mk_probe.sh (b12x 파일까지 마운트하는 일반 러너), moe_decode_stream_probe.py (b12x static U 스윕 vs MK 레인 vs L2-warm), osar_build_check.py, oneshot_ar_disttest.py 힌트 팔. - RUNBOOK EXP-12/13/14. test_logic 44,471 checks. GPU 프로브는 srv2 가 비는 대로 돈다(다른 세션의 DRAFTW4 브래킷 뒤); 숫자는 원장에. Co-Authored-By: Claude Fable 5.1 --- RUNBOOK_KERNEL_CAMPAIGN2.md | 60 +++++ bench/bracket.py | 1 + build/dsv4/dsv4_oneshot_ar.cu | 93 +++++++- build/dsv4/dsv4_oneshot_shim.py | 114 +++++++++- build/dsv4/glm53_megakernel.py | 28 +++ build/glm53/dsv4_oneshot_ar.cu | 93 +++++++- build/glm53/dsv4_oneshot_shim.py | 114 +++++++++- build/glm53/glm53_megakernel.py | 28 +++ build/glm53/glm5next_model.py | 34 ++- launchers/ab-glm53.sh | 5 +- overlay/modules/glm53_megakernel/README.md | 7 +- .../glm53_megakernel/glm53_megakernel.py | 28 +++ .../glm53_model_wiring/glm5next_model.py | 34 ++- overlay/modules/tp_oneshot_ar/README.md | 31 +++ .../modules/tp_oneshot_ar/dsv4_oneshot_ar.cu | 93 +++++++- .../tp_oneshot_ar/dsv4_oneshot_shim.py | 114 +++++++++- probes/mk_pdl_graph_check.py | 184 +++++++++++++++ probes/moe_decode_stream_probe.py | 215 ++++++++++++++++++ probes/oneshot_ar_disttest.py | 40 ++++ probes/osar_build_check.py | 40 ++++ probes/run_mk_probe.sh | 76 +++++++ profiles/glm53.env | 18 ++ tests/test_logic.py | 152 +++++++++++++ 23 files changed, 1585 insertions(+), 17 deletions(-) create mode 100644 probes/mk_pdl_graph_check.py create mode 100644 probes/moe_decode_stream_probe.py create mode 100644 probes/osar_build_check.py create mode 100644 probes/run_mk_probe.sh diff --git a/RUNBOOK_KERNEL_CAMPAIGN2.md b/RUNBOOK_KERNEL_CAMPAIGN2.md index 2c814c48..ac1370ed 100644 --- a/RUNBOOK_KERNEL_CAMPAIGN2.md +++ b/RUNBOOK_KERNEL_CAMPAIGN2.md @@ -474,6 +474,63 @@ MLA(rope 64 · topk 512 · 압축기 · 슬라이딩 윈도)·GEMM(dense 가 블 프로브의 `--stock dispatch` 가 부팅이 실제로 타는 팔을 재고 `hit` 열로 그걸 말한다. +## EXP-12 — 서빙 PDL (`VLLM_GLM53_MK_PDL=1`, 2026-09-04 추가) + +메가커널 발사는 PDL(programmatic dependent launch)로 튜닝돼 있다 — 다음 MK 커널이 +앞 커널이 비운 SM 에서 시작해 꼬리 동안 첫 W 타일을 당긴다(연속 2발 발사당 +−17~19%, 2차). 그런데 드라이버는 env 를 읽고 **프로브만 그것을 켰다**: 프로필에도 +`ab-glm53.sh` 의 cand 팔에도 없었으므로 지금까지의 무장 부팅은 전부 PDL 없이 +돌았다(173발/스텝). EXP-6 종단 무효과의 용의자 1번. + +- 수치 불변: 모든 MK 커널이 앞 커널 출력을 읽기 전에 `griddepcontrol.wait` 하고, + wait 앞의 채움은 가중치뿐(mk_gemm_phase 의 hoist). 그래프 캡처 안 체인 형태는 + `probes/mk_pdl_graph_check.py` 가 판정한다(gemm→gemm→gemm, mhc→gemm 리플레이 + 비트 동일 + PDL on/off 발사당 µs). +- 배선: `profiles/glm53.env` 기본 1, `ab-glm53.sh` cand 팔에 명시. 세그먼트가 + 하나도 무장되지 않으면 무효(inert). +- 상한: 173발 × 5~10 µs = −0.9~1.7 ms/스텝. 단독 부팅 금지 — EXP-6 브래킷의 cand + 팔에 얹는다(수치 축 하나 규칙: PDL 은 속도만 바꾼다). + +```bash +bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py # PDL=1 +VLLM_GLM53_MK_PDL=0 bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py # 대조 +``` + +## EXP-13 — AR 대기 중 다음 커널 가중치 L2 프리페치 (`VLLM_GLM53_AR_PREFETCH`, 2026-09-04 추가) + +`k_oneshot` 의 대기(회당 38.7/45.5 µs, 스텝당 ~100회)는 모든 랭크에서 DRAM 이 +노는 시간이다. 커널이 `HintArgs`(최대 8개 (포인터, 바이트))를 받아 워프 1~7 이 +`prefetch.global.L2` 로 걷고 스레드 0 은 그대로 플래그를 돈다. 힌트는 **학습**: +MK 드라이버의 발사(`_gemm_call`·`_mhc_call`·`_kda_launch`)가 읽을 가중치를 +`note_consumer` 로 알리고, shim 이 "타깃 forward 의 몇 번째 콜렉티브 뒤인가" 로 +파일해 두었다가 캡처 시 발사에 굳힌다. forward 경계는 컴파일 영역 위의 +`Glm5NextForConditionalGeneration.forward`. MK-GEMM 이 무장돼야 배울 것이 있다. + +- 게이트(순서대로): (1) `probes/osar_build_check.py` — 새 명령의 ptxas 통과(실패면 + 부팅이 NCCL 로 조용히 떨어진다), (2) `probes/oneshot_ar_disttest.py` 4랭크 — + 12 MB 힌트/무힌트의 `t_wait` 와 maxerr 0(힌트가 NIC 쓰기를 밀어 대기를 늘리면 + 손해), (3) `moe_decode_stream_probe.py` 의 `gemm cold` vs `gemm L2-warm` 행 — + 소비자 쪽 이득의 단위, (4) 플릿 브래킷: EXP-6+12 위에 `VLLM_GLM53_AR_PREFETCH=1` + (caller env). 수치 불변 → step/s 만. 부팅 로그에 `[osar] prefetch hints learned: + N collectives, X MB` 가 없으면 무장이 아니다. +- 상한: 임계 랭크의 대기 ≈ 전송 ~20 µs = 4.6 MB → −1.5~2.5 ms/스텝(전략 문서). +- 예산 노브: 1 = 12 MB/콜렉티브, N = N MB(1..20; L2 24 MB). + +## EXP-14 — MK_SEG_MOE go/no-go (2026-09-04 추가, 프로브만) + +21차의 "MoE 는 대역폭 바닥" 은 190 GB/s 를 **추정한 고유 전문가 수 ~40** 에서 +역산한 값이라 순환이다. `probes/moe_decode_stream_probe.py` 가 서빙이 만드는 +b12x 래퍼(같은 기하, 같은 디스패치 오버레이; C=1 은 64 pairs 라 static 백엔드)를 +디코드 형상(8토큰·top-8)에서 고유 전문가 U=8..64 별로, 가중치 DRAM-cold(8 세트 +순환), 그래프 리플레이로 재고 같은 바이트를 MK W4 레인이 스트리밍하는 속도(팩 +12개 연속, PDL)와 견준다. + +- 판정 규칙: b12x 가 레인의 90% 이상이면 축을 닫는다(원장에 기록). 아래면 + (레인 − b12x) 비율 × 31 ms 가 세그먼트의 상한이고, 설계는 전략 문서 4장(48블록 + persistent, FC1 (전문가, n타일) 유닛 → 전문가별 완료 카운터로 열리는 FC2 동적 + 큐, 공유 전문가 = 41번째 전문가, b12x nvfp4 레이아웃 제자리 읽기, A4→A8). +- 실제 서빙 U 는 다음 부팅에서 로그 한 줄로 확정한다(프로브는 U 별 곡선만 준다). + ## 브래킷 자동화 — `bench/bracket.py` (도구, 판정 아님) `leg`(살아있는 서버에 rep 기록) + `judge`(기록 판정) 2중 명령. 원장 규율을 코드로 @@ -497,6 +554,9 @@ MLA(rope 64 · topk 512 · 압축기 · 슬라이딩 윈도)·GEMM(dense 가 블 9. **EXP-9 (head-gate split-K)** — 단독 부팅 금지, EXP-7 부팅에 얹는다. 10. **EXP-7 이 붙은 뒤 드래프터 D 를 다시 잰다** — 9월 1일 트레이스에서 드래프터 ~4.3 ms 는 다음 스텝의 호스트 준비 유휴 뒤에 숨어 있었다(그래서 D≈0). 은신처가 사라지면 임계경로에 올라온다(천장 ~6%; fc GEMM 809 us 는 K=20480 직렬 스케줄이라 split-K 후보). #104 를 지금 재론하는 것이 아니라 조건이 바뀐 뒤의 재측정이다. 오프라인으로 닫을 >1% 레버는 더 없다(STEP_KERNEL_MAP 보충 분해 3). 11. **EXP-10 (드래프터 GEMM → MK W4)** — 프로브가 돌았다(2026-09-04): 스텝 합 bf16 3.23 → MK W4 1.25 ms(−3.0%), 배선 랜딩(`VLLM_DFLASH2_FP8_DENSE=1`, 불투명 op, fc K-chunk). 남은 것은 브래킷 부팅 하나(수용률 ±2 pct 게이트, 프리필 동반). 10번의 "EXP-7 뒤 재측정" 조건은 무장 트레이스가 이미 답했다: 꼬리는 지금 임계경로다. +12. **EXP-12 (서빙 PDL)** — 프로브(그래프 체인) 뒤 EXP-6 브래킷의 cand 팔에 얹는다. 단독 부팅 없음. +13. **EXP-13 (AR 프리페치)** — 컴파일 → 4랭크 disttest → 브래킷(EXP-6+12 위). 수치 불변. +14. **EXP-14 (MK_SEG_MOE go/no-go)** — 프로브 하나가 착수 여부를 정한다. 90% 규칙. 12. **EXP-11 (dsv4 에 MK_SEG_MHC)** — 2단계는 **부팅 없음**: srv2 에서 서빙 컨테이너가 비었을 때(`docker ps`) `bash probes/run_megakernel_bench.sh --profile dsv4 --iters 20`. diff --git a/bench/bracket.py b/bench/bracket.py index 411522a6..d4a89510 100644 --- a/bench/bracket.py +++ b/bench/bracket.py @@ -39,6 +39,7 @@ "VLLM_GLM53_PREP_FUSED", "VLLM_GLM53_ASYNC_DFLASH", "VLLM_GLM53_MHC_SMALLM", "VLLM_DFLASH_PREP_WARMUP", "VLLM_GLM53_MK_PDL", "VLLM_GLM53_MK_KSR_OUT", + "VLLM_GLM53_AR_PREFETCH", ) if k in os.environ] + sorted(k for k in os.environ if k.startswith(("VLLM_GLM53_MK_", "VLLM_GLM53_KPOOL"))) diff --git a/build/dsv4/dsv4_oneshot_ar.cu b/build/dsv4/dsv4_oneshot_ar.cu index 6f6e6c19..38803e9d 100644 --- a/build/dsv4/dsv4_oneshot_ar.cu +++ b/build/dsv4/dsv4_oneshot_ar.cu @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -146,8 +147,48 @@ __device__ __forceinline__ void osar_backoff(int &n, unsigned &ns) { #endif } +// L2 prefetch hints for the peer-wait window (VLLM_GLM53_AR_PREFETCH): byte +// ranges of the weights the NEXT kernel after this collective streams. The +// wait is 20-40 us of idle DRAM on every rank (MEASUREMENTS 19차: wait 38.7 +// of 45.5 us per collective, ~100 per step); warps 1..7 of every block walk +// these ranges with prefetch.global.L2 while thread 0 polls the peer flags, +// so the consumer finds its first megabytes in L2 (24 MB on this part). +// Passed by value: a CUDA-graph capture bakes the hint with the launch, and +// the shim learns the ranges from the consumers that follow each collective +// during the eager warmups that precede capture. n == 0 is exactly the old +// kernel -- no branch of it touches memory. +#define OSAR_MAXHINT 8 +struct HintArgs { + unsigned long long ptr[OSAR_MAXHINT]; + unsigned int len[OSAR_MAXHINT]; + int n; +}; + +__device__ __forceinline__ void osar_prefetch(const HintArgs &h, + volatile int *landed) { + // One 32 B sector index space over the concatenated ranges, interleaved + // across the grid: block b, thread t takes sectors b*224 + (t-32) + k*10752, + // so every block warms a uniform slice of every range. A prefix walk warms + // a few blocks' tiles of the consumer and leaves its slowest block cold + // (MEASUREMENTS 11차). Owning blocks stop as soon as the peers landed; the + // work is bounded either way (budget <= 20 MB is ~2 us of issue per block). + const int tid = (int)threadIdx.x - 32; + const int stride = ARGRID * (ARTHREADS - 32); + int idx = (int)blockIdx.x * (ARTHREADS - 32) + tid; + for (int r = 0; r < h.n; ++r) { + const unsigned long long base = h.ptr[r]; + const int nsec = (int)((h.len[r] + 31u) >> 5); + for (; idx < nsec; idx += stride) { + asm volatile("prefetch.global.L2 [%0];" ::"l"( + base + ((unsigned long long)idx << 5))); + if (*landed) return; + } + idx -= nsec; + } +} + __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, - int nbytes) { + int nbytes, const HintArgs h) { // The grid is fixed at ARGRID for the counter invariant, so at decode sizes // the smallest plain call has n = hidden and many blocks fall entirely past // the payload: with blockDim 256 and n 4096, blocks 16..47 copy nothing and @@ -212,6 +253,16 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, // Peer wait: rxf is only ever written by the peers' NICs, never by a block // of this kernel -- same independence argument as the guard above. Fence // stays where it always was: after the wait, before reading peer data. + // The prefetch hints ride the wait: thread 0 polls, warps 1..7 warm L2 + // with the next kernel's weights until the peers land (s_landed) or their + // slice is done. A non-owning block has no flag to wait for and simply + // issues its slice; warp 0's other lanes go straight to the barrier. The + // phase timer below still brackets thread 0's poll alone, so t_wait keeps + // measuring the collective -- and shows any DRAM contention the prefetch + // puts on the NIC's writes. + __shared__ volatile int s_landed; + if (threadIdx.x == 0) s_landed = 0; + __syncthreads(); long long t2 = timer ? clock64() : 0; if (owns && threadIdx.x == 0) { // The old form re-read every peer's flag on every pass, including peers @@ -233,6 +284,9 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, } } } + s_landed = 1; + } else if (h.n > 0 && threadIdx.x >= 32) { + osar_prefetch(h, &s_landed); } long long t3 = timer ? clock64() : 0; __syncthreads(); @@ -548,24 +602,55 @@ static void py_connect(std::vector all) { pthread_create(&g_proxy, nullptr, proxy_fn, nullptr); g_started = true; } -static torch::Tensor py_oneshot(torch::Tensor input) { +static torch::Tensor py_oneshot_impl(torch::Tensor input, + const std::vector &ptrs, + const std::vector &lens) { TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kBFloat16); TORCH_CHECK(input.is_contiguous()); int64_t n = input.numel(); TORCH_CHECK(n <= MAXEL, "oneshot: tensor too large"); + TORCH_CHECK(ptrs.size() == lens.size() && ptrs.size() <= OSAR_MAXHINT, + "oneshot: hint lists must pair up, at most OSAR_MAXHINT"); auto out = torch::empty_like(input); const bf16 *src = reinterpret_cast(input.data_ptr()); bf16 *dst = reinterpret_cast(out.data_ptr()); cudaStream_t st = c10::cuda::getCurrentCUDAStream(); + // Prefetch hints: (device pointer, bytes) pairs the shim learned for this + // collective's ordinal. Baked into the launch, so a captured graph replays + // them without any host code. + HintArgs h; + h.n = 0; + for (size_t i = 0; i < ptrs.size(); ++i) { + if (ptrs[i] == 0 || lens[i] <= 0) continue; + h.ptr[h.n] = (unsigned long long)ptrs[i]; + h.len[h.n] = (unsigned int)std::min(lens[i], 0x7fffffff); + ++h.n; + } // One launch, and the grid is FIXED at ARGRID however small n is: the // last-block detection in k_oneshot is (done_ctr % ARGRID == ARGRID-1), // which is only sound if every launch contributes exactly ARGRID // increments. The 48-block grid fills GB10 once and covers MAXEL through // the kernel's grid-stride loops; empty decode blocks only sync/increment. k_oneshot<<>>(g_ctrl, src, dst, (int)n, - (int)(n * 2)); + (int)(n * 2), h); return out; } +static torch::Tensor py_oneshot(torch::Tensor input) { + return py_oneshot_impl(input, {}, {}); +} +static torch::Tensor py_oneshot_hint(torch::Tensor input, + std::vector ptrs, + std::vector lens) { + return py_oneshot_impl(input, ptrs, lens); +} +// The phase counters (SM cycles, monotonic) for a probe that wants the wait +// per collective with and without hints: [guard, copy, wait, reduce, calls]. +static std::vector py_phase_counters() { + if (!g_ctrl) return {}; + return {(int64_t)g_ctrl->t_guard, (int64_t)g_ctrl->t_copy, + (int64_t)g_ctrl->t_wait, (int64_t)g_ctrl->t_reduce, + (int64_t)g_ctrl->t_calls}; +} static bool py_healthy() { if (!g_started) return false; static uint64_t last = 0; @@ -586,6 +671,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("local_infos", &py_local_infos); m.def("connect", &py_connect); m.def("oneshot_ar", &py_oneshot); + m.def("oneshot_ar_hint", &py_oneshot_hint); + m.def("phase_counters", &py_phase_counters); m.def("healthy", &py_healthy); m.def("shutdown", &py_shutdown); } diff --git a/build/dsv4/dsv4_oneshot_shim.py b/build/dsv4/dsv4_oneshot_shim.py index a76992fb..15657fbd 100644 --- a/build/dsv4/dsv4_oneshot_shim.py +++ b/build/dsv4/dsv4_oneshot_shim.py @@ -374,13 +374,121 @@ def _eligible(t): ) +# ---- L2 prefetch hints for the peer wait (VLLM_GLM53_AR_PREFETCH) ---------- +# The collective's wait (38.7 of 45.5 us, MEASUREMENTS 19차, ~100 per decode +# step) is idle DRAM on every rank. The kernel can spend it pulling the NEXT +# kernel's weights into L2 (k_oneshot's HintArgs); what it should pull is +# learned here, not declared: every weight-streaming launch of the megakernel +# driver notes the tensors it reads, the note is attributed to the most +# recent collective of the current target forward, and the table so learned +# in the eager warmups is what the captured launches bake in. Nothing here +# runs inside a compiled region: begin/end come from the model class above +# the traced graph, and the notes come from custom-op bodies that execute +# eagerly during capture. +def _resolve_prefetch_budget() -> int: + """Bytes of the next consumer's weights one collective may warm. 0 = off + (the kernel is byte-identical to before), 1 = 12 MB, or an integer MB in + 1..20 -- L2 is 24 MB and the consumer's activations must fit beside it.""" + raw = (os.environ.get("VLLM_GLM53_AR_PREFETCH") or "0").strip().lower() + if raw in ("", "0", "false", "off", "no"): + return 0 + if raw == "1": + return 12 << 20 + try: + mb = int(raw) + except ValueError: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%r is not 0, 1 or an MB count; off", + raw, + ) + return 0 + if not 1 <= mb <= 20: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%d MB is outside 1..20; off", mb + ) + return 0 + return mb << 20 + + +_PREFETCH_BUDGET = _resolve_prefetch_budget() +_HINT_MAX = 8 # OSAR_MAXHINT in the .cu +_in_forward = False +_ordinal = 0 # collectives seen so far in the current target forward +_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes +_cand: dict = {} # the same, being learned in this forward +_cand_bytes: dict = {} + + +def begin_forward(): + """Called by the target model above the compiled region, once per forward + (eager warmup and capture alike). This forward's collectives use the + adopted table; the consumers noting themselves build the candidate.""" + global _in_forward, _ordinal, _cand, _cand_bytes + if not _PREFETCH_BUDGET: + return + _in_forward = True + _ordinal = 0 + _cand = {} + _cand_bytes = {} + + +def end_forward(): + global _in_forward, _table + if not _PREFETCH_BUDGET: + return + _in_forward = False + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + # a decode-shaped forward saw more consumers than the last one did + # (a prefill forward routes M > 32 to stock and notes nothing) + _table = _cand + logger.warning( + "[osar] prefetch hints learned: %d collectives, %.1f MB in total", + len(_table), + sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + ) + + +def note_consumer(tensors): + """A weight-streaming launch reports the tensors it is about to read. + Attributed to the most recent collective of this forward, up to the + budget; the first notes after a collective are the ones that count.""" + if not _PREFETCH_BUDGET or not _in_forward or _ordinal == 0: + return + lst = _cand.setdefault(_ordinal, []) + used = _cand_bytes.get(_ordinal, 0) + for t in tensors: + if t is None or len(lst) >= _HINT_MAX or used >= _PREFETCH_BUDGET: + break + try: + if not t.is_cuda: + continue + nb = t.numel() * t.element_size() + except AttributeError: + continue + if nb <= 0: + continue + take = min(nb, _PREFETCH_BUDGET - used) + lst.append((int(t.data_ptr()), int(take))) + used += take + _cand_bytes[_ordinal] = used + + +def prefetch_hint_table(): + """The adopted table, for probes and the boot log.""" + return dict(_table) + + def maybe_all_reduce(comm, input_, orig): """Return a reduced tensor if handled here, else None (caller uses NCCL). One-shot only ever serves in REAL mode (shadow=0), where it replaces NCCL at exactly the AR call sites — 4-rank lockstep is automatic. shadow=1 runs the boot self-test then stays permanently on NCCL (observe-only).""" - global _disabled + global _disabled, _ordinal + if _in_forward: + # every collective of the forward counts, whichever path serves it: + # the ordinal is the key the learned hints are filed under + _ordinal += 1 if _disabled: return None if not _connected: @@ -397,6 +505,10 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: + hint = _table.get(_ordinal) if _in_forward else None + if hint: + return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager + input_, [p for p, _ in hint], [n for _, n in hint]) return _ext.oneshot_ar(input_) # real path (works in graph + eager) except Exception as e: raise OneShotFatal( diff --git a/build/dsv4/glm53_megakernel.py b/build/dsv4/glm53_megakernel.py index e840fc4f..642c60be 100644 --- a/build/dsv4/glm53_megakernel.py +++ b/build/dsv4/glm53_megakernel.py @@ -531,12 +531,38 @@ def build_mk_weight_w4_kchunks(weight): for c in range(0, k, MK_GEMM_KMAX)] +_AR_NOTE = None + + +def _ar_note(*tensors) -> None: + """Tell the one-shot AR shim which weights this launch streams. + + The shim (tp_oneshot_ar) attributes the note to the collective that + preceded this launch and, on the next capture, has that collective's + kernel warm these bytes into L2 while it waits for the peers + (VLLM_GLM53_AR_PREFETCH). Resolved once; a lane without the shim + mounted, or with the knob off, pays one attribute read per launch. + Only the eager Python of a launch runs this -- graph replay never does. + """ + global _AR_NOTE + if _AR_NOTE is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as _shim) + _AR_NOTE = _shim.note_consumer + except Exception: + _AR_NOTE = False + if _AR_NOTE: + _AR_NOTE(tensors) + + def _gemm_call(x, mk_pack, n_rows): """mk_pack is (wq4, ws4, gscale) from build_mk_weight_w4.""" import torch out = torch.empty(x.shape[0], n_rows, dtype=torch.bfloat16, device=x.device) + _ar_note(mk_pack[0], mk_pack[1]) _EXT.run_gemm(x.contiguous(), mk_pack[0], mk_pack[1], out, n_rows, float(mk_pack[2])) return out @@ -606,6 +632,7 @@ def _mhc_call(x_flat, residual_flat, pm_flat, cm_flat, fn, hc_scale, layer_input_cur = torch.empty(num_tokens, hidden, dtype=torch.bfloat16, device=x_flat.device) ws = _ensure_workspace(x_flat.device) + _ar_note(fn) _EXT.run_mhc( [x_flat.data_ptr(), residual_flat.data_ptr(), pm_flat.data_ptr(), cm_flat.data_ptr(), fn.data_ptr(), hc_scale.data_ptr(), @@ -887,6 +914,7 @@ def _kda_launch(layer, hidden_states, meta, conv_state, rec_state, out, ow = getattr(layer.o_norm, "weight", None) onorm_w = ow if isinstance(ow, torch.Tensor) else torch.ones( KDA_D, dtype=torch.bfloat16, device=hidden_states.device) + _ar_note(layer._mk_in_pack[0], layer._mk_in_pack[1]) _EXT.run_kda( [hidden_states.data_ptr(), layer._mk_in_pack[0].data_ptr(), diff --git a/build/glm53/dsv4_oneshot_ar.cu b/build/glm53/dsv4_oneshot_ar.cu index 6f6e6c19..38803e9d 100644 --- a/build/glm53/dsv4_oneshot_ar.cu +++ b/build/glm53/dsv4_oneshot_ar.cu @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -146,8 +147,48 @@ __device__ __forceinline__ void osar_backoff(int &n, unsigned &ns) { #endif } +// L2 prefetch hints for the peer-wait window (VLLM_GLM53_AR_PREFETCH): byte +// ranges of the weights the NEXT kernel after this collective streams. The +// wait is 20-40 us of idle DRAM on every rank (MEASUREMENTS 19차: wait 38.7 +// of 45.5 us per collective, ~100 per step); warps 1..7 of every block walk +// these ranges with prefetch.global.L2 while thread 0 polls the peer flags, +// so the consumer finds its first megabytes in L2 (24 MB on this part). +// Passed by value: a CUDA-graph capture bakes the hint with the launch, and +// the shim learns the ranges from the consumers that follow each collective +// during the eager warmups that precede capture. n == 0 is exactly the old +// kernel -- no branch of it touches memory. +#define OSAR_MAXHINT 8 +struct HintArgs { + unsigned long long ptr[OSAR_MAXHINT]; + unsigned int len[OSAR_MAXHINT]; + int n; +}; + +__device__ __forceinline__ void osar_prefetch(const HintArgs &h, + volatile int *landed) { + // One 32 B sector index space over the concatenated ranges, interleaved + // across the grid: block b, thread t takes sectors b*224 + (t-32) + k*10752, + // so every block warms a uniform slice of every range. A prefix walk warms + // a few blocks' tiles of the consumer and leaves its slowest block cold + // (MEASUREMENTS 11차). Owning blocks stop as soon as the peers landed; the + // work is bounded either way (budget <= 20 MB is ~2 us of issue per block). + const int tid = (int)threadIdx.x - 32; + const int stride = ARGRID * (ARTHREADS - 32); + int idx = (int)blockIdx.x * (ARTHREADS - 32) + tid; + for (int r = 0; r < h.n; ++r) { + const unsigned long long base = h.ptr[r]; + const int nsec = (int)((h.len[r] + 31u) >> 5); + for (; idx < nsec; idx += stride) { + asm volatile("prefetch.global.L2 [%0];" ::"l"( + base + ((unsigned long long)idx << 5))); + if (*landed) return; + } + idx -= nsec; + } +} + __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, - int nbytes) { + int nbytes, const HintArgs h) { // The grid is fixed at ARGRID for the counter invariant, so at decode sizes // the smallest plain call has n = hidden and many blocks fall entirely past // the payload: with blockDim 256 and n 4096, blocks 16..47 copy nothing and @@ -212,6 +253,16 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, // Peer wait: rxf is only ever written by the peers' NICs, never by a block // of this kernel -- same independence argument as the guard above. Fence // stays where it always was: after the wait, before reading peer data. + // The prefetch hints ride the wait: thread 0 polls, warps 1..7 warm L2 + // with the next kernel's weights until the peers land (s_landed) or their + // slice is done. A non-owning block has no flag to wait for and simply + // issues its slice; warp 0's other lanes go straight to the barrier. The + // phase timer below still brackets thread 0's poll alone, so t_wait keeps + // measuring the collective -- and shows any DRAM contention the prefetch + // puts on the NIC's writes. + __shared__ volatile int s_landed; + if (threadIdx.x == 0) s_landed = 0; + __syncthreads(); long long t2 = timer ? clock64() : 0; if (owns && threadIdx.x == 0) { // The old form re-read every peer's flag on every pass, including peers @@ -233,6 +284,9 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, } } } + s_landed = 1; + } else if (h.n > 0 && threadIdx.x >= 32) { + osar_prefetch(h, &s_landed); } long long t3 = timer ? clock64() : 0; __syncthreads(); @@ -548,24 +602,55 @@ static void py_connect(std::vector all) { pthread_create(&g_proxy, nullptr, proxy_fn, nullptr); g_started = true; } -static torch::Tensor py_oneshot(torch::Tensor input) { +static torch::Tensor py_oneshot_impl(torch::Tensor input, + const std::vector &ptrs, + const std::vector &lens) { TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kBFloat16); TORCH_CHECK(input.is_contiguous()); int64_t n = input.numel(); TORCH_CHECK(n <= MAXEL, "oneshot: tensor too large"); + TORCH_CHECK(ptrs.size() == lens.size() && ptrs.size() <= OSAR_MAXHINT, + "oneshot: hint lists must pair up, at most OSAR_MAXHINT"); auto out = torch::empty_like(input); const bf16 *src = reinterpret_cast(input.data_ptr()); bf16 *dst = reinterpret_cast(out.data_ptr()); cudaStream_t st = c10::cuda::getCurrentCUDAStream(); + // Prefetch hints: (device pointer, bytes) pairs the shim learned for this + // collective's ordinal. Baked into the launch, so a captured graph replays + // them without any host code. + HintArgs h; + h.n = 0; + for (size_t i = 0; i < ptrs.size(); ++i) { + if (ptrs[i] == 0 || lens[i] <= 0) continue; + h.ptr[h.n] = (unsigned long long)ptrs[i]; + h.len[h.n] = (unsigned int)std::min(lens[i], 0x7fffffff); + ++h.n; + } // One launch, and the grid is FIXED at ARGRID however small n is: the // last-block detection in k_oneshot is (done_ctr % ARGRID == ARGRID-1), // which is only sound if every launch contributes exactly ARGRID // increments. The 48-block grid fills GB10 once and covers MAXEL through // the kernel's grid-stride loops; empty decode blocks only sync/increment. k_oneshot<<>>(g_ctrl, src, dst, (int)n, - (int)(n * 2)); + (int)(n * 2), h); return out; } +static torch::Tensor py_oneshot(torch::Tensor input) { + return py_oneshot_impl(input, {}, {}); +} +static torch::Tensor py_oneshot_hint(torch::Tensor input, + std::vector ptrs, + std::vector lens) { + return py_oneshot_impl(input, ptrs, lens); +} +// The phase counters (SM cycles, monotonic) for a probe that wants the wait +// per collective with and without hints: [guard, copy, wait, reduce, calls]. +static std::vector py_phase_counters() { + if (!g_ctrl) return {}; + return {(int64_t)g_ctrl->t_guard, (int64_t)g_ctrl->t_copy, + (int64_t)g_ctrl->t_wait, (int64_t)g_ctrl->t_reduce, + (int64_t)g_ctrl->t_calls}; +} static bool py_healthy() { if (!g_started) return false; static uint64_t last = 0; @@ -586,6 +671,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("local_infos", &py_local_infos); m.def("connect", &py_connect); m.def("oneshot_ar", &py_oneshot); + m.def("oneshot_ar_hint", &py_oneshot_hint); + m.def("phase_counters", &py_phase_counters); m.def("healthy", &py_healthy); m.def("shutdown", &py_shutdown); } diff --git a/build/glm53/dsv4_oneshot_shim.py b/build/glm53/dsv4_oneshot_shim.py index a76992fb..15657fbd 100644 --- a/build/glm53/dsv4_oneshot_shim.py +++ b/build/glm53/dsv4_oneshot_shim.py @@ -374,13 +374,121 @@ def _eligible(t): ) +# ---- L2 prefetch hints for the peer wait (VLLM_GLM53_AR_PREFETCH) ---------- +# The collective's wait (38.7 of 45.5 us, MEASUREMENTS 19차, ~100 per decode +# step) is idle DRAM on every rank. The kernel can spend it pulling the NEXT +# kernel's weights into L2 (k_oneshot's HintArgs); what it should pull is +# learned here, not declared: every weight-streaming launch of the megakernel +# driver notes the tensors it reads, the note is attributed to the most +# recent collective of the current target forward, and the table so learned +# in the eager warmups is what the captured launches bake in. Nothing here +# runs inside a compiled region: begin/end come from the model class above +# the traced graph, and the notes come from custom-op bodies that execute +# eagerly during capture. +def _resolve_prefetch_budget() -> int: + """Bytes of the next consumer's weights one collective may warm. 0 = off + (the kernel is byte-identical to before), 1 = 12 MB, or an integer MB in + 1..20 -- L2 is 24 MB and the consumer's activations must fit beside it.""" + raw = (os.environ.get("VLLM_GLM53_AR_PREFETCH") or "0").strip().lower() + if raw in ("", "0", "false", "off", "no"): + return 0 + if raw == "1": + return 12 << 20 + try: + mb = int(raw) + except ValueError: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%r is not 0, 1 or an MB count; off", + raw, + ) + return 0 + if not 1 <= mb <= 20: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%d MB is outside 1..20; off", mb + ) + return 0 + return mb << 20 + + +_PREFETCH_BUDGET = _resolve_prefetch_budget() +_HINT_MAX = 8 # OSAR_MAXHINT in the .cu +_in_forward = False +_ordinal = 0 # collectives seen so far in the current target forward +_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes +_cand: dict = {} # the same, being learned in this forward +_cand_bytes: dict = {} + + +def begin_forward(): + """Called by the target model above the compiled region, once per forward + (eager warmup and capture alike). This forward's collectives use the + adopted table; the consumers noting themselves build the candidate.""" + global _in_forward, _ordinal, _cand, _cand_bytes + if not _PREFETCH_BUDGET: + return + _in_forward = True + _ordinal = 0 + _cand = {} + _cand_bytes = {} + + +def end_forward(): + global _in_forward, _table + if not _PREFETCH_BUDGET: + return + _in_forward = False + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + # a decode-shaped forward saw more consumers than the last one did + # (a prefill forward routes M > 32 to stock and notes nothing) + _table = _cand + logger.warning( + "[osar] prefetch hints learned: %d collectives, %.1f MB in total", + len(_table), + sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + ) + + +def note_consumer(tensors): + """A weight-streaming launch reports the tensors it is about to read. + Attributed to the most recent collective of this forward, up to the + budget; the first notes after a collective are the ones that count.""" + if not _PREFETCH_BUDGET or not _in_forward or _ordinal == 0: + return + lst = _cand.setdefault(_ordinal, []) + used = _cand_bytes.get(_ordinal, 0) + for t in tensors: + if t is None or len(lst) >= _HINT_MAX or used >= _PREFETCH_BUDGET: + break + try: + if not t.is_cuda: + continue + nb = t.numel() * t.element_size() + except AttributeError: + continue + if nb <= 0: + continue + take = min(nb, _PREFETCH_BUDGET - used) + lst.append((int(t.data_ptr()), int(take))) + used += take + _cand_bytes[_ordinal] = used + + +def prefetch_hint_table(): + """The adopted table, for probes and the boot log.""" + return dict(_table) + + def maybe_all_reduce(comm, input_, orig): """Return a reduced tensor if handled here, else None (caller uses NCCL). One-shot only ever serves in REAL mode (shadow=0), where it replaces NCCL at exactly the AR call sites — 4-rank lockstep is automatic. shadow=1 runs the boot self-test then stays permanently on NCCL (observe-only).""" - global _disabled + global _disabled, _ordinal + if _in_forward: + # every collective of the forward counts, whichever path serves it: + # the ordinal is the key the learned hints are filed under + _ordinal += 1 if _disabled: return None if not _connected: @@ -397,6 +505,10 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: + hint = _table.get(_ordinal) if _in_forward else None + if hint: + return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager + input_, [p for p, _ in hint], [n for _, n in hint]) return _ext.oneshot_ar(input_) # real path (works in graph + eager) except Exception as e: raise OneShotFatal( diff --git a/build/glm53/glm53_megakernel.py b/build/glm53/glm53_megakernel.py index e840fc4f..642c60be 100644 --- a/build/glm53/glm53_megakernel.py +++ b/build/glm53/glm53_megakernel.py @@ -531,12 +531,38 @@ def build_mk_weight_w4_kchunks(weight): for c in range(0, k, MK_GEMM_KMAX)] +_AR_NOTE = None + + +def _ar_note(*tensors) -> None: + """Tell the one-shot AR shim which weights this launch streams. + + The shim (tp_oneshot_ar) attributes the note to the collective that + preceded this launch and, on the next capture, has that collective's + kernel warm these bytes into L2 while it waits for the peers + (VLLM_GLM53_AR_PREFETCH). Resolved once; a lane without the shim + mounted, or with the knob off, pays one attribute read per launch. + Only the eager Python of a launch runs this -- graph replay never does. + """ + global _AR_NOTE + if _AR_NOTE is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as _shim) + _AR_NOTE = _shim.note_consumer + except Exception: + _AR_NOTE = False + if _AR_NOTE: + _AR_NOTE(tensors) + + def _gemm_call(x, mk_pack, n_rows): """mk_pack is (wq4, ws4, gscale) from build_mk_weight_w4.""" import torch out = torch.empty(x.shape[0], n_rows, dtype=torch.bfloat16, device=x.device) + _ar_note(mk_pack[0], mk_pack[1]) _EXT.run_gemm(x.contiguous(), mk_pack[0], mk_pack[1], out, n_rows, float(mk_pack[2])) return out @@ -606,6 +632,7 @@ def _mhc_call(x_flat, residual_flat, pm_flat, cm_flat, fn, hc_scale, layer_input_cur = torch.empty(num_tokens, hidden, dtype=torch.bfloat16, device=x_flat.device) ws = _ensure_workspace(x_flat.device) + _ar_note(fn) _EXT.run_mhc( [x_flat.data_ptr(), residual_flat.data_ptr(), pm_flat.data_ptr(), cm_flat.data_ptr(), fn.data_ptr(), hc_scale.data_ptr(), @@ -887,6 +914,7 @@ def _kda_launch(layer, hidden_states, meta, conv_state, rec_state, out, ow = getattr(layer.o_norm, "weight", None) onorm_w = ow if isinstance(ow, torch.Tensor) else torch.ones( KDA_D, dtype=torch.bfloat16, device=hidden_states.device) + _ar_note(layer._mk_in_pack[0], layer._mk_in_pack[1]) _EXT.run_kda( [hidden_states.data_ptr(), layer._mk_in_pack[0].data_ptr(), diff --git a/build/glm53/glm5next_model.py b/build/glm53/glm5next_model.py index e5d4d99e..61223b6d 100644 --- a/build/glm53/glm5next_model.py +++ b/build/glm53/glm5next_model.py @@ -1336,7 +1336,39 @@ def forward(self, *args, **kwargs): maybe_free_fp8_dense_bf16(self) except Exception: logger.exception("[fp8-dense] bf16 release skipped") - return super().forward(*args, **kwargs) + # AR prefetch hints (tp_oneshot_ar, VLLM_GLM53_AR_PREFETCH): the shim + # keys its learned hints by "which collective of the target forward", + # so the forward boundary has to come from here -- the same reason + # the release above lives here: this class is above the compiled + # region, and the drafter is a different class, so its collectives + # never see the target's table. No-op unless the knob is set. + osar = _osar_shim() + if osar is not None: + osar.begin_forward() + try: + return super().forward(*args, **kwargs) + finally: + if osar is not None: + osar.end_forward() + + +_OSAR = None + + +def _osar_shim(): + """The one-shot AR shim, resolved once; None when it is not mounted or + predates the prefetch hints.""" + global _OSAR + if _OSAR is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as shim, + ) + + _OSAR = shim if hasattr(shim, "begin_forward") else False + except Exception: + _OSAR = False + return _OSAR or None def get_spec_layer_idx_from_weight_name( diff --git a/launchers/ab-glm53.sh b/launchers/ab-glm53.sh index c267eaf6..0367e48d 100755 --- a/launchers/ab-glm53.sh +++ b/launchers/ab-glm53.sh @@ -26,7 +26,10 @@ HEAD=${HEAD_IP:-10.10.10.2} case "$ARM" in base) ARM_ENV="" ;; - cand) ARM_ENV="VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_MHC=1 VLLM_GLM53_MK_GEMM=1 VLLM_GLM53_MK_MLA=1" ;; + # MK_PDL rides with the segments: the launch form the kernels were tuned + # for (17-19 pct per launch, bench probe), and the profile default since + # 2026-09-04 -- named here so an older profile cannot drop it silently. + cand) ARM_ENV="VLLM_GLM53_MEGAKERNEL=1 VLLM_GLM53_MK_MHC=1 VLLM_GLM53_MK_GEMM=1 VLLM_GLM53_MK_MLA=1 VLLM_GLM53_MK_PDL=1" ;; *) echo "ABORT: arm must be base or cand"; exit 1 ;; esac diff --git a/overlay/modules/glm53_megakernel/README.md b/overlay/modules/glm53_megakernel/README.md index 22e60114..b4ec5894 100644 --- a/overlay/modules/glm53_megakernel/README.md +++ b/overlay/modules/glm53_megakernel/README.md @@ -96,8 +96,11 @@ cold tax by the share this module covers. waits before its first read of the previous kernel's output, so the next MK launch starts on the SMs this one frees and pulls its first W tiles during this one's tail. Two launches back to back measure 17-19% less - per launch than one alone. Behind `VLLM_GLM53_MK_PDL=1` (default off - until the serving bracket; the probe sets it). + per launch than one alone. Behind `VLLM_GLM53_MK_PDL=1` -- on in + `profiles/glm53.env` and in `ab-glm53.sh`'s cand arm since 2026-09-04; + until then serving had never carried it (the driver reads the env, only + the bench probe set it), so every armed boot ran the lane without PDL. + The captured-chain form is checked by `probes/mk_pdl_graph_check.py`. - Fixed 48-block grid everywhere; the never-reset monotonic ticket barrier is what keeps CUDA-graph replay with baked pointers exact (the osar `done_ctr` trick). A larger grid deadlocked on this part (#150). diff --git a/overlay/modules/glm53_megakernel/glm53_megakernel.py b/overlay/modules/glm53_megakernel/glm53_megakernel.py index e840fc4f..642c60be 100644 --- a/overlay/modules/glm53_megakernel/glm53_megakernel.py +++ b/overlay/modules/glm53_megakernel/glm53_megakernel.py @@ -531,12 +531,38 @@ def build_mk_weight_w4_kchunks(weight): for c in range(0, k, MK_GEMM_KMAX)] +_AR_NOTE = None + + +def _ar_note(*tensors) -> None: + """Tell the one-shot AR shim which weights this launch streams. + + The shim (tp_oneshot_ar) attributes the note to the collective that + preceded this launch and, on the next capture, has that collective's + kernel warm these bytes into L2 while it waits for the peers + (VLLM_GLM53_AR_PREFETCH). Resolved once; a lane without the shim + mounted, or with the knob off, pays one attribute read per launch. + Only the eager Python of a launch runs this -- graph replay never does. + """ + global _AR_NOTE + if _AR_NOTE is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as _shim) + _AR_NOTE = _shim.note_consumer + except Exception: + _AR_NOTE = False + if _AR_NOTE: + _AR_NOTE(tensors) + + def _gemm_call(x, mk_pack, n_rows): """mk_pack is (wq4, ws4, gscale) from build_mk_weight_w4.""" import torch out = torch.empty(x.shape[0], n_rows, dtype=torch.bfloat16, device=x.device) + _ar_note(mk_pack[0], mk_pack[1]) _EXT.run_gemm(x.contiguous(), mk_pack[0], mk_pack[1], out, n_rows, float(mk_pack[2])) return out @@ -606,6 +632,7 @@ def _mhc_call(x_flat, residual_flat, pm_flat, cm_flat, fn, hc_scale, layer_input_cur = torch.empty(num_tokens, hidden, dtype=torch.bfloat16, device=x_flat.device) ws = _ensure_workspace(x_flat.device) + _ar_note(fn) _EXT.run_mhc( [x_flat.data_ptr(), residual_flat.data_ptr(), pm_flat.data_ptr(), cm_flat.data_ptr(), fn.data_ptr(), hc_scale.data_ptr(), @@ -887,6 +914,7 @@ def _kda_launch(layer, hidden_states, meta, conv_state, rec_state, out, ow = getattr(layer.o_norm, "weight", None) onorm_w = ow if isinstance(ow, torch.Tensor) else torch.ones( KDA_D, dtype=torch.bfloat16, device=hidden_states.device) + _ar_note(layer._mk_in_pack[0], layer._mk_in_pack[1]) _EXT.run_kda( [hidden_states.data_ptr(), layer._mk_in_pack[0].data_ptr(), diff --git a/overlay/modules/glm53_model_wiring/glm5next_model.py b/overlay/modules/glm53_model_wiring/glm5next_model.py index e5d4d99e..61223b6d 100644 --- a/overlay/modules/glm53_model_wiring/glm5next_model.py +++ b/overlay/modules/glm53_model_wiring/glm5next_model.py @@ -1336,7 +1336,39 @@ def forward(self, *args, **kwargs): maybe_free_fp8_dense_bf16(self) except Exception: logger.exception("[fp8-dense] bf16 release skipped") - return super().forward(*args, **kwargs) + # AR prefetch hints (tp_oneshot_ar, VLLM_GLM53_AR_PREFETCH): the shim + # keys its learned hints by "which collective of the target forward", + # so the forward boundary has to come from here -- the same reason + # the release above lives here: this class is above the compiled + # region, and the drafter is a different class, so its collectives + # never see the target's table. No-op unless the knob is set. + osar = _osar_shim() + if osar is not None: + osar.begin_forward() + try: + return super().forward(*args, **kwargs) + finally: + if osar is not None: + osar.end_forward() + + +_OSAR = None + + +def _osar_shim(): + """The one-shot AR shim, resolved once; None when it is not mounted or + predates the prefetch hints.""" + global _OSAR + if _OSAR is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as shim, + ) + + _OSAR = shim if hasattr(shim, "begin_forward") else False + except Exception: + _OSAR = False + return _OSAR or None def get_spec_layer_idx_from_weight_name( diff --git a/overlay/modules/tp_oneshot_ar/README.md b/overlay/modules/tp_oneshot_ar/README.md index 6f5a40dc..6da485a2 100644 --- a/overlay/modules/tp_oneshot_ar/README.md +++ b/overlay/modules/tp_oneshot_ar/README.md @@ -38,3 +38,34 @@ the existing all-rank bootstrap vote, so every rank stays on NCCL. The fixed graph behavior. A block barrier between each thread's system fence and thread 0's completion atomic also guarantees that the last block cannot publish a partially copied RDMA payload. + +## L2 prefetch during the peer wait (2026-09-04, `VLLM_GLM53_AR_PREFETCH`) + +The collective's wait is 38.7 of its 45.5 us (MEASUREMENTS 19차), ~100 times +per decode step, and during it DRAM is idle on every rank. `k_oneshot` now +takes `HintArgs` -- up to 8 (pointer, bytes) ranges of the weights the NEXT +kernel streams -- and warps 1..7 of every block walk them with +`prefetch.global.L2` (32 B sectors, interleaved across the grid so every +consumer block gets a uniform slice) while thread 0 polls the peer flags. +Owning blocks stop the moment the peers land; the budget (default 12 MB, +knob value = MB, 1..20; L2 is 24 MB) bounds the work either way. `n == 0` is +the old kernel byte for byte. + +What to warm is LEARNED, not declared. The megakernel driver's launches +(`_gemm_call`, `_mhc_call`, `_kda_launch`) call `note_consumer(tensors)`; +the shim files the note under the ordinal of the most recent collective of +the current target forward; `begin_forward`/`end_forward` come from +`Glm5NextForConditionalGeneration.forward`, the class above the compiled +region (the drafter is a different class and never sees the table). A +forward that noted more consumers than the adopted table replaces it -- a +decode-shaped eager warmup does, a prefill forward (M > 32, stock GEMMs) +does not -- and the captured launches bake the adopted table's ranges into +their `HintArgs`. Every collective of the forward advances the ordinal, +whichever path serves it, so NCCL-served prefill collectives keep the keys +aligned. Needs MK-GEMM armed to have anything to learn. + +`oneshot_ar_hint(x, ptrs, lens)` and `phase_counters()` are the probe-facing +bindings; `probes/oneshot_ar_disttest.py` times a 12 MB hint against none and +reads `t_wait` for both, which is where any DRAM contention with the NIC's +writes would show. Ceiling on the critical rank (its wait is the transfer, +~20 us = 4.6 MB at 230 GB/s): ~1.5-2.5 ms/step. Fleet bracket only. diff --git a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_ar.cu b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_ar.cu index 6f6e6c19..38803e9d 100644 --- a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_ar.cu +++ b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_ar.cu @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -146,8 +147,48 @@ __device__ __forceinline__ void osar_backoff(int &n, unsigned &ns) { #endif } +// L2 prefetch hints for the peer-wait window (VLLM_GLM53_AR_PREFETCH): byte +// ranges of the weights the NEXT kernel after this collective streams. The +// wait is 20-40 us of idle DRAM on every rank (MEASUREMENTS 19차: wait 38.7 +// of 45.5 us per collective, ~100 per step); warps 1..7 of every block walk +// these ranges with prefetch.global.L2 while thread 0 polls the peer flags, +// so the consumer finds its first megabytes in L2 (24 MB on this part). +// Passed by value: a CUDA-graph capture bakes the hint with the launch, and +// the shim learns the ranges from the consumers that follow each collective +// during the eager warmups that precede capture. n == 0 is exactly the old +// kernel -- no branch of it touches memory. +#define OSAR_MAXHINT 8 +struct HintArgs { + unsigned long long ptr[OSAR_MAXHINT]; + unsigned int len[OSAR_MAXHINT]; + int n; +}; + +__device__ __forceinline__ void osar_prefetch(const HintArgs &h, + volatile int *landed) { + // One 32 B sector index space over the concatenated ranges, interleaved + // across the grid: block b, thread t takes sectors b*224 + (t-32) + k*10752, + // so every block warms a uniform slice of every range. A prefix walk warms + // a few blocks' tiles of the consumer and leaves its slowest block cold + // (MEASUREMENTS 11차). Owning blocks stop as soon as the peers landed; the + // work is bounded either way (budget <= 20 MB is ~2 us of issue per block). + const int tid = (int)threadIdx.x - 32; + const int stride = ARGRID * (ARTHREADS - 32); + int idx = (int)blockIdx.x * (ARTHREADS - 32) + tid; + for (int r = 0; r < h.n; ++r) { + const unsigned long long base = h.ptr[r]; + const int nsec = (int)((h.len[r] + 31u) >> 5); + for (; idx < nsec; idx += stride) { + asm volatile("prefetch.global.L2 [%0];" ::"l"( + base + ((unsigned long long)idx << 5))); + if (*landed) return; + } + idx -= nsec; + } +} + __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, - int nbytes) { + int nbytes, const HintArgs h) { // The grid is fixed at ARGRID for the counter invariant, so at decode sizes // the smallest plain call has n = hidden and many blocks fall entirely past // the payload: with blockDim 256 and n 4096, blocks 16..47 copy nothing and @@ -212,6 +253,16 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, // Peer wait: rxf is only ever written by the peers' NICs, never by a block // of this kernel -- same independence argument as the guard above. Fence // stays where it always was: after the wait, before reading peer data. + // The prefetch hints ride the wait: thread 0 polls, warps 1..7 warm L2 + // with the next kernel's weights until the peers land (s_landed) or their + // slice is done. A non-owning block has no flag to wait for and simply + // issues its slice; warp 0's other lanes go straight to the barrier. The + // phase timer below still brackets thread 0's poll alone, so t_wait keeps + // measuring the collective -- and shows any DRAM contention the prefetch + // puts on the NIC's writes. + __shared__ volatile int s_landed; + if (threadIdx.x == 0) s_landed = 0; + __syncthreads(); long long t2 = timer ? clock64() : 0; if (owns && threadIdx.x == 0) { // The old form re-read every peer's flag on every pass, including peers @@ -233,6 +284,9 @@ __global__ void k_oneshot(Ctrl *c, const bf16 *src, bf16 *dst, int n, } } } + s_landed = 1; + } else if (h.n > 0 && threadIdx.x >= 32) { + osar_prefetch(h, &s_landed); } long long t3 = timer ? clock64() : 0; __syncthreads(); @@ -548,24 +602,55 @@ static void py_connect(std::vector all) { pthread_create(&g_proxy, nullptr, proxy_fn, nullptr); g_started = true; } -static torch::Tensor py_oneshot(torch::Tensor input) { +static torch::Tensor py_oneshot_impl(torch::Tensor input, + const std::vector &ptrs, + const std::vector &lens) { TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kBFloat16); TORCH_CHECK(input.is_contiguous()); int64_t n = input.numel(); TORCH_CHECK(n <= MAXEL, "oneshot: tensor too large"); + TORCH_CHECK(ptrs.size() == lens.size() && ptrs.size() <= OSAR_MAXHINT, + "oneshot: hint lists must pair up, at most OSAR_MAXHINT"); auto out = torch::empty_like(input); const bf16 *src = reinterpret_cast(input.data_ptr()); bf16 *dst = reinterpret_cast(out.data_ptr()); cudaStream_t st = c10::cuda::getCurrentCUDAStream(); + // Prefetch hints: (device pointer, bytes) pairs the shim learned for this + // collective's ordinal. Baked into the launch, so a captured graph replays + // them without any host code. + HintArgs h; + h.n = 0; + for (size_t i = 0; i < ptrs.size(); ++i) { + if (ptrs[i] == 0 || lens[i] <= 0) continue; + h.ptr[h.n] = (unsigned long long)ptrs[i]; + h.len[h.n] = (unsigned int)std::min(lens[i], 0x7fffffff); + ++h.n; + } // One launch, and the grid is FIXED at ARGRID however small n is: the // last-block detection in k_oneshot is (done_ctr % ARGRID == ARGRID-1), // which is only sound if every launch contributes exactly ARGRID // increments. The 48-block grid fills GB10 once and covers MAXEL through // the kernel's grid-stride loops; empty decode blocks only sync/increment. k_oneshot<<>>(g_ctrl, src, dst, (int)n, - (int)(n * 2)); + (int)(n * 2), h); return out; } +static torch::Tensor py_oneshot(torch::Tensor input) { + return py_oneshot_impl(input, {}, {}); +} +static torch::Tensor py_oneshot_hint(torch::Tensor input, + std::vector ptrs, + std::vector lens) { + return py_oneshot_impl(input, ptrs, lens); +} +// The phase counters (SM cycles, monotonic) for a probe that wants the wait +// per collective with and without hints: [guard, copy, wait, reduce, calls]. +static std::vector py_phase_counters() { + if (!g_ctrl) return {}; + return {(int64_t)g_ctrl->t_guard, (int64_t)g_ctrl->t_copy, + (int64_t)g_ctrl->t_wait, (int64_t)g_ctrl->t_reduce, + (int64_t)g_ctrl->t_calls}; +} static bool py_healthy() { if (!g_started) return false; static uint64_t last = 0; @@ -586,6 +671,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("local_infos", &py_local_infos); m.def("connect", &py_connect); m.def("oneshot_ar", &py_oneshot); + m.def("oneshot_ar_hint", &py_oneshot_hint); + m.def("phase_counters", &py_phase_counters); m.def("healthy", &py_healthy); m.def("shutdown", &py_shutdown); } diff --git a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py index a76992fb..15657fbd 100644 --- a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py +++ b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py @@ -374,13 +374,121 @@ def _eligible(t): ) +# ---- L2 prefetch hints for the peer wait (VLLM_GLM53_AR_PREFETCH) ---------- +# The collective's wait (38.7 of 45.5 us, MEASUREMENTS 19차, ~100 per decode +# step) is idle DRAM on every rank. The kernel can spend it pulling the NEXT +# kernel's weights into L2 (k_oneshot's HintArgs); what it should pull is +# learned here, not declared: every weight-streaming launch of the megakernel +# driver notes the tensors it reads, the note is attributed to the most +# recent collective of the current target forward, and the table so learned +# in the eager warmups is what the captured launches bake in. Nothing here +# runs inside a compiled region: begin/end come from the model class above +# the traced graph, and the notes come from custom-op bodies that execute +# eagerly during capture. +def _resolve_prefetch_budget() -> int: + """Bytes of the next consumer's weights one collective may warm. 0 = off + (the kernel is byte-identical to before), 1 = 12 MB, or an integer MB in + 1..20 -- L2 is 24 MB and the consumer's activations must fit beside it.""" + raw = (os.environ.get("VLLM_GLM53_AR_PREFETCH") or "0").strip().lower() + if raw in ("", "0", "false", "off", "no"): + return 0 + if raw == "1": + return 12 << 20 + try: + mb = int(raw) + except ValueError: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%r is not 0, 1 or an MB count; off", + raw, + ) + return 0 + if not 1 <= mb <= 20: + logger.warning( + "[osar] VLLM_GLM53_AR_PREFETCH=%d MB is outside 1..20; off", mb + ) + return 0 + return mb << 20 + + +_PREFETCH_BUDGET = _resolve_prefetch_budget() +_HINT_MAX = 8 # OSAR_MAXHINT in the .cu +_in_forward = False +_ordinal = 0 # collectives seen so far in the current target forward +_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes +_cand: dict = {} # the same, being learned in this forward +_cand_bytes: dict = {} + + +def begin_forward(): + """Called by the target model above the compiled region, once per forward + (eager warmup and capture alike). This forward's collectives use the + adopted table; the consumers noting themselves build the candidate.""" + global _in_forward, _ordinal, _cand, _cand_bytes + if not _PREFETCH_BUDGET: + return + _in_forward = True + _ordinal = 0 + _cand = {} + _cand_bytes = {} + + +def end_forward(): + global _in_forward, _table + if not _PREFETCH_BUDGET: + return + _in_forward = False + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + # a decode-shaped forward saw more consumers than the last one did + # (a prefill forward routes M > 32 to stock and notes nothing) + _table = _cand + logger.warning( + "[osar] prefetch hints learned: %d collectives, %.1f MB in total", + len(_table), + sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + ) + + +def note_consumer(tensors): + """A weight-streaming launch reports the tensors it is about to read. + Attributed to the most recent collective of this forward, up to the + budget; the first notes after a collective are the ones that count.""" + if not _PREFETCH_BUDGET or not _in_forward or _ordinal == 0: + return + lst = _cand.setdefault(_ordinal, []) + used = _cand_bytes.get(_ordinal, 0) + for t in tensors: + if t is None or len(lst) >= _HINT_MAX or used >= _PREFETCH_BUDGET: + break + try: + if not t.is_cuda: + continue + nb = t.numel() * t.element_size() + except AttributeError: + continue + if nb <= 0: + continue + take = min(nb, _PREFETCH_BUDGET - used) + lst.append((int(t.data_ptr()), int(take))) + used += take + _cand_bytes[_ordinal] = used + + +def prefetch_hint_table(): + """The adopted table, for probes and the boot log.""" + return dict(_table) + + def maybe_all_reduce(comm, input_, orig): """Return a reduced tensor if handled here, else None (caller uses NCCL). One-shot only ever serves in REAL mode (shadow=0), where it replaces NCCL at exactly the AR call sites — 4-rank lockstep is automatic. shadow=1 runs the boot self-test then stays permanently on NCCL (observe-only).""" - global _disabled + global _disabled, _ordinal + if _in_forward: + # every collective of the forward counts, whichever path serves it: + # the ordinal is the key the learned hints are filed under + _ordinal += 1 if _disabled: return None if not _connected: @@ -397,6 +505,10 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: + hint = _table.get(_ordinal) if _in_forward else None + if hint: + return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager + input_, [p for p, _ in hint], [n for _, n in hint]) return _ext.oneshot_ar(input_) # real path (works in graph + eager) except Exception as e: raise OneShotFatal( diff --git a/probes/mk_pdl_graph_check.py b/probes/mk_pdl_graph_check.py new file mode 100644 index 00000000..98b2b015 --- /dev/null +++ b/probes/mk_pdl_graph_check.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Programmatic dependent launch (PDL) of the megakernel segments UNDER CUDA +graph capture -- the form serving actually replays. + +The bench probe measures PDL on a stream (x2 pairs, MEASUREMENTS 2차: 17-19 +pct less per launch); the drafter probe captured ONE MK op with PDL on and +replayed it bitwise. Neither captured a CHAIN of programmatic launches, and +that chain is what VLLM_GLM53_MK_PDL=1 puts into every decode graph: each MK +kernel triggers its dependents at entry and `griddepcontrol.wait`s before it +reads the previous kernel's output, so a chain is correct only if every +pre-wait read is of static data (weights) -- which is what this checks, on +the real kernels, in the real launch form. + +Gates (one process = one PDL state, the .so caches the env once): + chain gemm -> gemm -> gemm and mhc -> gemm, eager == graph replay bitwise + replay two replays of the same graph agree bitwise (monotonic barriers) + timing per-launch us of a 24-launch graph (8 rounds x 3 different packs, + DRAM-cold by rotation) -- run once with VLLM_GLM53_MK_PDL=1 and once + with =0 and read the two rows side by side + +Run inside the glm53 image with the composed overlays mounted: + bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py + VLLM_GLM53_MK_PDL=0 bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py +""" +from __future__ import annotations + +import os +import sys + +os.environ.setdefault("VLLM_GLM53_MEGAKERNEL", "1") +os.environ.setdefault("VLLM_GLM53_MK_MHC", "1") +os.environ.setdefault("VLLM_GLM53_MK_GEMM", "1") +os.environ.setdefault("VLLM_GLM53_MK_KDA", "0") +os.environ.setdefault("VLLM_GLM53_MK_MLA", "0") +os.environ.setdefault("VLLM_GLM53_MK_PDL", "1") +sys.path.insert(0, os.environ.get("MK_PKG_PATH", + "/usr/local/lib/python3.12/dist-packages")) + +import torch # noqa: E402 + +DEV = "cuda" +_FAIL = [] + + +def check(cond: bool, msg: str) -> None: + print((" ok " if cond else " FAIL ") + msg, flush=True) + if not cond: + _FAIL.append(msg) + + +def _same(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool( + torch.equal(a.view(torch.int16), b.view(torch.int16))) + + +def main() -> int: + from vllm.model_executor.layers import glm53_megakernel as mk + + torch.cuda.init() + ext = mk._build() + pdl = bool(ext.pdl_enabled()) if hasattr(ext, "pdl_enabled") else None + env_pdl = os.environ.get("VLLM_GLM53_MK_PDL") + print(f"env VLLM_GLM53_MK_PDL={env_pdl} (.so reports {pdl})") + mk.maybe_arm() + check(mk._ARMED["gemm"], "MK-GEMM armed") + check(mk._ARMED["mhc"], "MK-MHC armed") + if not (mk._ARMED["gemm"] and mk._ARMED["mhc"]): + return 1 + + torch.manual_seed(0) + T = 8 + # --- gemm -> gemm -> gemm: n = k = 4096 so each output feeds the next + packs = [mk.build_mk_weight_w4( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=DEV) * 0.05) + for _ in range(3)] + x0 = torch.randn(T, 4096, dtype=torch.bfloat16, device=DEV) + + def chain_gemm(x): + y = x + for p in packs: + y = mk._gemm_call(y, p, 4096) + return y + + # --- mhc -> gemm: the MHC tail writes layer_input, the GEMM reads it + H, HC, NOUT = mk.HIDDEN, mk.HC, mk.NOUT + res = torch.randn(T, HC, H, dtype=torch.bfloat16, device=DEV) * 0.1 + pm = torch.rand(T, HC, dtype=torch.float32, device=DEV) + cm = torch.rand(T, HC * HC, dtype=torch.float32, device=DEV).contiguous() + fn = torch.randn(NOUT, HC * H, dtype=torch.float32, device=DEV) * 0.02 + hc_scale = torch.ones(3, dtype=torch.float32, device=DEV) + hc_base = torch.zeros(NOUT, dtype=torch.float32, device=DEV) + nw = torch.randn(H, dtype=torch.bfloat16, device=DEV) + p_in = mk.build_mk_weight_w4( + torch.randn(6416, 4096, dtype=torch.bfloat16, device=DEV) * 0.05) + + def chain_mhc(x): + rc, pmc, cmc, li = mk._mhc_call( + x, res, pm, cm, fn, hc_scale, hc_base, nw, T, 1e-6, 1e-6, 1e-6, + 1.0, 1e-6, mk.SINKHORN_SERVED) + return mk._gemm_call(li, p_in, 6416), rc, pmc, cmc + + # eager references (warm the arm first) + for _ in range(3): + chain_gemm(x0) + chain_mhc(x0) + torch.cuda.synchronize() + ref_g = chain_gemm(x0) + ref_m = chain_mhc(x0) + torch.cuda.synchronize() + + # capture both chains into one graph, replay twice + st = torch.cuda.Stream() + st.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(st): + for _ in range(2): + chain_gemm(x0) + chain_mhc(x0) + torch.cuda.current_stream().wait_stream(st) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(g, stream=st): + out_g = chain_gemm(x0) + out_m = chain_mhc(x0) + except Exception as e: # noqa: BLE001 + check(False, f"CUDA-graph capture of the MK chains raised: {e!r}"[:300]) + return 1 + torch.cuda.synchronize() + g.replay() + torch.cuda.synchronize() + first = (out_g.clone(), [t.clone() for t in out_m]) + g.replay() + torch.cuda.synchronize() + check(_same(first[0], ref_g), "gemm->gemm->gemm: graph replay == eager bitwise") + check(all(_same(a, b) for a, b in zip(first[1], ref_m)), + "mhc->gemm: graph replay == eager bitwise (out, residual, mixes)") + check(_same(out_g, first[0]) and all(_same(a, b) for a, b in zip(out_m, first[1])), + "second replay == first replay bitwise (monotonic barriers)") + + # --- timing: a 24-launch graph, DRAM-cold by rotation (8 x 3 packs of + # 8.4 MB = 200 MB per replay, L2 is 24 MB), per-launch us + rounds = 8 + rot = [mk.build_mk_weight_w4( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=DEV) * 0.05) + for _ in range(rounds * 3)] + xs = torch.randn(T, 4096, dtype=torch.bfloat16, device=DEV) + + def many(): + y = xs + for p in rot: + y = mk._gemm_call(y, p, 4096) + return y + + with torch.cuda.stream(st): + for _ in range(3): + many() + torch.cuda.current_stream().wait_stream(st) + torch.cuda.synchronize() + g2 = torch.cuda.CUDAGraph() + with torch.cuda.graph(g2, stream=st): + many() + torch.cuda.synchronize() + for _ in range(5): + g2.replay() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + N = 40 + s.record() + for _ in range(N): + g2.replay() + e.record() + torch.cuda.synchronize() + per = s.elapsed_time(e) * 1e3 / (N * len(rot)) + nb = rot[0][0].numel() + rot[0][1].numel() + print(f"timing pdl={env_pdl}: {len(rot)}-launch graph, n=k=4096, " + f"{per:.1f} us/launch, {nb / per / 1e3:.0f} GB/s over the pack bytes") + print("VERDICT:", "PASS" if not _FAIL else f"FAIL ({len(_FAIL)})") + return 0 if not _FAIL else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/probes/moe_decode_stream_probe.py b/probes/moe_decode_stream_probe.py new file mode 100644 index 00000000..c4d3cf67 --- /dev/null +++ b/probes/moe_decode_stream_probe.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Go/no-go for MK_SEG_MOE: how fast does the SERVED b12x MoE kernel stream +expert bytes at the decode shape, against what this part can stream. + +The 21차 trace put the routed MoE at 31 ms/step (44 pct) and called it +"bandwidth floor" -- but that 190 GB/s was back-computed from a GUESSED +~40 unique experts per layer, and the MK W4 lane streams 244 GB/s in-loop +on the same part. So the floor claim is circular. This probe measures the +kernel directly: the b12x wrapper the serving path constructs (same +geometry, same dispatch overlay, same static/micro cutover), at C=1 decode +(8 tokens x top-8 = 64 routed pairs -> the STATIC backend), with the number +of unique experts U controlled, weights DRAM-cold by rotating 8 full expert +sets (8 x 1 GB, L2 is 24 MB), CUDA-graph replay timed with events. + +Bytes per unique expert per rank (GLM-5.3 TP=4, nvfp4): w13 [1024 x 2048 B] +2.10 MB + sf 0.26 MB + w2 [4096 x 256 B] 1.05 MB + sf 0.13 MB = 3.54 MB. + +Rows: + b12x U=.. us per MoE call, effective GB/s = U x 3.54 MB / t + lane the MK W4 GEMM streaming 12 different [6416 x 4096] packs in + one graph (161 MB, PDL on): the rate a persistent MK segment + would stream the same bytes at + torch read fp32 sum over 161 MB: a library read-only reference + warm gemm one n=6416 launch cold vs after its pack was pulled into L2 + (what an AR-wait prefetch of a consumer's pack buys) + +Verdict rule (strategy doc): b12x >= 90 pct of the lane rate closes the +MK_SEG_MOE axis; below it, the gap x 31 ms is the segment's ceiling. + + bash probes/run_mk_probe.sh probes/moe_decode_stream_probe.py +""" +from __future__ import annotations + +import os +import sys + +os.environ.setdefault("VLLM_GLM53_MEGAKERNEL", "1") +os.environ.setdefault("VLLM_GLM53_MK_GEMM", "1") +os.environ.setdefault("VLLM_GLM53_MK_MHC", "0") +os.environ.setdefault("VLLM_GLM53_MK_KDA", "0") +os.environ.setdefault("VLLM_GLM53_MK_MLA", "0") +os.environ.setdefault("VLLM_GLM53_MK_PDL", "1") +sys.path.insert(0, os.environ.get("MK_PKG_PATH", + "/usr/local/lib/python3.12/dist-packages")) + +import torch # noqa: E402 + +DEV = "cuda" +E, TOPK, HID, INTER = 288, 8, 4096, 512 # per-rank intermediate 2048 / 4 +T = 8 # C=1 verify batch (k=7 + 1) +SETS = 8 # weight sets rotated per replay +BYTES_PER_EXPERT = 1024 * 2048 + 1024 * 256 + 4096 * 256 + 4096 * 32 + + +def _graph(fn, stream): + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + fn() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=stream): + fn() + torch.cuda.synchronize() + return g + + +def _time_graph(g, reps: int) -> float: + for _ in range(3): + g.replay() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(reps): + g.replay() + e.record() + torch.cuda.synchronize() + return s.elapsed_time(e) * 1e3 / reps # us per replay + + +def _routing(U: int): + """[T, TOPK] int32 with exactly U distinct experts, 8 distinct per token.""" + gen = torch.Generator().manual_seed(U) + pool = torch.randperm(E, generator=gen)[:U] + flat = torch.arange(T * TOPK) % U + ids = pool[flat].view(T, TOPK).to(torch.int32) + w = torch.rand(T, TOPK, generator=gen, dtype=torch.float32) + w = w / w.sum(dim=1, keepdim=True) + return ids.to(DEV), w.to(DEV) + + +def _weight_set(gen): + w13 = torch.randint(0, 256, (E, 2 * INTER, HID // 2), dtype=torch.uint8, + generator=gen).to(DEV) + w2 = torch.randint(0, 256, (E, HID, INTER // 2), dtype=torch.uint8, + generator=gen).to(DEV) + s13 = (torch.rand(E, 2 * INTER, HID // 16, generator=gen) * 0.05 + 0.01) + s2 = (torch.rand(E, HID, INTER // 16, generator=gen) * 0.05 + 0.01) + return w13, w2, s13.to(torch.float8_e4m3fn).to(DEV), s2.to(torch.float8_e4m3fn).to(DEV) + + +def main() -> int: + from flashinfer.fused_moe import B12xMoEWrapper + from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout + from vllm.model_executor.layers import glm53_megakernel as mk + + torch.cuda.init() + print(f"device {torch.cuda.get_device_name()} sets={SETS} T={T} topk={TOPK}") + stream = torch.cuda.Stream() + + # ---- the served wrapper: exact GLM TP geometry, static workspace only + wrapper = B12xMoEWrapper( + num_experts=E, top_k=TOPK, hidden_size=HID, intermediate_size=INTER, + use_cuda_graph=True, max_num_tokens=64, num_local_experts=E, + activation="swigluoai_uninterleave", swiglu_alpha=1.0, + swiglu_beta=0.0, swiglu_limit=10.0) + gen = torch.Generator().manual_seed(1) + sets = [] + for _ in range(SETS): + w13, w2, s13, s2 = _weight_set(gen) + sf13 = flashinfer_convert_sf_to_mma_layout( + s13.reshape(E * 2 * INTER, HID // 16), m=2 * INTER, k=HID, + num_groups=E) + sf2 = flashinfer_convert_sf_to_mma_layout( + s2.reshape(E * HID, INTER // 16), m=HID, k=INTER, num_groups=E) + sets.append((w13, sf13, w2, sf2)) + ones = torch.ones(E, dtype=torch.float32, device=DEV) + x = torch.randn(T, HID, dtype=torch.bfloat16, device=DEV) * 0.5 + out = torch.empty(T, HID, dtype=torch.bfloat16, device=DEV) + print(f"{'row':<14}{'us/call':>9}{'MB':>8}{'GB/s':>8}") + + rates = {} + for U in (8, 16, 24, 32, 40, 48, 56, 64): + ids, w = _routing(U) + + def moe_all(): + for w13, sf13, w2, sf2 in sets: + wrapper.run(x, w13, sf13, w2, sf2, ids, w, w1_alpha=ones, + w2_alpha=ones, fc2_input_scale=ones, out=out) + + g = _graph(moe_all, stream) + us = _time_graph(g, 20) / SETS + mb = U * BYTES_PER_EXPERT / 1e6 + rates[U] = mb * 1e6 / us / 1e3 + print(f"{'b12x U=' + str(U):<14}{us:>9.1f}{mb:>8.1f}{rates[U]:>8.0f}") + del g + + # ---- the MK W4 lane on the same class of bytes: 12 packs of [6416 x + # 4096] (13.4 MB each, 161 MB per replay), PDL on, one graph + mk.maybe_arm() + assert mk._ARMED["gemm"], "MK-GEMM did not arm" + packs = [mk.build_mk_weight_w4( + torch.randn(6416, HID, dtype=torch.bfloat16, device=DEV) * 0.05) + for _ in range(12)] + xg = torch.randn(T, HID, dtype=torch.bfloat16, device=DEV) + + def lane_all(): + for p in packs: + mk._gemm_call(xg, p, 6416) + + g = _graph(lane_all, stream) + us = _time_graph(g, 20) / len(packs) + nb = packs[0][0].numel() + packs[0][1].numel() + lane = nb / us / 1e3 + print(f"{'lane n=6416':<14}{us:>9.1f}{nb / 1e6:>8.1f}{lane:>8.0f}") + del g + + # ---- a library read-only reference over the same 161 MB + big = torch.randn(len(packs) * nb // 4, dtype=torch.float32, device=DEV) + g = _graph(lambda: big.sum(), stream) + us = _time_graph(g, 20) + print(f"{'torch sum':<14}{us:>9.1f}{big.numel() * 4 / 1e6:>8.1f}" + f"{big.numel() * 4 / us / 1e3:>8.0f}") + del g + + # ---- what an L2 prefetch of a consumer's pack buys: one n=6416 launch, + # cold (after a 2 x 24 MB flush) vs after the pack was read once + flush = torch.empty(48 << 20, dtype=torch.int8, device=DEV) + drain = torch.zeros(16 << 20, dtype=torch.float32, device=DEV) + p = packs[0] + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + cold, warm = [], [] + for i in range(20): + flush.zero_() + drain.sum() + drain.sum() + xg.add_(0) + if i % 2: + p[0].view(torch.int32).sum() # pull nibbles into L2 (clean) + p[1].view(torch.int32).sum() # and the group scales + s.record() + mk._gemm_call(xg, p, 6416) + e.record() + torch.cuda.synchronize() + (warm if i % 2 else cold).append(s.elapsed_time(e) * 1e3) + cold.sort() + warm.sort() + print(f"{'gemm cold':<14}{cold[len(cold) // 2]:>9.1f}{nb / 1e6:>8.1f}" + f"{nb / cold[len(cold) // 2] / 1e3:>8.0f}") + print(f"{'gemm L2-warm':<14}{warm[len(warm) // 2]:>9.1f}{nb / 1e6:>8.1f}" + f"{nb / warm[len(warm) // 2] / 1e3:>8.0f}") + + best = max(rates.values()) + print(f"b12x best {best:.0f} GB/s = {100 * best / lane:.0f} pct of the " + f"lane's {lane:.0f} GB/s; U=40 -> {rates[40]:.0f} GB/s " + f"({100 * rates[40] / lane:.0f} pct)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/probes/oneshot_ar_disttest.py b/probes/oneshot_ar_disttest.py index d3c95baf..d18b34e8 100644 --- a/probes/oneshot_ar_disttest.py +++ b/probes/oneshot_ar_disttest.py @@ -64,6 +64,46 @@ print(f"rank{RANK} GRAPH one-shot AR: {us:.2f} us/call | healthy={ext.healthy()}", flush=True) + +def _wait_us(c0, c1): + return (c1[2] - c0[2]) / max(1, c1[4] - c0[4]) / 1592.0 # SM_CLK_MHZ + + +# --- the peer wait as L2 prefetch time: a 12 MB hint (one KDA in_proj W4 +# pack's worth) against none, same graph shape, per-call us and the kernel's +# own t_wait for both. A hint that costs the wait more than it warms is a +# loss; the consumer-side gain is measured by moe_decode_stream_probe.py +# (gemm cold vs L2-warm) and the step by the fleet bracket. +x8 = torch.full((8, 4096), float(RANK + 1), dtype=torch.bfloat16, device="cuda") +hint = torch.empty(12 << 20, dtype=torch.uint8, device="cuda") +flush = torch.empty(48 << 20, dtype=torch.uint8, device="cuda") +for label, ptrs, lens in (("no hint", [], []), + ("12 MB hint", [hint.data_ptr()], [hint.numel()])): + gh = torch.cuda.CUDAGraph() + with torch.cuda.graph(gh): + for _ in range(8): + out = ext.oneshot_ar_hint(x8, ptrs, lens) + torch.cuda.synchronize() + for _ in range(20): + gh.replay() + torch.cuda.synchronize() + err = (out.float() - 10.0).abs().max().item() + dist.barrier() + c0 = ext.phase_counters() + t0 = time.time() + for _ in range(500): + flush.zero_() # the hint must be cold every time, as in a step + gh.replay() + torch.cuda.synchronize() + dt = time.time() - t0 + c1 = ext.phase_counters() + # the flush (48 MB at ~200 GB/s, ~250 us) is inside the loop; report the + # kernel's own wait, which excludes it, beside the wall time that does not + print(f"rank{RANK} {label:>11}: wall {dt / (500 * 8) * 1e6:.1f} us/call " + f"(incl. flush) | t_wait {_wait_us(c0, c1):.1f} us/call | " + f"maxerr={err:.4f} (expect 0)", flush=True) + dist.barrier() + dist.barrier() ext.shutdown() dist.destroy_process_group() diff --git a/probes/osar_build_check.py b/probes/osar_build_check.py new file mode 100644 index 00000000..ccef745b --- /dev/null +++ b/probes/osar_build_check.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Compile the one-shot AR extension the way the shim does and confirm the +prefetch-hint bindings exist. A ptxas rejection of a new instruction would +otherwise surface as "[osar] local setup failed" on every rank of a real +boot -- and that boot falls back to NCCL, silently losing the AR module. + + bash probes/run_mk_probe.sh probes/osar_build_check.py +""" +import os +import sys +import tempfile + +sys.path.insert(0, os.environ.get("MK_PKG_PATH", + "/usr/local/lib/python3.12/dist-packages")) + +import torch # noqa: E402 +from torch.utils.cpp_extension import load # noqa: E402 + +SRC = "/repo/overlay/modules/tp_oneshot_ar/dsv4_oneshot_ar.cu" + + +def main() -> int: + build = tempfile.mkdtemp(prefix="osar_build_") + ext = load(name="dsv4_oneshot_ar_check", sources=[SRC], + extra_cuda_cflags=["-O2", "-arch=sm_121a"], + extra_ldflags=["-libverbs"], build_directory=build, + verbose=False) + names = [n for n in dir(ext) if not n.startswith("_")] + print("bound:", " ".join(sorted(names))) + ok = all(n in names for n in ("oneshot_ar", "oneshot_ar_hint", + "phase_counters", "healthy")) + print("osar build:", "PASS" if ok else "FAIL (missing binding)") + # no RDMA here: init() needs the fabric. The compile is the gate. + print("torch", torch.__version__, "cuda", torch.version.cuda) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/probes/run_mk_probe.sh b/probes/run_mk_probe.sh new file mode 100644 index 00000000..42aaae3f --- /dev/null +++ b/probes/run_mk_probe.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# Run ANY offline probe against the composed glm53 overlay sources, in a fresh +# container -- the general form of run_megakernel_bench.sh (which is pinned to +# the ladder-step-2 probe and its six files). Mounts the megakernel driver AND +# the b12x MoE files, so a probe can drive the served MoE path and the MK lane +# side by side. +# +# bash probes/run_mk_probe.sh probes/.py [probe args...] +# +# Never the serving container (docker ps must be empty of glm53*), never while +# a TP=4 boot is up on any host: MEASUREMENTS 11차 lost two measurements to a +# serving boot on another node. VLLM_GLM53_*/VLLM_DSV4_* set in this shell are +# forwarded (a probe arms itself with setdefault, so a forwarded knob WINS). +set -euo pipefail + +REPO=$(cd "$(dirname "$0")/.." && pwd) +PROBE=${1:?usage: run_mk_probe.sh probes/.py [args...]} +shift +[ -f "$REPO/$PROBE" ] || { echo "ABORT: no such probe: $REPO/$PROBE" >&2; exit 1; } + +PROFILE=${PROFILE:-glm53} +ENVFILE="$REPO/profiles/$PROFILE.env" +[ -f "$ENVFILE" ] || { echo "ABORT: no such profile: $ENVFILE" >&2; exit 1; } +PROFILE_IMAGE="" +TARGET_PREFIX="/opt/venv/lib/python3.12/site-packages/" +eval "$( + # shellcheck disable=SC1090 + . "$ENVFILE" + printf 'PROFILE_IMAGE=%q\nTARGET_PREFIX=%q\n' \ + "${PROFILE_IMAGE:-}" "${TARGET_PREFIX:-/opt/venv/lib/python3.12/site-packages/}" +)" +IMAGE=${IMAGE:-$PROFILE_IMAGE} +[ -n "$IMAGE" ] || { echo "ABORT: $PROFILE names no image" >&2; exit 1; } + +BUILD="$REPO/build/$PROFILE" +MANIFEST="$BUILD/manifest.tsv" +bash "$REPO/launchers/compose-overlays.sh" "$PROFILE" >&2 + +# The megakernel driver and everything its arm touches, plus the served MoE +# path (vLLM's b12x experts, flashinfer's wrapper and its sm12x dispatch). +sources=(glm53_megakernel.py glm53_megakernel.cu + glm5next_kda.py tilelang.py tilelang_kernels.py glm53_fp8_dense.py + flashinfer_b12x_moe.py b12x_moe.py moe_dispatch.py moe_micro_kernel.py) +mounts=() +for source in "${sources[@]}"; do + target=$(awk -F '\t' -v source="$source" '$1 == source {print $2}' "$MANIFEST") + [ -n "$target" ] || { echo "ABORT: $source is missing from $MANIFEST" >&2; exit 1; } + [ "${target#*$'\n'}" = "$target" ] \ + || { echo "ABORT: duplicate $source rows in $MANIFEST" >&2; exit 1; } + [ -f "$BUILD/$source" ] \ + || { echo "ABORT: composed source missing: $BUILD/$source" >&2; exit 1; } + mounts+=(--mount "type=bind,src=$BUILD/$source,dst=$target,readonly") +done + +case "${VLLM_GLM53_MEGAKERNEL-1}" in + 0|""|false|FALSE|no|off) + echo "ABORT: VLLM_GLM53_MEGAKERNEL='${VLLM_GLM53_MEGAKERNEL}' would disarm the probe; unset it" >&2 + exit 1 ;; +esac + +envs=(-e "MK_PKG_PATH=${TARGET_PREFIX%/}") +_fwd="" +for v in $(compgen -v | grep -E '^VLLM_(GLM53|DSV4)_'); do + envs+=(-e "$v=${!v}") + _fwd="$_fwd $v=${!v}" +done +echo "profile=$PROFILE image=$IMAGE probe=$PROBE files=${#sources[@]} args=$*" >&2 +echo "forwarded:${_fwd:- (none)}" >&2 + +cmd="python3 /repo/$PROBE" +for a in "$@"; do cmd="$cmd $(printf %q "$a")"; done + +exec docker run --rm --gpus all --entrypoint /bin/bash \ + --mount "type=bind,src=$REPO,dst=/repo,readonly" \ + "${envs[@]}" "${mounts[@]}" \ + "$IMAGE" -lc "$cmd" diff --git a/profiles/glm53.env b/profiles/glm53.env index 64b4e96d..560b7e84 100644 --- a/profiles/glm53.env +++ b/profiles/glm53.env @@ -188,6 +188,24 @@ VLLM_GLM53_MK_MHC=0 VLLM_GLM53_MK_GEMM=0 VLLM_GLM53_MK_KDA=0 VLLM_GLM53_MK_KDA_SHADOW=0 +# Programmatic dependent launch for every MK launch: the next MK kernel +# starts on the SMs this one frees and pulls its first W tiles during its +# tail (17-19 pct less per launch back to back, MEASUREMENTS 2차). Serving +# never carried it: the driver reads this env, the bench probe sets it, and +# neither this profile nor ab-glm53.sh did -- every armed boot to date ran +# the lane without PDL. Numerics are unchanged (each kernel +# griddepcontrol.waits before its first read of the previous output; the +# pre-wait fill is weights only) and the captured-chain form is checked by +# probes/mk_pdl_graph_check.py. Inert while no segment is armed. +VLLM_GLM53_MK_PDL=1 +# tp_oneshot_ar: while a collective waits for its peers (38.7 of 45.5 us, +# ~100 per step, DRAM idle), warm the next kernel's weights into L2. What +# to warm is learned from the megakernel launches that follow each +# collective in the eager warmups, so this needs MK-GEMM armed to have +# anything to learn. 0 = off (kernel byte-identical), 1 = 12 MB per +# collective, N = N MB (1..20). Speed only; ceiling ~2-3 pct/step +# (strategy doc 2026-09-04), bracket on the fleet like any arm. +VLLM_GLM53_AR_PREFETCH=0 # MK_SEG_MLA: our own sparse MLA decode kernel (bf16 mma; see the segment # comment in glm53_megakernel.cu). glm53_mk_mla_wiring routes decode # shapes (T<=32) to it; isolated 1.33x the FlashInfer wrapper at C=1. A diff --git a/tests/test_logic.py b/tests/test_logic.py index e132626f..f96301d3 100644 --- a/tests/test_logic.py +++ b/tests/test_logic.py @@ -6925,6 +6925,157 @@ def test_osar_wait_is_split_by_message_size() -> None: print(" osar wait is split by message size .. OK") +def test_osar_prefetch_hints_contract() -> None: + """The one-shot AR's peer wait doubles as an L2 prefetch of the next + kernel's weights (VLLM_GLM53_AR_PREFETCH), and what it warms is learned + from the megakernel launches that follow each collective. + + Contract: the kernel takes the hints by value (a captured graph bakes + them), only warps 1..7 issue them and only inside the wait window (thread + 0's poll and its timer are untouched, so t_wait keeps measuring the + collective), an empty hint is the old kernel; the shim counts EVERY + collective of the target forward so NCCL-served prefill calls keep the + ordinals aligned, and the driver notes its weights before each launch.""" + cu = open(os.path.join(REPO, "overlay/modules/tp_oneshot_ar/" + "dsv4_oneshot_ar.cu"), encoding="utf-8").read() + check("#define OSAR_MAXHINT 8" in cu and "struct HintArgs {" in cu + and "int nbytes, const HintArgs h) {" in cu, + "k_oneshot takes up to 8 (ptr, bytes) hints by value") + check(cu.count('asm volatile("prefetch.global.L2 [%0];"') == 1 + and cu.index('asm volatile("prefetch.global.L2') + > cu.index("__device__ __forceinline__ void osar_prefetch("), + "one prefetch instruction, inside osar_prefetch") + call_at = cu.find("} else if (h.n > 0 && threadIdx.x >= 32) {\n" + " osar_prefetch(h, &s_landed);") + t2_at = cu.find("long long t2 = timer ? clock64() : 0;") + t3_at = cu.find("long long t3 = timer ? clock64() : 0;") + check(0 < t2_at < call_at < t3_at, + "the prefetch runs only in the peer-wait window, on warps 1..7, " + "while thread 0 polls") + check(" s_landed = 1;\n } else if (h.n > 0" in cu, + "owning blocks release their prefetch warps the moment the peers " + "land") + check("k_oneshot<<>>(g_ctrl, src, dst, (int)n," + in cu and "(int)(n * 2), h);" in cu, + "the fixed-geometry launch carries the hints") + check('m.def("oneshot_ar_hint", &py_oneshot_hint);' in cu + and 'm.def("phase_counters", &py_phase_counters);' in cu + and 'm.def("oneshot_ar", &py_oneshot);' in cu, + "hint and counter bindings beside the unchanged plain entry") + + shim_path = os.path.join(REPO, "overlay/modules/tp_oneshot_ar/" + "dsv4_oneshot_shim.py") + shim = open(shim_path, encoding="utf-8").read() + check("_PREFETCH_BUDGET = _resolve_prefetch_budget()" in shim + and "def begin_forward():" in shim and "def end_forward():" in shim + and "def note_consumer(tensors):" in shim, + "the shim learns the hint table per target forward") + mar = shim[shim.index("def maybe_all_reduce("):] + check(mar.index("_ordinal += 1") < mar.index("if _disabled:"), + "every collective advances the ordinal before any path decision") + check("hint = _table.get(_ordinal) if _in_forward else None" in mar + and "_ext.oneshot_ar_hint(" in mar and "_ext.oneshot_ar(input_)" in mar, + "hints apply only inside a target forward; the plain path stays") + + drv = open(os.path.join(REPO, "overlay/modules/glm53_megakernel/" + "glm53_megakernel.py"), encoding="utf-8").read() + for site, note, launch in ( + ("gemm", "_ar_note(mk_pack[0], mk_pack[1])", "_EXT.run_gemm("), + ("mhc", "_ar_note(fn)", "_EXT.run_mhc("), + ("kda", "_ar_note(layer._mk_in_pack[0], layer._mk_in_pack[1])", + "_EXT.run_kda(")): + n_at, l_at = drv.find(note), drv.find(launch) + check(0 < n_at < l_at, f"{site} launch notes its weights first") + + wiring = open(os.path.join(REPO, "overlay/modules/glm53_model_wiring/" + "glm5next_model.py"), encoding="utf-8").read() + cls_at = wiring.index("class Glm5NextForConditionalGeneration(") + layer_at = wiring.index("class Glm5NextDecoderLayer(") + b_at = wiring.index("osar.begin_forward()") + check(b_at > cls_at and "osar.end_forward()" in wiring[b_at:] + and "def _osar_shim():" in wiring, + "the forward boundary comes from the class above the compiled region") + check("begin_forward" not in wiring[layer_at:cls_at], + "no hint call inside the traced decoder layer") + + profile = open(os.path.join(REPO, "profiles", "glm53.env"), + encoding="utf-8").read() + check(re.search(r"^VLLM_GLM53_AR_PREFETCH=0$", profile, re.M) is not None, + "the prefetch knob is declared off in the profile") + check(re.search(r"^VLLM_GLM53_MK_PDL=1$", profile, re.M) is not None, + "PDL is the profile default for the MK launches (2026-09-04)") + ab = open(os.path.join(REPO, "launchers", "ab-glm53.sh"), + encoding="utf-8").read() + check("VLLM_GLM53_MK_PDL=1" in ab.split("cand)", 1)[1].split("\n", 1)[0], + "the A/B cand arm names MK_PDL explicitly") + bracket = open(os.path.join(REPO, "bench", "bracket.py"), + encoding="utf-8").read() + check('"VLLM_GLM53_AR_PREFETCH"' in bracket, + "bracket.py snapshots the prefetch knob") + + # the budget parser and the learning protocol, on the real module + import importlib.util + + def load(env: str, tag: str): + old = os.environ.get("VLLM_GLM53_AR_PREFETCH") + os.environ["VLLM_GLM53_AR_PREFETCH"] = env + try: + spec = importlib.util.spec_from_file_location( + f"_osar_shim_{tag}", shim_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + finally: + if old is None: + os.environ.pop("VLLM_GLM53_AR_PREFETCH", None) + else: + os.environ["VLLM_GLM53_AR_PREFETCH"] = old + return mod + + check(load("0", "a")._PREFETCH_BUDGET == 0 + and load("1", "b")._PREFETCH_BUDGET == 12 << 20 + and load("8", "c")._PREFETCH_BUDGET == 8 << 20 + and load("40", "d")._PREFETCH_BUDGET == 0 + and load("x", "e")._PREFETCH_BUDGET == 0, + "budget: 0 off, 1 = 12 MB, N = N MB in 1..20, else off") + + class _T: + is_cuda = True + + def __init__(self, ptr, nbytes): + self._p, self._n = ptr, nbytes + + def numel(self): + return self._n + + def element_size(self): + return 1 + + def data_ptr(self): + return self._p + + m = load("1", "f") + m.begin_forward() + m.note_consumer([_T(0x10, 100)]) # before any collective: dropped + m._ordinal = 1 + m.note_consumer([_T(0x1000, 10 << 20), _T(0x2000, 10 << 20)]) + m._ordinal = 2 + m.note_consumer([_T(0x3000, 64)]) + m.end_forward() + t = m.prefetch_hint_table() + check(0 not in t and t.get(1) == [(0x1000, 10 << 20), (0x2000, 2 << 20)] + and t.get(2) == [(0x3000, 64)], + "notes file under the preceding collective, capped at the budget") + m.begin_forward() + m._ordinal = 1 + m.note_consumer([_T(0x9000, 8)]) + m.end_forward() + check(m.prefetch_hint_table() == t, + "a forward with fewer notes (prefill-shaped) does not replace the " + "table") + check(m._in_forward is False, "end_forward closes the window") + print(" osar prefetch hints contract .. OK") + + def test_kv_cache_is_pinned_in_tokens() -> None: """KV is pinned in tokens, not left to take whatever GMU leaves. @@ -9556,6 +9707,7 @@ def test_common_tp4_library_is_the_one_implementation() -> None: test_fp8_dense_drafter_patterns_and_opaque_op() test_ab_runner_measures_both_channels() test_osar_wait_is_split_by_message_size() + test_osar_prefetch_hints_contract() test_glm53_megakernel_contracts() test_prefill_warmup_contracts() test_megakernel_w4_layout_functional() From 8450b7914e8187564182550c366632b571a4cf89 Mon Sep 17 00:00:00 2001 From: choiceoh Date: Fri, 4 Sep 2026 17:26:18 +0900 Subject: [PATCH 2/3] =?UTF-8?q?feat(osar,drafter):=20=ED=94=84=EB=A6=AC?= =?UTF-8?q?=ED=8E=98=EC=B9=98=20=ED=9E=8C=ED=8A=B8=20=ED=91=9C=EB=A5=BC=20?= =?UTF-8?q?=EB=AA=A8=EB=8D=B8=EB=B3=84=EB=A1=9C,=20=EB=93=9C=EB=9E=98?= =?UTF-8?q?=ED=94=84=ED=84=B0=20forward=20=EC=97=90=EB=8F=84=20=EA=B2=BD?= =?UTF-8?q?=EA=B3=84=20--=20=EA=B7=B8=EB=A6=AC=EA=B3=A0=2027=EC=B0=A8=20?= =?UTF-8?q?=EC=9B=90=EC=9E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - shim: 힌트 표를 scope("target"/"drafter")별로 든다. 하나의 표를 나눠 쓰면 타깃의 더 풍부한 표가 드래프터의 열 개 콜렉티브 행을 덮거나, 드래프터의 ordinal 이 타깃의 행을 읽는다. begin_forward(scope) / prefetch_hint_table(scope). - DFlash2Qwen3ForCausalLM.forward: 컴파일된 DFlashQwen3Model 위에서 "drafter" scope 의 경계를 연다 -- 드래프터의 층당 두 콜렉티브도 다음 GEMM 팩을 데운다. - MEASUREMENTS 27차: 서빙 PDL 은 꺼져 있었다(발사당 −7.6%), MoE 커널은 레인과 같은 197~200 GB/s(90% 규칙으로 닫힘), L2-warm 은 W4 GEMM 에 7%(프리페치 상한 −0.4~0.6 ms 로 정정), AR 빌드 PASS, 드래프터 꼬리 인구조사, 전략 상한 대 프로브 정정표. - RUNBOOK EXP-12/13/14 결과 줄. Co-Authored-By: Claude Fable 5.1 --- MEASUREMENTS.md | 73 +++++++++++++++++++ RUNBOOK_KERNEL_CAMPAIGN2.md | 9 +++ build/dsv4/dsv4_oneshot_shim.py | 44 ++++++----- build/glm53/dsv4_oneshot_shim.py | 44 ++++++----- build/glm53/glm5next_model.py | 2 +- build/glm53/qwen3_dflash2.py | 43 +++++++++++ .../glm53_dflash2_fp8_head/qwen3_dflash2.py | 43 +++++++++++ .../glm53_model_wiring/glm5next_model.py | 2 +- .../tp_oneshot_ar/dsv4_oneshot_shim.py | 44 ++++++----- tests/test_logic.py | 41 ++++++++--- 10 files changed, 283 insertions(+), 62 deletions(-) diff --git a/MEASUREMENTS.md b/MEASUREMENTS.md index be09e904..fdbc11da 100644 --- a/MEASUREMENTS.md +++ b/MEASUREMENTS.md @@ -1869,6 +1869,79 @@ C>1 은 요청마다 수락률이 달라 배치 구성이 흔들려 단일 정 **먹힌 것**: W4 확장의 로컬 메모리 바이트 배열 → 레지스터 워드 산술(−20~33%; 타일 우선 팩 + cp.async 3단, 더 깊으면 손해); MHC p2 청크 축약을 24 레인에 + 4×4 sinkhorn 16 레인(bar2 13.5 → 5.4 µs; T=8 46→36, T=32 76→65); PDL 연쇄 발사(gemm 발사당 −3~5%, `VLLM_GLM53_MK_PDL`, 기본 off); ksr 호스트 계산·폴드 나눗셈 제거; 단계 스탬프(`-DMK_PHASE_TS`, gemm/mhc, wait 누적, 상한 프로브 `MK_PROBE_SKIP`). **안 먹힌 것(기록)**: W 채움 끌어올림(배리어 대기로 자리만 옮김), exact wait + 깊이 4·5, last-arriver 폴드, 동적 유닛 배분(n=6416 +1~2%), x cp.async 스테이징(프롤로그가 채움 전체를 기다림), smem 예산 공유(W8 −4~7% → 별도 인스턴스), mhc p1 fn 1회 읽기·fn 재배열(토큰당 지연이 지배), m≤16 mma 가드(언롤 안 break), 스위즐 3종. +## ★★메가커널 27차 — 서빙 PDL 은 꺼져 있었다, MoE 커널은 우리 레인과 같은 속도, L2-warm 은 W4 GEMM 에 7% (2026-09-04, srv2) + +전략 문서(2026-09-04 아티팩트 "메가커널 대대적 개선 방안")의 후보 중 운영자가 고른 셋 — 서빙 +PDL, AR 대기 중 다음 커널 L2 프리페치, MK_SEG_MOE go/no-go — 를 `probes/run_mk_probe.sh`(b12x +파일까지 마운트하는 일반 러너)로 srv2 빈 창(다른 세션의 DRAFTW4 브래킷이 끝난 뒤)에서 쟀다. + +**1. 서빙 PDL (EXP-12).** 발견: `mk_pdl_enabled()` 는 env `VLLM_GLM53_MK_PDL` 을 읽는데 glm53.env +에도 `ab-glm53.sh` 의 cand 팔에도 없었고 벤치 프로브만 켰다 → 09-03 무장 부팅(21차 트레이스)의 +173발은 전부 PDL 없이 돌았다. `probes/mk_pdl_graph_check.py`(그래프 캡처 안 체인): gemm→gemm→gemm 과 +mhc→gemm 의 리플레이가 eager 와 비트 동일, 재리플레이 동일 — PDL on/off 모두. 24발 그래프 +(n=k=4096, 8.4 MB 팩 24개 회전, DRAM-cold): **off 58.0 → on 53.6 µs/발사(−7.6%)**. 173발/스텝이면 +−0.76 ms(1.1%). 프로필 기본 1, cand 팔에 명시. 수치 불변이라 EXP-6 브래킷의 cand 팔에 얹는다. + +**2. MK_SEG_MOE go/no-go (EXP-14).** 서빙이 만드는 b12x 래퍼(같은 기하·같은 디스패치 오버레이; +C=1 은 64 pairs 라 static 백엔드)를 디코드 형상(8토큰·top-8)에서 고유 전문가 U 별로, 가중치 +DRAM-cold(8 세트 순환, 세트당 1 GB), 그래프 리플레이로: + +| U | µs/호출 | MB | GB/s | +|---|---|---|---| +| 8 / 16 / 24 | 148 / 311 / 451 | 28 / 57 / 85 | 191 / 182 / 188 | +| 32 / **40** / 48 | 573 / **720** / 859 | 113 / 142 / 170 | 198 / **197** / 198 | +| 56 / 64 | 985 / 1135 | 198 / 227 | 201 / 200 | +| MK W4 레인 n=6416 (12팩 그래프, PDL) | 76.7 | 15.0 | **196** | +| torch sum (읽기만, 180 MB) | 736 | 180 | 245 | + +U=40 의 720 µs 는 21차의 층당 738 µs 와 맞는다 — 서빙의 고유 전문가 수가 ~40 이라는 것과 +실효 ~197 GB/s 가 둘 다 확인됐다. **판정: b12x = 레인의 103% → 90% 규칙으로 축을 닫는다.** +남은 20%(196~200 vs 읽기 245)는 두 커널에 공통인 "발사 안 구조" 몫이다: 레인은 루프만 244 +GB/s 이고 발사당 ~15 µs 의 고정비(4차)가 있으며, MoE 커널은 층당 한 발사라 고정비가 아니라 +타일 구조다. persistent 스트림이 240 에 닿는다는 증거는 아직 없다(레인의 발사당 196~222). 그러니 +MK_SEG_MOE 의 상한 −4~5 ms 는 그 가정 위에 있고, 착수하려면 먼저 2~3일짜리 "persistent 전문가 +타일 스트림" 마이크로커널로 240 도달 여부를 재야 한다. 지금은 보류. + +**3. AR 프리페치의 소비자 이득 (EXP-13 게이트 3).** n=6416 W4 GEMM 한 발: DRAM-cold 86.0 µs → +팩 15 MB 를 L2 에 미리 읽어 둔 뒤 79.9 µs(**−7%**). W4 GEMM 은 L2-warm 이어도 80 µs 다 — DRAM 이 +아니라 발행(cp.async/LSU, 2차)에 묶여 있어서, 프리페치가 사는 것은 DRAM 지연분뿐이다. 회당 +~10 MB 를 데워도 −4 µs/AR, 스텝당 100회면 **−0.4~0.6 ms(0.6~0.9%)**. 전략 문서의 −1.5~2.5 ms 는 +"L2 는 DRAM 의 4배" 가정이었고 그 가정이 틀렸다. 구현은 랜딩했다(기본 off, 힌트 없으면 커널 +byte-identical): 4랭크 disttest(`t_wait` 힌트/무힌트)와 브래킷은 부팅이고, 단독으로는 CV 1.7% +아래라 EXP-6+12 위에 얹어서만 판정한다. + +**4. AR 커널 빌드 (EXP-13 게이트 1).** `probes/osar_build_check.py`: `prefetch.global.L2` 가 sm_121a +ptxas 를 통과하고 `oneshot_ar_hint`·`phase_counters` 가 묶였다. PASS. + +**5. 드래프터 꼬리 인구조사 (09-03 트레이스, 깨끗한 디코드 스텝 5·6).** 샘플러 앵커에서 다음 스텝의 +준비 커널까지 136 커널, 커널 합 6.28 ms(스팬 6.1~7.1): bf16 cutlass GEMM 33 × 122 = 4.0 ms(EXP-10 이 +서빙되면 ~1.25), 타깃 헤드 fp8 814, `k_oneshot` 11 × 72 = 788(forward 안의 45 보다 느리다 — 꼬리 +구간의 랭크 편차), `kernel_mha` 5 × 29 = 145, AllGather 2 × 54 = 108, 글루 ~50개 0.33 ms. 드래프터 +메가커널의 상한 = 글루 0.33 + 발사 간극 ~0.35 + MK 발사 고정비 30 × ~10 µs ≈ **−0.8~1.0 ms +(1.2~1.5%)**; AR 0.79 ms 는 융합해도 남는다. + +**6. 전략 문서의 상한 대 프로브 (정정).** 운영자 물음 "실질 이득이 적다면 왜 처음엔 크다고 +봤나" 에 대한 답: + +| 항목 | 전략 상한 | 프로브 뒤 | 틀린 가정 | +|---|---|---|---| +| AR 프리페치 | −1.5~2.5 ms | −0.4~0.6 | 소비자가 DRAM 바운드라는 가정. 2차의 "발행 병목" 을 L2-warm 에 적용하지 않았다 | +| MK_SEG_MOE | −4~8 ms | 닫힘 | MoE 190 을 레인의 루프 전용 244 와 비교. 같은 형태(그래프 안 발사당)면 레인 196 = MoE 197 | +| 샤딩 샘플링 | −0.4 ms | −0.1 | 25차의 AllGather 409~567 µs 는 프리필 문맥 스텝의 값. 깨끗한 디코드는 2 × 54 µs → **후보에서 내림** | +| 드래프터 메가커널 | −1~1.5 ms | −0.8~1.0 | 대체로 맞음(위 5) | +| 서빙 PDL | −0.9~1.7 ms | −0.76 | 범위 안 | + +공통 원인: 상한을 목표와 같은 측정 형태(그래프 안 발사당, 깨끗한 디코드 스텝, 확인된 병목)로 +만들지 않았다. 남은 큰 항목 — k 축소(SPEC=0 대 k=7 실측에 앵커), W4 팔(발사당 실측), EXP-7(호스트 +유휴 실측) — 은 그 형태의 측정 위에 있다. + +**교훈**: (1) 노브에 독자가 있어도 서빙이 그 값을 나른다는 보장은 없다 — 프로브가 켜는 env 는 +프로필에도 있어야 하고, 무장 부팅의 fingerprint 에 찍혀야 한다. (2) "대역폭 바닥" 판정은 같은 +바이트를 우리 커널이 얼마나 빨리 흘리는지와 견줘야 한다 — 역산값끼리 맞는 것은 검증이 아니다. +(3) L2-warm 이 이득이 되려면 소비자가 DRAM 바운드여야 한다 — W4 GEMM 은 아니었다. +(4) 상한은 목표와 같은 측정 형태로 — 루프 전용 스탬프·프로파일러 꼬리·미확인 병목 가정으로 만든 +상한은 프로브 하나에 절반 이하로 준다. + ## ★★★26차 — 부팅의 메모리 절벽: fp8-dense 패스가 로드 중에 여러 번 돌았고, W4 팩 빌더가 예약 메모리를 흘렸다 (2026-09-04, 4노드 계측) 드래프터 W4 브래킷(25차 승인)을 띄우자 base 팔 첫 부팅이 로드 시작 13초 뒤 죽었다. srv1 은 diff --git a/RUNBOOK_KERNEL_CAMPAIGN2.md b/RUNBOOK_KERNEL_CAMPAIGN2.md index ac1370ed..ebf844b1 100644 --- a/RUNBOOK_KERNEL_CAMPAIGN2.md +++ b/RUNBOOK_KERNEL_CAMPAIGN2.md @@ -496,6 +496,9 @@ bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py # P VLLM_GLM53_MK_PDL=0 bash probes/run_mk_probe.sh probes/mk_pdl_graph_check.py # 대조 ``` +**결과(2026-09-04, srv2, 27차)**: 체인 리플레이 비트 동일(on/off 모두), 24발 그래프 off 58.0 → +on 53.6 µs/발사(−7.6%). 프로필 기본 1. 남은 것 = EXP-6 브래킷의 cand 팔. + ## EXP-13 — AR 대기 중 다음 커널 가중치 L2 프리페치 (`VLLM_GLM53_AR_PREFETCH`, 2026-09-04 추가) `k_oneshot` 의 대기(회당 38.7/45.5 µs, 스텝당 ~100회)는 모든 랭크에서 DRAM 이 @@ -515,6 +518,9 @@ MK 드라이버의 발사(`_gemm_call`·`_mhc_call`·`_kda_launch`)가 읽을 N collectives, X MB` 가 없으면 무장이 아니다. - 상한: 임계 랭크의 대기 ≈ 전송 ~20 µs = 4.6 MB → −1.5~2.5 ms/스텝(전략 문서). - 예산 노브: 1 = 12 MB/콜렉티브, N = N MB(1..20; L2 24 MB). +- **결과(2026-09-04, 27차)**: 게이트 1 PASS(빌드), 게이트 3 = n=6416 W4 GEMM cold 86.0 → L2-warm + 79.9 µs(−7%): 소비자가 DRAM 바운드가 아니라 상한은 **−0.4~0.6 ms/스텝**으로 내려갔다(전략 문서의 + −1.5~2.5 는 철회). 단독 판정 불가 — EXP-6+12 위에 얹어서만. 게이트 2(4랭크 disttest)는 부팅. ## EXP-14 — MK_SEG_MOE go/no-go (2026-09-04 추가, 프로브만) @@ -530,6 +536,9 @@ b12x 래퍼(같은 기하, 같은 디스패치 오버레이; C=1 은 64 pairs persistent, FC1 (전문가, n타일) 유닛 → 전문가별 완료 카운터로 열리는 FC2 동적 큐, 공유 전문가 = 41번째 전문가, b12x nvfp4 레이아웃 제자리 읽기, A4→A8). - 실제 서빙 U 는 다음 부팅에서 로그 한 줄로 확정한다(프로브는 U 별 곡선만 준다). +- **결과(2026-09-04, 27차)**: b12x static U=40 = 197 GB/s, 레인 n=6416 = 196 GB/s → **103%, 닫힘**. + 읽기 전용 참조(torch sum) 245 GB/s 와의 20% 는 두 커널 공통의 발사 안 구조 몫이라, 재개 조건은 + "persistent 전문가 타일 스트림 마이크로커널이 240 에 닿는다" 는 2~3일 프로브의 양성이다. 보류. ## 브래킷 자동화 — `bench/bracket.py` (도구, 판정 아님) diff --git a/build/dsv4/dsv4_oneshot_shim.py b/build/dsv4/dsv4_oneshot_shim.py index 15657fbd..581f44ff 100644 --- a/build/dsv4/dsv4_oneshot_shim.py +++ b/build/dsv4/dsv4_oneshot_shim.py @@ -413,38 +413,48 @@ def _resolve_prefetch_budget() -> int: _PREFETCH_BUDGET = _resolve_prefetch_budget() _HINT_MAX = 8 # OSAR_MAXHINT in the .cu _in_forward = False -_ordinal = 0 # collectives seen so far in the current target forward -_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes -_cand: dict = {} # the same, being learned in this forward +_scope = "target" # which model's forward: "target" or "drafter" +_ordinal = 0 # collectives seen so far in the current forward +# scope -> ordinal -> [(ptr, nbytes)], adopted; what capture bakes. One table +# per model: the drafter's ten collectives have their own consumers, and a +# shared table would let the richer target forward overwrite them (or the +# drafter's ordinals read the target's rows). +_tables: dict = {} +_cand: dict = {} # the current forward's candidate, ordinal -> ranges _cand_bytes: dict = {} -def begin_forward(): - """Called by the target model above the compiled region, once per forward +def begin_forward(scope: str = "target"): + """Called by a model class above its compiled region, once per forward (eager warmup and capture alike). This forward's collectives use the - adopted table; the consumers noting themselves build the candidate.""" - global _in_forward, _ordinal, _cand, _cand_bytes + scope's adopted table; the consumers noting themselves build the + candidate. `scope` names the model -- the target and the drafter each + keep their own table.""" + global _in_forward, _scope, _ordinal, _cand, _cand_bytes if not _PREFETCH_BUDGET: return _in_forward = True + _scope = scope _ordinal = 0 _cand = {} _cand_bytes = {} def end_forward(): - global _in_forward, _table + global _in_forward if not _PREFETCH_BUDGET: return _in_forward = False - if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + table = _tables.get(_scope, {}) + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in table.values()): # a decode-shaped forward saw more consumers than the last one did # (a prefill forward routes M > 32 to stock and notes nothing) - _table = _cand + _tables[_scope] = _cand logger.warning( - "[osar] prefetch hints learned: %d collectives, %.1f MB in total", - len(_table), - sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + "[osar] prefetch hints learned (%s): %d collectives, %.1f MB in total", + _scope, + len(_cand), + sum(sum(n for _, n in v) for v in _cand.values()) / 1e6, ) @@ -473,9 +483,9 @@ def note_consumer(tensors): _cand_bytes[_ordinal] = used -def prefetch_hint_table(): - """The adopted table, for probes and the boot log.""" - return dict(_table) +def prefetch_hint_table(scope: str = "target"): + """The adopted table of one scope, for probes and the boot log.""" + return dict(_tables.get(scope, {})) def maybe_all_reduce(comm, input_, orig): @@ -505,7 +515,7 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: - hint = _table.get(_ordinal) if _in_forward else None + hint = _tables.get(_scope, {}).get(_ordinal) if _in_forward else None if hint: return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager input_, [p for p, _ in hint], [n for _, n in hint]) diff --git a/build/glm53/dsv4_oneshot_shim.py b/build/glm53/dsv4_oneshot_shim.py index 15657fbd..581f44ff 100644 --- a/build/glm53/dsv4_oneshot_shim.py +++ b/build/glm53/dsv4_oneshot_shim.py @@ -413,38 +413,48 @@ def _resolve_prefetch_budget() -> int: _PREFETCH_BUDGET = _resolve_prefetch_budget() _HINT_MAX = 8 # OSAR_MAXHINT in the .cu _in_forward = False -_ordinal = 0 # collectives seen so far in the current target forward -_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes -_cand: dict = {} # the same, being learned in this forward +_scope = "target" # which model's forward: "target" or "drafter" +_ordinal = 0 # collectives seen so far in the current forward +# scope -> ordinal -> [(ptr, nbytes)], adopted; what capture bakes. One table +# per model: the drafter's ten collectives have their own consumers, and a +# shared table would let the richer target forward overwrite them (or the +# drafter's ordinals read the target's rows). +_tables: dict = {} +_cand: dict = {} # the current forward's candidate, ordinal -> ranges _cand_bytes: dict = {} -def begin_forward(): - """Called by the target model above the compiled region, once per forward +def begin_forward(scope: str = "target"): + """Called by a model class above its compiled region, once per forward (eager warmup and capture alike). This forward's collectives use the - adopted table; the consumers noting themselves build the candidate.""" - global _in_forward, _ordinal, _cand, _cand_bytes + scope's adopted table; the consumers noting themselves build the + candidate. `scope` names the model -- the target and the drafter each + keep their own table.""" + global _in_forward, _scope, _ordinal, _cand, _cand_bytes if not _PREFETCH_BUDGET: return _in_forward = True + _scope = scope _ordinal = 0 _cand = {} _cand_bytes = {} def end_forward(): - global _in_forward, _table + global _in_forward if not _PREFETCH_BUDGET: return _in_forward = False - if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + table = _tables.get(_scope, {}) + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in table.values()): # a decode-shaped forward saw more consumers than the last one did # (a prefill forward routes M > 32 to stock and notes nothing) - _table = _cand + _tables[_scope] = _cand logger.warning( - "[osar] prefetch hints learned: %d collectives, %.1f MB in total", - len(_table), - sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + "[osar] prefetch hints learned (%s): %d collectives, %.1f MB in total", + _scope, + len(_cand), + sum(sum(n for _, n in v) for v in _cand.values()) / 1e6, ) @@ -473,9 +483,9 @@ def note_consumer(tensors): _cand_bytes[_ordinal] = used -def prefetch_hint_table(): - """The adopted table, for probes and the boot log.""" - return dict(_table) +def prefetch_hint_table(scope: str = "target"): + """The adopted table of one scope, for probes and the boot log.""" + return dict(_tables.get(scope, {})) def maybe_all_reduce(comm, input_, orig): @@ -505,7 +515,7 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: - hint = _table.get(_ordinal) if _in_forward else None + hint = _tables.get(_scope, {}).get(_ordinal) if _in_forward else None if hint: return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager input_, [p for p, _ in hint], [n for _, n in hint]) diff --git a/build/glm53/glm5next_model.py b/build/glm53/glm5next_model.py index 61223b6d..30750d1d 100644 --- a/build/glm53/glm5next_model.py +++ b/build/glm53/glm5next_model.py @@ -1344,7 +1344,7 @@ def forward(self, *args, **kwargs): # never see the target's table. No-op unless the knob is set. osar = _osar_shim() if osar is not None: - osar.begin_forward() + osar.begin_forward("target") try: return super().forward(*args, **kwargs) finally: diff --git a/build/glm53/qwen3_dflash2.py b/build/glm53/qwen3_dflash2.py index 5df224e3..e66c46c9 100644 --- a/build/glm53/qwen3_dflash2.py +++ b/build/glm53/qwen3_dflash2.py @@ -271,6 +271,25 @@ def forward( +_OSAR = None + + +def _osar_shim(): + """The one-shot AR shim, resolved once; None when it is not mounted or + predates the prefetch hints (the same resolver the target model uses).""" + global _OSAR + if _OSAR is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as shim, + ) + + _OSAR = shim if hasattr(shim, "begin_forward") else False + except Exception: + _OSAR = False + return _OSAR or None + + def dflash2_selector_load_verdict(stats): """Did the path selector's weights actually load? Pure predicate. @@ -351,6 +370,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # AR prefetch hints for the drafter's own collectives (tp_oneshot_ar, + # VLLM_GLM53_AR_PREFETCH): the drafter runs TP=4 with two collectives + # per layer, and each one can warm the next GEMM's W4 pack while it + # waits for the peers. The shim keys its learned hints by "which + # collective of which model's forward", so the boundary has to come + # from THIS class -- above the compiled DFlashQwen3Model, like the + # target's boundary sits above its compiled region -- and under its + # own scope, so the target's richer table never overwrites the + # drafter's. No-op unless the knob is set or the shim is not mounted. + osar = _osar_shim() + if osar is not None: + osar.begin_forward("drafter") + try: + return super().forward(input_ids, positions, inputs_embeds) + finally: + if osar is not None: + osar.end_forward() + def verify_selector_loaded(self) -> None: """Say out loud whether the path selector's weights actually arrived. diff --git a/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py b/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py index 5df224e3..e66c46c9 100644 --- a/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py +++ b/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py @@ -271,6 +271,25 @@ def forward( +_OSAR = None + + +def _osar_shim(): + """The one-shot AR shim, resolved once; None when it is not mounted or + predates the prefetch hints (the same resolver the target model uses).""" + global _OSAR + if _OSAR is None: + try: + from vllm.distributed.device_communicators import ( + dsv4_oneshot_shim as shim, + ) + + _OSAR = shim if hasattr(shim, "begin_forward") else False + except Exception: + _OSAR = False + return _OSAR or None + + def dflash2_selector_load_verdict(stats): """Did the path selector's weights actually load? Pure predicate. @@ -351,6 +370,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # AR prefetch hints for the drafter's own collectives (tp_oneshot_ar, + # VLLM_GLM53_AR_PREFETCH): the drafter runs TP=4 with two collectives + # per layer, and each one can warm the next GEMM's W4 pack while it + # waits for the peers. The shim keys its learned hints by "which + # collective of which model's forward", so the boundary has to come + # from THIS class -- above the compiled DFlashQwen3Model, like the + # target's boundary sits above its compiled region -- and under its + # own scope, so the target's richer table never overwrites the + # drafter's. No-op unless the knob is set or the shim is not mounted. + osar = _osar_shim() + if osar is not None: + osar.begin_forward("drafter") + try: + return super().forward(input_ids, positions, inputs_embeds) + finally: + if osar is not None: + osar.end_forward() + def verify_selector_loaded(self) -> None: """Say out loud whether the path selector's weights actually arrived. diff --git a/overlay/modules/glm53_model_wiring/glm5next_model.py b/overlay/modules/glm53_model_wiring/glm5next_model.py index 61223b6d..30750d1d 100644 --- a/overlay/modules/glm53_model_wiring/glm5next_model.py +++ b/overlay/modules/glm53_model_wiring/glm5next_model.py @@ -1344,7 +1344,7 @@ def forward(self, *args, **kwargs): # never see the target's table. No-op unless the knob is set. osar = _osar_shim() if osar is not None: - osar.begin_forward() + osar.begin_forward("target") try: return super().forward(*args, **kwargs) finally: diff --git a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py index 15657fbd..581f44ff 100644 --- a/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py +++ b/overlay/modules/tp_oneshot_ar/dsv4_oneshot_shim.py @@ -413,38 +413,48 @@ def _resolve_prefetch_budget() -> int: _PREFETCH_BUDGET = _resolve_prefetch_budget() _HINT_MAX = 8 # OSAR_MAXHINT in the .cu _in_forward = False -_ordinal = 0 # collectives seen so far in the current target forward -_table: dict = {} # ordinal -> [(ptr, nbytes)], adopted; what capture bakes -_cand: dict = {} # the same, being learned in this forward +_scope = "target" # which model's forward: "target" or "drafter" +_ordinal = 0 # collectives seen so far in the current forward +# scope -> ordinal -> [(ptr, nbytes)], adopted; what capture bakes. One table +# per model: the drafter's ten collectives have their own consumers, and a +# shared table would let the richer target forward overwrite them (or the +# drafter's ordinals read the target's rows). +_tables: dict = {} +_cand: dict = {} # the current forward's candidate, ordinal -> ranges _cand_bytes: dict = {} -def begin_forward(): - """Called by the target model above the compiled region, once per forward +def begin_forward(scope: str = "target"): + """Called by a model class above its compiled region, once per forward (eager warmup and capture alike). This forward's collectives use the - adopted table; the consumers noting themselves build the candidate.""" - global _in_forward, _ordinal, _cand, _cand_bytes + scope's adopted table; the consumers noting themselves build the + candidate. `scope` names the model -- the target and the drafter each + keep their own table.""" + global _in_forward, _scope, _ordinal, _cand, _cand_bytes if not _PREFETCH_BUDGET: return _in_forward = True + _scope = scope _ordinal = 0 _cand = {} _cand_bytes = {} def end_forward(): - global _in_forward, _table + global _in_forward if not _PREFETCH_BUDGET: return _in_forward = False - if sum(len(v) for v in _cand.values()) > sum(len(v) for v in _table.values()): + table = _tables.get(_scope, {}) + if sum(len(v) for v in _cand.values()) > sum(len(v) for v in table.values()): # a decode-shaped forward saw more consumers than the last one did # (a prefill forward routes M > 32 to stock and notes nothing) - _table = _cand + _tables[_scope] = _cand logger.warning( - "[osar] prefetch hints learned: %d collectives, %.1f MB in total", - len(_table), - sum(sum(n for _, n in v) for v in _table.values()) / 1e6, + "[osar] prefetch hints learned (%s): %d collectives, %.1f MB in total", + _scope, + len(_cand), + sum(sum(n for _, n in v) for v in _cand.values()) / 1e6, ) @@ -473,9 +483,9 @@ def note_consumer(tensors): _cand_bytes[_ordinal] = used -def prefetch_hint_table(): - """The adopted table, for probes and the boot log.""" - return dict(_table) +def prefetch_hint_table(scope: str = "target"): + """The adopted table of one scope, for probes and the boot log.""" + return dict(_tables.get(scope, {})) def maybe_all_reduce(comm, input_, orig): @@ -505,7 +515,7 @@ def maybe_all_reduce(comm, input_, orig): "rank-local NCCL fallback" ) try: - hint = _table.get(_ordinal) if _in_forward else None + hint = _tables.get(_scope, {}).get(_ordinal) if _in_forward else None if hint: return _ext.oneshot_ar_hint( # real path + L2 hints, graph + eager input_, [p for p, _ in hint], [n for _, n in hint]) diff --git a/tests/test_logic.py b/tests/test_logic.py index f96301d3..9533541f 100644 --- a/tests/test_logic.py +++ b/tests/test_logic.py @@ -6967,15 +6967,19 @@ def test_osar_prefetch_hints_contract() -> None: "dsv4_oneshot_shim.py") shim = open(shim_path, encoding="utf-8").read() check("_PREFETCH_BUDGET = _resolve_prefetch_budget()" in shim - and "def begin_forward():" in shim and "def end_forward():" in shim - and "def note_consumer(tensors):" in shim, - "the shim learns the hint table per target forward") + and 'def begin_forward(scope: str = "target"):' in shim + and "def end_forward():" in shim + and "def note_consumer(tensors):" in shim + and "_tables[_scope] = _cand" in shim, + "the shim learns one hint table per model forward (scope)") mar = shim[shim.index("def maybe_all_reduce("):] check(mar.index("_ordinal += 1") < mar.index("if _disabled:"), "every collective advances the ordinal before any path decision") - check("hint = _table.get(_ordinal) if _in_forward else None" in mar - and "_ext.oneshot_ar_hint(" in mar and "_ext.oneshot_ar(input_)" in mar, - "hints apply only inside a target forward; the plain path stays") + check("hint = _tables.get(_scope, {}).get(_ordinal) if _in_forward else None" + in mar and "_ext.oneshot_ar_hint(" in mar + and "_ext.oneshot_ar(input_)" in mar, + "hints apply only inside a forward, from that model's own table; " + "the plain path stays") drv = open(os.path.join(REPO, "overlay/modules/glm53_megakernel/" "glm53_megakernel.py"), encoding="utf-8").read() @@ -6993,8 +6997,19 @@ def test_osar_prefetch_hints_contract() -> None: layer_at = wiring.index("class Glm5NextDecoderLayer(") b_at = wiring.index("osar.begin_forward()") check(b_at > cls_at and "osar.end_forward()" in wiring[b_at:] - and "def _osar_shim():" in wiring, + and "def _osar_shim():" in wiring + and 'osar.begin_forward("target")' in wiring, "the forward boundary comes from the class above the compiled region") + drafter = open(os.path.join(REPO, "overlay/modules/glm53_dflash2_fp8_head/" + "qwen3_dflash2.py"), encoding="utf-8").read() + d_cls = drafter.index("class DFlash2Qwen3ForCausalLM(") + d_b = drafter.index('osar.begin_forward("drafter")') + check(d_b > d_cls and "osar.end_forward()" in drafter[d_b:] + and "def _osar_shim():" in drafter + and "class DFlash2Qwen3Model(DFlashQwen3Model)" in drafter + and 'begin_forward' not in drafter[drafter.index("class DFlash2Qwen3Model("):d_cls], + "the drafter's boundary sits on its ForCausalLM class (above the " + "compiled DFlashQwen3Model) under its own scope") check("begin_forward" not in wiring[layer_at:cls_at], "no hint call inside the traced decoder layer") @@ -7054,7 +7069,7 @@ def data_ptr(self): return self._p m = load("1", "f") - m.begin_forward() + m.begin_forward("target") m.note_consumer([_T(0x10, 100)]) # before any collective: dropped m._ordinal = 1 m.note_consumer([_T(0x1000, 10 << 20), _T(0x2000, 10 << 20)]) @@ -7065,7 +7080,7 @@ def data_ptr(self): check(0 not in t and t.get(1) == [(0x1000, 10 << 20), (0x2000, 2 << 20)] and t.get(2) == [(0x3000, 64)], "notes file under the preceding collective, capped at the budget") - m.begin_forward() + m.begin_forward("target") m._ordinal = 1 m.note_consumer([_T(0x9000, 8)]) m.end_forward() @@ -7073,6 +7088,14 @@ def data_ptr(self): "a forward with fewer notes (prefill-shaped) does not replace the " "table") check(m._in_forward is False, "end_forward closes the window") + # the drafter's forward learns into ITS table; the target's is untouched + m.begin_forward("drafter") + m._ordinal = 1 + m.note_consumer([_T(0x7000, 8)]) + m.end_forward() + check(m.prefetch_hint_table("drafter") == {1: [(0x7000, 8)]} + and m.prefetch_hint_table("target") == t, + "the drafter's collectives keep a table of their own") print(" osar prefetch hints contract .. OK") From 7dc5edf7127793a1c1b9bcf4667f1f31f4b0c1d4 Mon Sep 17 00:00:00 2001 From: choiceoh Date: Fri, 4 Sep 2026 17:28:43 +0900 Subject: [PATCH 3/3] =?UTF-8?q?feat(drafter):=20fc=20=EB=A5=BC=20=ED=83=80?= =?UTF-8?q?=EA=B9=83=20=ED=97=A4=EB=93=9C=C2=B7=EC=83=98=ED=94=8C=EB=9F=AC?= =?UTF-8?q?=20=EC=95=84=EB=9E=98=20side=20stream=20=EC=9C=BC=EB=A1=9C=20(g?= =?UTF-8?q?lm53=5Fdflash=5Fearly=5Ffc)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 드래프터의 첫 GEMM fc 는 타깃 forward 가 끝나는 순간 입력(aux 은닉)이 다 있는데 stock 은 propose() 안, 샘플러 뒤에 계산한다. 그 사이는 DRAM 이 논다(AllGather 는 패브릭, 샘플러는 소형 커널). 생산자 = GPUModelRunner.execute_model 래퍼(forward 뒤 side stream 에서 cat + fc 를 영속 버퍼로, 이벤트 기록), 소비자 = DFlash2Qwen3ForCausalLM.combine_hidden_states(같은 토큰 수의 대기 결과만, 이벤트 대기 뒤; precompute_and_store_context_kv 와 드래프터 그래프보다 앞이라 MK 발사끼리 겹치지 않음). 수치 동일. 노브 VLLM_GLM53_DFLASH_EARLY_FC(기본 0), 생산자 실패는 부팅 동안 자동 해제. 상한 ~0.3 ms/스텝, EXP-10 위에서만 의미. RUNBOOK EXP-15(이것), EXP-16(드래프터 메가커널: 정정 상한 −0.8~1.0 ms, 착수 전 승인). test_logic 44,558 checks. Co-Authored-By: Claude Fable 5.1 --- RUNBOOK_KERNEL_CAMPAIGN2.md | 25 +++ bench/bracket.py | 2 +- build/glm53/glm53_dflash_early_fc.py | 166 ++++++++++++++++++ build/glm53/glm5next_model.py | 15 ++ build/glm53/manifest.tsv | 1 + build/glm53/qwen3_dflash2.py | 32 ++++ .../glm53_dflash2_fp8_head/qwen3_dflash2.py | 32 ++++ .../modules/glm53_dflash_early_fc/README.md | 24 +++ .../glm53_dflash_early_fc.py | 166 ++++++++++++++++++ .../glm53_dflash_early_fc/manifest.tsv | 2 + .../modules/glm53_dflash_early_fc/requires | 2 + .../glm53_model_wiring/glm5next_model.py | 15 ++ profiles/glm53.env | 8 +- tests/test_logic.py | 135 +++++++++++++- 14 files changed, 622 insertions(+), 3 deletions(-) create mode 100644 build/glm53/glm53_dflash_early_fc.py create mode 100644 overlay/modules/glm53_dflash_early_fc/README.md create mode 100644 overlay/modules/glm53_dflash_early_fc/glm53_dflash_early_fc.py create mode 100644 overlay/modules/glm53_dflash_early_fc/manifest.tsv create mode 100644 overlay/modules/glm53_dflash_early_fc/requires diff --git a/RUNBOOK_KERNEL_CAMPAIGN2.md b/RUNBOOK_KERNEL_CAMPAIGN2.md index ebf844b1..96b15c20 100644 --- a/RUNBOOK_KERNEL_CAMPAIGN2.md +++ b/RUNBOOK_KERNEL_CAMPAIGN2.md @@ -540,6 +540,31 @@ b12x 래퍼(같은 기하, 같은 디스패치 오버레이; C=1 은 64 pairs 읽기 전용 참조(torch sum) 245 GB/s 와의 20% 는 두 커널 공통의 발사 안 구조 몫이라, 재개 조건은 "persistent 전문가 타일 스트림 마이크로커널이 240 에 닿는다" 는 2~3일 프로브의 양성이다. 보류. +## EXP-15 — 드래프터 fc 를 타깃 헤드·샘플러 아래로 (`glm53_dflash_early_fc`, 2026-09-04 추가) + +디코드 꼬리는 타깃 헤드 → 로짓 AllGather → 거부 샘플러 → 드래프터 순인데, 드래프터의 첫 +GEMM `fc`(aux 은닉 [토큰, 5×4096] → 4096; 레인에서 301 µs)는 타깃 forward 가 끝나는 순간 +입력이 다 있다. stock 은 `propose()` 안에서 샘플러 뒤에 계산한다. 그 사이 구간(AllGather 는 +패브릭, 샘플러는 소형 커널)은 DRAM 이 놀아 fc 가 공짜로 흐른다. 생산자 = `GPUModelRunner. +execute_model` 래퍼(forward 뒤 side stream 에서 cat + fc, 영속 버퍼 + 이벤트), 소비자 = +드래프터 오버레이의 `combine_hidden_states`(같은 토큰 수의 대기 결과만, 이벤트 대기 뒤). 소비가 +`precompute_and_store_context_kv` 와 드래프터 그래프보다 앞이라 MK 발사끼리 겹치지 않는다. + +- 수치 동일(같은 커널·같은 입력). 노브 `VLLM_GLM53_DFLASH_EARLY_FC=1`, 기본 0, 생산자 실패 시 + 부팅 동안 자동 해제. +- 상한 ~0.3 ms/스텝(27차 인구조사: fc 만 그 시점에 입력이 준비된 꼬리 GEMM). EXP-10 위에서만 + 의미(fc 가 레인에 있어야 함). 단독 판정 불가 — EXP-10 브래킷의 cand 에 얹는다. + +## EXP-16 — 드래프터 메가커널 (제안, 착수 전 승인; 2026-09-04) + +27차 인구조사(깨끗한 디코드 스텝): 꼬리 136 커널 중 GEMM 33(EXP-10 뒤 MK 발사 ~31), AR 11, +`kernel_mha` 5, 글루 ~50개 0.33 ms, 발사 간극 ~0.35 ms. 융합으로 없앨 수 있는 것은 글루 + 간극 + +MK 발사 고정비(30 × ~10 µs) ≈ **−0.8~1.0 ms(1.2~1.5%)**; AR 0.79 ms 와 mha 는 남는다. 브래킷 +해상도(CV 1.7%) 아래라 단독으로는 판정할 수 없고, 공사는 MK-KDA 급(2주). 설계는 층당 두 발사: +[norm → conv.prepare → qkv GEMM] 과 [conv.finish → post-norm → mlp_conv.prepare → gate_up → +act → mlp_conv.finish → down], 어텐션(`kernel_mha`)은 stock 유지. **운영자 승인 뒤 착수** — 정정된 +상한을 본 뒤의 결정이어야 한다. + ## 브래킷 자동화 — `bench/bracket.py` (도구, 판정 아님) `leg`(살아있는 서버에 rep 기록) + `judge`(기록 판정) 2중 명령. 원장 규율을 코드로 diff --git a/bench/bracket.py b/bench/bracket.py index d4a89510..76b071fc 100644 --- a/bench/bracket.py +++ b/bench/bracket.py @@ -39,7 +39,7 @@ "VLLM_GLM53_PREP_FUSED", "VLLM_GLM53_ASYNC_DFLASH", "VLLM_GLM53_MHC_SMALLM", "VLLM_DFLASH_PREP_WARMUP", "VLLM_GLM53_MK_PDL", "VLLM_GLM53_MK_KSR_OUT", - "VLLM_GLM53_AR_PREFETCH", + "VLLM_GLM53_AR_PREFETCH", "VLLM_GLM53_DFLASH_EARLY_FC", ) if k in os.environ] + sorted(k for k in os.environ if k.startswith(("VLLM_GLM53_MK_", "VLLM_GLM53_KPOOL"))) diff --git a/build/glm53/glm53_dflash_early_fc.py b/build/glm53/glm53_dflash_early_fc.py new file mode 100644 index 00000000..2486de71 --- /dev/null +++ b/build/glm53/glm53_dflash_early_fc.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DFlash drafter: run the aux-hidden projection (`fc`) under the target's +head + sampler instead of after them (VLLM_GLM53_DFLASH_EARLY_FC). + +The decode step's tail (MEASUREMENTS 25차/27차) is target head -> logits +AllGather -> rejection sampler -> drafter. The drafter's first GEMM, `fc` +([tokens, 5 x 4096] -> [tokens, 4096], 168 MB bf16 / 301 us on the MK W4 +lane), only needs the target's aux hidden states, which exist the moment the +target forward returns -- but the stock speculator computes it inside +`propose()`, after the sampler. Everything between is not DRAM-bound (the +AllGather is the fabric, the sampler is small kernels), so the fc can stream +its weights there for free. + +Two halves, both above any compiled or captured region: + + producer a wrapper on `GPUModelRunner.execute_model`: once the target + forward has returned (execute_model_state carries the aux hidden + states), a side stream cats the aux states and runs the drafter's + `fc` into a persistent buffer, and records an event + consumer `DFlash2Qwen3ForCausalLM.combine_hidden_states` (the drafter + overlay): if a result is pending for this step and this token + count, wait on the event and return the buffer; otherwise the + stock computation runs (so a missed step costs nothing but the + overlap) + +The consumer waits BEFORE `precompute_and_store_context_kv` and before the +drafter graph, so no megakernel launch ever runs concurrently with the fc's +own MK launch (the lane's monotonic ticket barrier assumes one launch at a +time). Numerics are identical: same kernel, same inputs, one stream apart. + +Knob: VLLM_GLM53_DFLASH_EARLY_FC=1 (exact). Anything else is off, and the +installer then does not touch the runner at all. +""" +from __future__ import annotations + +import os + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_INSTALLED = False +_ORIG_EXECUTE_MODEL = None +_DISABLED = False +_SEQ = 0 +_STREAM = None + + +def early_fc_enabled() -> bool: + return (os.environ.get("VLLM_GLM53_DFLASH_EARLY_FC") or "0").strip() == "1" + + +def _drafter_of(runner): + """The drafter ForCausalLM that owns `fc`, or None.""" + spec = getattr(runner, "speculator", None) + model = getattr(spec, "model", None) + inner = getattr(model, "model", None) + if inner is None or getattr(inner, "fc", None) is None: + return None + if not getattr(inner, "use_aux_hidden_state", False): + return None + if not hasattr(model, "combine_hidden_states"): + return None + return model + + +def launch_early_fc(runner) -> bool: + """Producer. Called after execute_model; returns whether it launched.""" + global _SEQ, _STREAM + import torch + + st = getattr(runner, "execute_model_state", None) + if st is None: + return False + aux = getattr(st, "aux_hidden_states", None) + if not aux: + return False + drafter = _drafter_of(runner) + if drafter is None: + return False + input_batch = getattr(st, "input_batch", None) + n = int(getattr(input_batch, "num_tokens", 0) or 0) + if n <= 0 or n > aux[0].shape[0]: + return False + fc = drafter.model.fc + width = sum(int(a.shape[-1]) for a in aux) + if width != int(fc.input_size): + return False # a draft/target pair the stock path would reject too + cur = torch.cuda.current_stream() + if _STREAM is None: + _STREAM = torch.cuda.Stream() + cat_buf = getattr(drafter, "_deneb_early_fc_cat", None) + if cat_buf is None or cat_buf.shape[0] < aux[0].shape[0]: + rows = int(aux[0].shape[0]) + cat_buf = torch.empty(rows, width, dtype=aux[0].dtype, device=aux[0].device) + out_buf = torch.empty(rows, int(fc.output_size), dtype=aux[0].dtype, + device=aux[0].device) + drafter._deneb_early_fc_cat = cat_buf + drafter._deneb_early_fc_out = out_buf + drafter._deneb_early_fc_event = torch.cuda.Event() + out_buf = drafter._deneb_early_fc_out + _STREAM.wait_stream(cur) + with torch.cuda.stream(_STREAM): + torch.cat([a[:n] for a in aux], dim=-1, out=cat_buf[:n]) + h = fc(cat_buf[:n]) + if isinstance(h, tuple): + h = h[0] + out_buf[:n].copy_(h) + drafter._deneb_early_fc_event.record(_STREAM) + _SEQ += 1 + drafter._deneb_early_fc_pending = (n, _SEQ) + return True + + +def take_early_fc(drafter, num_tokens: int): + """Consumer. The pending buffer for this step, or None (stock path).""" + import torch + + pend = getattr(drafter, "_deneb_early_fc_pending", None) + if pend is None: + return None + drafter._deneb_early_fc_pending = None # consumed once, whatever happens + n, _seq = pend + if n != num_tokens: + return None + torch.cuda.current_stream().wait_event(drafter._deneb_early_fc_event) + return drafter._deneb_early_fc_out[:n] + + +def _patched_execute_model(self, *args, **kwargs): + global _DISABLED + out = _ORIG_EXECUTE_MODEL(self, *args, **kwargs) + if not _DISABLED: + try: + launch_early_fc(self) + except Exception: + _DISABLED = True + logger.exception("[dflash-early-fc] producer failed -> stock path " + "for the rest of the boot") + return out + + +def install_glm53_dflash_early_fc() -> bool: + """Wrap the runner's execute_model once. Safe to call from model import.""" + global _INSTALLED, _ORIG_EXECUTE_MODEL + if _INSTALLED or not early_fc_enabled(): + return _INSTALLED + import sys + + mr = sys.modules.get("vllm.v1.worker.gpu.model_runner") + if mr is None: + try: + import vllm.v1.worker.gpu.model_runner as mr # noqa: F811 + except Exception: + logger.exception("[dflash-early-fc] runner module not importable -> off") + return False + Runner = getattr(mr, "GPUModelRunner", None) + if Runner is None or not hasattr(Runner, "execute_model"): + logger.warning("[dflash-early-fc] GPUModelRunner.execute_model missing -> off") + return False + _ORIG_EXECUTE_MODEL = Runner.execute_model + Runner.execute_model = _patched_execute_model + _INSTALLED = True + logger.warning("[dflash-early-fc] installed: the drafter's fc runs under the " + "target head + sampler on a side stream") + return True diff --git a/build/glm53/glm5next_model.py b/build/glm53/glm5next_model.py index 30750d1d..038ddcad 100644 --- a/build/glm53/glm5next_model.py +++ b/build/glm53/glm5next_model.py @@ -133,6 +133,21 @@ except Exception: logger.exception("[prep-fused] install failed -> stock path") +# deneb fork (glm53_dflash_early_fc): the drafter's fc under the target head +# + sampler. Same shape of install as prep_fused: import only, inert unless +# VLLM_GLM53_DFLASH_EARLY_FC=1, loud on failure, never fatal. +try: + from .glm53_dflash_early_fc import install_glm53_dflash_early_fc +except ImportError as _e: + install_glm53_dflash_early_fc = None + if _e.name != f"{__package__}.glm53_dflash_early_fc": + logger.exception("[dflash-early-fc] module import failed -> stock path") +if install_glm53_dflash_early_fc is not None: + try: + install_glm53_dflash_early_fc() + except Exception: + logger.exception("[dflash-early-fc] install failed -> stock path") + def _validate_decodable_vocab_bound( decodable_vocab: int, diff --git a/build/glm53/manifest.tsv b/build/glm53/manifest.tsv index 96f8b0a4..de3f608a 100644 --- a/build/glm53/manifest.tsv +++ b/build/glm53/manifest.tsv @@ -33,3 +33,4 @@ glm5next_attention.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5ne glm53_indexer_gate.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/glm53_indexer_gate.py absent deneb_boot_stamps.py /usr/local/lib/python3.12/dist-packages/deneb_boot_stamps.py absent zz_deneb_boot_stamps.pth /usr/local/lib/python3.12/dist-packages/zz_deneb_boot_stamps.pth absent +glm53_dflash_early_fc.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/glm53_dflash_early_fc.py absent diff --git a/build/glm53/qwen3_dflash2.py b/build/glm53/qwen3_dflash2.py index e66c46c9..9c615c4b 100644 --- a/build/glm53/qwen3_dflash2.py +++ b/build/glm53/qwen3_dflash2.py @@ -271,6 +271,25 @@ def forward( +_EARLY_FC = None + + +def _early_fc_take(): + """`take_early_fc` of glm53_dflash_early_fc, resolved once; None when that + module is not mounted (stock projection every step).""" + global _EARLY_FC + if _EARLY_FC is None: + try: + from vllm.models.glm5next.nvidia.glm53_dflash_early_fc import ( + take_early_fc, + ) + + _EARLY_FC = take_early_fc + except Exception: + _EARLY_FC = False + return _EARLY_FC or None + + _OSAR = None @@ -394,6 +413,19 @@ def forward( if osar is not None: osar.end_forward() + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + # deneb fork (glm53_dflash_early_fc): the fc may already have run on a + # side stream under the target's head + sampler. Take that result for + # this step's token count, else the stock projection. Waits on the + # producer's event first, before precompute_and_store_context_kv and + # the drafter graph -- no megakernel launch overlaps the fc's. + early = _early_fc_take() + if early is not None and hidden_states.dim() == 2: + got = early(self, int(hidden_states.shape[0])) + if got is not None: + return got + return super().combine_hidden_states(hidden_states) + def verify_selector_loaded(self) -> None: """Say out loud whether the path selector's weights actually arrived. diff --git a/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py b/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py index e66c46c9..9c615c4b 100644 --- a/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py +++ b/overlay/modules/glm53_dflash2_fp8_head/qwen3_dflash2.py @@ -271,6 +271,25 @@ def forward( +_EARLY_FC = None + + +def _early_fc_take(): + """`take_early_fc` of glm53_dflash_early_fc, resolved once; None when that + module is not mounted (stock projection every step).""" + global _EARLY_FC + if _EARLY_FC is None: + try: + from vllm.models.glm5next.nvidia.glm53_dflash_early_fc import ( + take_early_fc, + ) + + _EARLY_FC = take_early_fc + except Exception: + _EARLY_FC = False + return _EARLY_FC or None + + _OSAR = None @@ -394,6 +413,19 @@ def forward( if osar is not None: osar.end_forward() + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + # deneb fork (glm53_dflash_early_fc): the fc may already have run on a + # side stream under the target's head + sampler. Take that result for + # this step's token count, else the stock projection. Waits on the + # producer's event first, before precompute_and_store_context_kv and + # the drafter graph -- no megakernel launch overlaps the fc's. + early = _early_fc_take() + if early is not None and hidden_states.dim() == 2: + got = early(self, int(hidden_states.shape[0])) + if got is not None: + return got + return super().combine_hidden_states(hidden_states) + def verify_selector_loaded(self) -> None: """Say out loud whether the path selector's weights actually arrived. diff --git a/overlay/modules/glm53_dflash_early_fc/README.md b/overlay/modules/glm53_dflash_early_fc/README.md new file mode 100644 index 00000000..f1ff1016 --- /dev/null +++ b/overlay/modules/glm53_dflash_early_fc/README.md @@ -0,0 +1,24 @@ +# glm53_dflash_early_fc + +The DFlash drafter's `fc` (aux hidden states [tokens, 5 x 4096] -> [tokens, +4096]; 168 MB bf16, 301 us on the MK W4 lane) needs only the target's aux +hidden states, which exist when the target forward returns -- but the stock +speculator computes it inside `propose()`, after the head GEMM, the logits +AllGather and the rejection sampler. That window is not DRAM-bound, so the fc +streams its weights there for free (ceiling ~0.3 ms/step, MEASUREMENTS 27차 +census: fc is the only tail GEMM whose inputs are ready that early). + +- producer: a wrapper on `GPUModelRunner.execute_model` (installed from the + GLM model module import, like `glm53_prep_fused`); after the forward it + cats the aux states and runs the drafter's `fc` on a side stream into a + persistent buffer, recording an event +- consumer: `DFlash2Qwen3ForCausalLM.combine_hidden_states` (the drafter + overlay) takes the pending buffer for this step and token count after + waiting on the event, else runs the stock computation + +The consumer runs before `precompute_and_store_context_kv` and the drafter +graph, so the fc's MK launch never overlaps another megakernel launch (the +lane's ticket barrier is one-launch-at-a-time). Numerics identical: same +kernel, same inputs. Knob `VLLM_GLM53_DFLASH_EARLY_FC=1`, default 0; any +producer failure disables it for the boot and logs. Speed only: bracket on +C=1 step/s stacked on the EXP-10 arm (the fc must be on the lane to matter). diff --git a/overlay/modules/glm53_dflash_early_fc/glm53_dflash_early_fc.py b/overlay/modules/glm53_dflash_early_fc/glm53_dflash_early_fc.py new file mode 100644 index 00000000..2486de71 --- /dev/null +++ b/overlay/modules/glm53_dflash_early_fc/glm53_dflash_early_fc.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DFlash drafter: run the aux-hidden projection (`fc`) under the target's +head + sampler instead of after them (VLLM_GLM53_DFLASH_EARLY_FC). + +The decode step's tail (MEASUREMENTS 25차/27차) is target head -> logits +AllGather -> rejection sampler -> drafter. The drafter's first GEMM, `fc` +([tokens, 5 x 4096] -> [tokens, 4096], 168 MB bf16 / 301 us on the MK W4 +lane), only needs the target's aux hidden states, which exist the moment the +target forward returns -- but the stock speculator computes it inside +`propose()`, after the sampler. Everything between is not DRAM-bound (the +AllGather is the fabric, the sampler is small kernels), so the fc can stream +its weights there for free. + +Two halves, both above any compiled or captured region: + + producer a wrapper on `GPUModelRunner.execute_model`: once the target + forward has returned (execute_model_state carries the aux hidden + states), a side stream cats the aux states and runs the drafter's + `fc` into a persistent buffer, and records an event + consumer `DFlash2Qwen3ForCausalLM.combine_hidden_states` (the drafter + overlay): if a result is pending for this step and this token + count, wait on the event and return the buffer; otherwise the + stock computation runs (so a missed step costs nothing but the + overlap) + +The consumer waits BEFORE `precompute_and_store_context_kv` and before the +drafter graph, so no megakernel launch ever runs concurrently with the fc's +own MK launch (the lane's monotonic ticket barrier assumes one launch at a +time). Numerics are identical: same kernel, same inputs, one stream apart. + +Knob: VLLM_GLM53_DFLASH_EARLY_FC=1 (exact). Anything else is off, and the +installer then does not touch the runner at all. +""" +from __future__ import annotations + +import os + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_INSTALLED = False +_ORIG_EXECUTE_MODEL = None +_DISABLED = False +_SEQ = 0 +_STREAM = None + + +def early_fc_enabled() -> bool: + return (os.environ.get("VLLM_GLM53_DFLASH_EARLY_FC") or "0").strip() == "1" + + +def _drafter_of(runner): + """The drafter ForCausalLM that owns `fc`, or None.""" + spec = getattr(runner, "speculator", None) + model = getattr(spec, "model", None) + inner = getattr(model, "model", None) + if inner is None or getattr(inner, "fc", None) is None: + return None + if not getattr(inner, "use_aux_hidden_state", False): + return None + if not hasattr(model, "combine_hidden_states"): + return None + return model + + +def launch_early_fc(runner) -> bool: + """Producer. Called after execute_model; returns whether it launched.""" + global _SEQ, _STREAM + import torch + + st = getattr(runner, "execute_model_state", None) + if st is None: + return False + aux = getattr(st, "aux_hidden_states", None) + if not aux: + return False + drafter = _drafter_of(runner) + if drafter is None: + return False + input_batch = getattr(st, "input_batch", None) + n = int(getattr(input_batch, "num_tokens", 0) or 0) + if n <= 0 or n > aux[0].shape[0]: + return False + fc = drafter.model.fc + width = sum(int(a.shape[-1]) for a in aux) + if width != int(fc.input_size): + return False # a draft/target pair the stock path would reject too + cur = torch.cuda.current_stream() + if _STREAM is None: + _STREAM = torch.cuda.Stream() + cat_buf = getattr(drafter, "_deneb_early_fc_cat", None) + if cat_buf is None or cat_buf.shape[0] < aux[0].shape[0]: + rows = int(aux[0].shape[0]) + cat_buf = torch.empty(rows, width, dtype=aux[0].dtype, device=aux[0].device) + out_buf = torch.empty(rows, int(fc.output_size), dtype=aux[0].dtype, + device=aux[0].device) + drafter._deneb_early_fc_cat = cat_buf + drafter._deneb_early_fc_out = out_buf + drafter._deneb_early_fc_event = torch.cuda.Event() + out_buf = drafter._deneb_early_fc_out + _STREAM.wait_stream(cur) + with torch.cuda.stream(_STREAM): + torch.cat([a[:n] for a in aux], dim=-1, out=cat_buf[:n]) + h = fc(cat_buf[:n]) + if isinstance(h, tuple): + h = h[0] + out_buf[:n].copy_(h) + drafter._deneb_early_fc_event.record(_STREAM) + _SEQ += 1 + drafter._deneb_early_fc_pending = (n, _SEQ) + return True + + +def take_early_fc(drafter, num_tokens: int): + """Consumer. The pending buffer for this step, or None (stock path).""" + import torch + + pend = getattr(drafter, "_deneb_early_fc_pending", None) + if pend is None: + return None + drafter._deneb_early_fc_pending = None # consumed once, whatever happens + n, _seq = pend + if n != num_tokens: + return None + torch.cuda.current_stream().wait_event(drafter._deneb_early_fc_event) + return drafter._deneb_early_fc_out[:n] + + +def _patched_execute_model(self, *args, **kwargs): + global _DISABLED + out = _ORIG_EXECUTE_MODEL(self, *args, **kwargs) + if not _DISABLED: + try: + launch_early_fc(self) + except Exception: + _DISABLED = True + logger.exception("[dflash-early-fc] producer failed -> stock path " + "for the rest of the boot") + return out + + +def install_glm53_dflash_early_fc() -> bool: + """Wrap the runner's execute_model once. Safe to call from model import.""" + global _INSTALLED, _ORIG_EXECUTE_MODEL + if _INSTALLED or not early_fc_enabled(): + return _INSTALLED + import sys + + mr = sys.modules.get("vllm.v1.worker.gpu.model_runner") + if mr is None: + try: + import vllm.v1.worker.gpu.model_runner as mr # noqa: F811 + except Exception: + logger.exception("[dflash-early-fc] runner module not importable -> off") + return False + Runner = getattr(mr, "GPUModelRunner", None) + if Runner is None or not hasattr(Runner, "execute_model"): + logger.warning("[dflash-early-fc] GPUModelRunner.execute_model missing -> off") + return False + _ORIG_EXECUTE_MODEL = Runner.execute_model + Runner.execute_model = _patched_execute_model + _INSTALLED = True + logger.warning("[dflash-early-fc] installed: the drafter's fc runs under the " + "target head + sampler on a side stream") + return True diff --git a/overlay/modules/glm53_dflash_early_fc/manifest.tsv b/overlay/modules/glm53_dflash_early_fc/manifest.tsv new file mode 100644 index 00000000..d2f0c663 --- /dev/null +++ b/overlay/modules/glm53_dflash_early_fc/manifest.tsv @@ -0,0 +1,2 @@ +# source container_target(relative to TARGET_PREFIX) base_preimage_sha256 +glm53_dflash_early_fc.py vllm/models/glm5next/nvidia/glm53_dflash_early_fc.py absent diff --git a/overlay/modules/glm53_dflash_early_fc/requires b/overlay/modules/glm53_dflash_early_fc/requires new file mode 100644 index 00000000..7b14be4f --- /dev/null +++ b/overlay/modules/glm53_dflash_early_fc/requires @@ -0,0 +1,2 @@ +glm53_dflash2_fp8_head +glm53_model_wiring diff --git a/overlay/modules/glm53_model_wiring/glm5next_model.py b/overlay/modules/glm53_model_wiring/glm5next_model.py index 30750d1d..038ddcad 100644 --- a/overlay/modules/glm53_model_wiring/glm5next_model.py +++ b/overlay/modules/glm53_model_wiring/glm5next_model.py @@ -133,6 +133,21 @@ except Exception: logger.exception("[prep-fused] install failed -> stock path") +# deneb fork (glm53_dflash_early_fc): the drafter's fc under the target head +# + sampler. Same shape of install as prep_fused: import only, inert unless +# VLLM_GLM53_DFLASH_EARLY_FC=1, loud on failure, never fatal. +try: + from .glm53_dflash_early_fc import install_glm53_dflash_early_fc +except ImportError as _e: + install_glm53_dflash_early_fc = None + if _e.name != f"{__package__}.glm53_dflash_early_fc": + logger.exception("[dflash-early-fc] module import failed -> stock path") +if install_glm53_dflash_early_fc is not None: + try: + install_glm53_dflash_early_fc() + except Exception: + logger.exception("[dflash-early-fc] install failed -> stock path") + def _validate_decodable_vocab_bound( decodable_vocab: int, diff --git a/profiles/glm53.env b/profiles/glm53.env index 560b7e84..adf53d67 100644 --- a/profiles/glm53.env +++ b/profiles/glm53.env @@ -60,7 +60,7 @@ TARGET_PREFIX="/usr/local/lib/python3.12/dist-packages/" # The image runs V2 Model Runner. glm53_drop_audit and glm53_sparse_q replace # V1-only paths and are deliberately absent: mounting them advertised guards # and sparse-q behavior that the live DFlash2 path never executed. -MODULES="b12x_shared_workspace b12x_zero_weight_micro glm53_tail_slot_persistent glm53_kpool_tail_select glm53_v2_sampler_guards moe_gate_sm121 glm53_model_wiring tp_oneshot_ar glm53_oneshot_wiring fp8_lm_head glm53_dflash2_fp8_head glm53_dflash_loader_fp8 glm53_fp8_dense glm53_b12x_out glm53_mhc_tilelang glm53_sm121_mla_prefill glm53_kda_prefill_regime glm53_dflash_warmup glm53_megakernel glm53_mk_mla_wiring glm53_mk_kda_wiring glm53_prep_fused glm53_indexer_gate_splitk glm53_boot_stamps" +MODULES="b12x_shared_workspace b12x_zero_weight_micro glm53_tail_slot_persistent glm53_kpool_tail_select glm53_v2_sampler_guards moe_gate_sm121 glm53_model_wiring tp_oneshot_ar glm53_oneshot_wiring fp8_lm_head glm53_dflash2_fp8_head glm53_dflash_loader_fp8 glm53_fp8_dense glm53_b12x_out glm53_mhc_tilelang glm53_sm121_mla_prefill glm53_kda_prefill_regime glm53_dflash_warmup glm53_megakernel glm53_mk_mla_wiring glm53_mk_kda_wiring glm53_prep_fused glm53_indexer_gate_splitk glm53_boot_stamps glm53_dflash_early_fc" # --- serving knobs ----------------------------------------------------------- # Dense-MHA for GLM's short/full-attention sparse-MLA prefill region on SM121. @@ -206,6 +206,12 @@ VLLM_GLM53_MK_PDL=1 # collective, N = N MB (1..20). Speed only; ceiling ~2-3 pct/step # (strategy doc 2026-09-04), bracket on the fleet like any arm. VLLM_GLM53_AR_PREFETCH=0 +# glm53_dflash_early_fc: the drafter's fc (aux hidden -> 4096; 301 us on the +# lane) runs on a side stream under the target head + AllGather + sampler +# instead of after them. Same kernel, same inputs -> identical bytes; speed +# only, ceiling ~0.3 ms/step, matters when EXP-10 puts the fc on the lane. +# Exact 1 arms; the producer disables itself for the boot on any failure. +VLLM_GLM53_DFLASH_EARLY_FC=0 # MK_SEG_MLA: our own sparse MLA decode kernel (bf16 mma; see the segment # comment in glm53_megakernel.cu). glm53_mk_mla_wiring routes decode # shapes (T<=32) to it; isolated 1.33x the FlashInfer wrapper at C=1. A diff --git a/tests/test_logic.py b/tests/test_logic.py index 9533541f..0ca6ad1a 100644 --- a/tests/test_logic.py +++ b/tests/test_logic.py @@ -6995,7 +6995,7 @@ def test_osar_prefetch_hints_contract() -> None: "glm5next_model.py"), encoding="utf-8").read() cls_at = wiring.index("class Glm5NextForConditionalGeneration(") layer_at = wiring.index("class Glm5NextDecoderLayer(") - b_at = wiring.index("osar.begin_forward()") + b_at = wiring.index('osar.begin_forward("target")') check(b_at > cls_at and "osar.end_forward()" in wiring[b_at:] and "def _osar_shim():" in wiring and 'osar.begin_forward("target")' in wiring, @@ -7099,6 +7099,138 @@ def data_ptr(self): print(" osar prefetch hints contract .. OK") +def test_glm53_dflash_early_fc_contracts() -> None: + """glm53_dflash_early_fc: the drafter's fc under the target head + sampler. + + Producer = a wrapper on GPUModelRunner.execute_model installed from the + GLM model import (inert unless VLLM_GLM53_DFLASH_EARLY_FC=1); consumer = + the drafter's combine_hidden_states, which takes a pending result only + for this step's token count and waits on the producer's event first -- + before precompute_and_store_context_kv and the drafter graph, so the fc's + megakernel launch never overlaps another one.""" + mod_dir = os.path.join(REPO, "overlay", "modules", "glm53_dflash_early_fc") + src = open(os.path.join(mod_dir, "glm53_dflash_early_fc.py"), encoding="utf-8").read() + profile = open(os.path.join(REPO, "profiles", "glm53.env"), encoding="utf-8").read() + modules = re.search(r'^MODULES="([^"]+)"', profile, re.M).group(1).split() + check("glm53_dflash_early_fc" in modules, "glm53 profile mounts glm53_dflash_early_fc") + check(re.search(r"^VLLM_GLM53_DFLASH_EARLY_FC=0$", profile, re.M) is not None, + "the knob ships off") + rows = [l.split("\t") for l in open(os.path.join(mod_dir, "manifest.tsv"), encoding="utf-8") + .read().splitlines() if l and not l.startswith("#")] + check(rows == [["glm53_dflash_early_fc.py", + "vllm/models/glm5next/nvidia/glm53_dflash_early_fc.py", "absent"]], + f"manifest binds the module as a new file next to the model: {rows}") + req = open(os.path.join(mod_dir, "requires"), encoding="utf-8").read().split() + check({"glm53_dflash2_fp8_head", "glm53_model_wiring"} <= set(req), + "requires names the drafter overlay (consumer) and the wiring (installer)") + check('(os.environ.get("VLLM_GLM53_DFLASH_EARLY_FC") or "0").strip() == "1"' in src + and "def install_glm53_dflash_early_fc() -> bool:" in src + and "Runner.execute_model = _patched_execute_model" in src + and "_ORIG_EXECUTE_MODEL = Runner.execute_model" in src, + "exact-1 knob; the installer wraps execute_model and keeps the original") + prod = src[src.index("def launch_early_fc("):src.index("def take_early_fc(")] + check("_STREAM.wait_stream(cur)" in prod and "with torch.cuda.stream(_STREAM):" in prod + and "drafter._deneb_early_fc_event.record(_STREAM)" in prod + and "torch.cat([a[:n] for a in aux], dim=-1, out=cat_buf[:n])" in prod + and "width != int(fc.input_size)" in prod, + "the producer runs on a side stream after the forward, into a persistent " + "buffer, and refuses a width the stock path would refuse") + cons = src[src.index("def take_early_fc("):src.index("def _patched_execute_model(")] + check("drafter._deneb_early_fc_pending = None # consumed once" in cons + and "if n != num_tokens:\n return None" in cons + and "torch.cuda.current_stream().wait_event(drafter._deneb_early_fc_event)" in cons, + "the consumer takes a result once, only for this token count, after the event") + patched = src[src.index("def _patched_execute_model("):src.index("def install_glm53_dflash_early_fc(")] + check("out = _ORIG_EXECUTE_MODEL(self, *args, **kwargs)" in patched + and "_DISABLED = True" in patched and "return out" in patched, + "a producer failure disables the arm for the boot and never breaks the step") + wiring = open(os.path.join(REPO, "overlay/modules/glm53_model_wiring/" + "glm5next_model.py"), encoding="utf-8").read() + check("from .glm53_dflash_early_fc import install_glm53_dflash_early_fc" in wiring + and 'if _e.name != f"{__package__}.glm53_dflash_early_fc":' in wiring + and wiring.index("install_glm53_dflash_early_fc()") > wiring.index("install_glm53_prep_fused()"), + "installed from the wiring like prep_fused: silent without the module, loud when broken") + drafter = open(os.path.join(REPO, "overlay/modules/glm53_dflash2_fp8_head/" + "qwen3_dflash2.py"), encoding="utf-8").read() + body = drafter[drafter.index("def combine_hidden_states(self, hidden_states"):] + body = body[:body.index("def verify_selector_loaded")] + check("early = _early_fc_take()" in body + and "return super().combine_hidden_states(hidden_states)" in body + and "from vllm.models.glm5next.nvidia.glm53_dflash_early_fc import" in drafter, + "the drafter overlay consumes the pending result and falls back to stock") + bracket = open(os.path.join(REPO, "bench", "bracket.py"), encoding="utf-8").read() + check('"VLLM_GLM53_DFLASH_EARLY_FC"' in bracket, "bracket.py snapshots the knob") + + # the producer/consumer protocol on fakes: taken once, only for the same + # token count, and never when nothing is pending + import importlib.util + spec = importlib.util.spec_from_file_location( + "_early_fc_mod", os.path.join(mod_dir, "glm53_dflash_early_fc.py")) + import types + fake_vllm = types.ModuleType("vllm") + fake_logger = types.ModuleType("vllm.logger") + fake_logger.init_logger = lambda name: _CapturingLogger() + fake_vllm.logger = fake_logger + saved = {k: sys.modules.get(k) for k in ("vllm", "vllm.logger")} + sys.modules["vllm"] = fake_vllm + sys.modules["vllm.logger"] = fake_logger + try: + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + class _Ev: + pass + + class _Stream: + def wait_event(self, ev): + self.waited = ev + + class _Torch: + class cuda: + _s = _Stream() + + @staticmethod + def current_stream(): + return _Torch.cuda._s + + mod.torch = _Torch # take_early_fc imports torch lazily; give it the fake + import builtins + real_import = builtins.__import__ + + def fake_import(name, *a, **k): + if name == "torch": + return _Torch + return real_import(name, *a, **k) + + class _Drafter: + pass + + d = _Drafter() + d._deneb_early_fc_event = _Ev() + d._deneb_early_fc_out = list(range(16)) + builtins.__import__ = fake_import + try: + check(mod.take_early_fc(d, 8) is None, "nothing pending -> stock path") + d._deneb_early_fc_pending = (8, 1) + check(mod.take_early_fc(d, 5) is None and d._deneb_early_fc_pending is None, + "a token-count mismatch drops the pending result (stock path)") + d._deneb_early_fc_pending = (8, 2) + got = mod.take_early_fc(d, 8) + check(got == list(range(8)) and _Torch.cuda._s.waited is d._deneb_early_fc_event + and d._deneb_early_fc_pending is None, + "a match waits on the event, returns the first n rows, and is consumed") + check(mod.take_early_fc(d, 8) is None, "consumed once") + finally: + builtins.__import__ = real_import + print(" glm53_dflash_early_fc contracts .. OK") + + def test_kv_cache_is_pinned_in_tokens() -> None: """KV is pinned in tokens, not left to take whatever GMU leaves. @@ -9731,6 +9863,7 @@ def test_common_tp4_library_is_the_one_implementation() -> None: test_ab_runner_measures_both_channels() test_osar_wait_is_split_by_message_size() test_osar_prefetch_hints_contract() + test_glm53_dflash_early_fc_contracts() test_glm53_megakernel_contracts() test_prefill_warmup_contracts() test_megakernel_w4_layout_functional()