From c8407351d74edd37673c3aed84c8a7af0526b33b Mon Sep 17 00:00:00 2001 From: choiceoh Date: Sat, 5 Sep 2026 01:29:30 +0900 Subject: [PATCH] review(mk-gemm v2): 13 findings from the PR #305 review -- sweep-derived ksr rule, device SM count, exact reciprocal, unrolled ring copy, static epilogue indices, one ts reader; bench judges the served lane by default; trace tool picks the main stream by kernel time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kernel (mk_gemm2_kernel): - mk_choose_ksr2 from the 30차 sweeps: one exact wave of the resident slots (48 tiles x 2, 32 x 3, 16 x 6), the longest slices under one wave when the tile count does not divide it, the finest slices above half the slots (51 tiles: a 102-unit plan measured +30 us). SM count from the device (was a hard-coded 48). v2 x2 now >= v1 x2 on every m <= 16 shape. - __frcp_rn for the row scale's reciprocal (exact for a power of two; the IEEE divide compiled to a slow-path call per row per k-block), the ring copy loop unrolled (the swizzled destinations are per-thread constants), static acc indices in the epilogue (a runtime bound put acc in local memory: 64 B stack frame -> 0), one store walk with two store lambdas. ptxas: 124 registers, 0 spills, 0 B stack. - One read-and-clear template for the four stamp readers; the v2 header states the cross-stream no-overlap contract its shared buffers inherit. Bench: --gemm2 defaults to the served lane (env), 'both' is explicit (the VERDICT authorises boots); the ksr sweep measures only at the ksr the host launches (clamps otherwise mislabelled a cell) and reuses the main loop's tensors; probe_exact builds the fixture, pack and reference once for all lanes; --stamps2 takes medians over iters launches. Trace tool: main stream = the stream with the most kernel time (the 09-01 stock trace has the forward on 210 and the prep/side pair on 17, the armed trace the other way), per-stream summary line, overlap lookback = the segment's longest kernel, finite-gap medians, exposure medians over all steps with the warm-up step skipped. Probe: the MoE-overlap probe reuses moe_decode_stream_probe's fixtures and calls the wrapper without out= (the stock flashinfer wrapper has no such keyword; the bench runner mounts the megakernel files only). Ledger 30차 §4-5: first bench round, the ksr=3 sweep, the rule, the review facts. README/runbook updated. Co-Authored-By: Claude Fable 5.1 --- MEASUREMENTS.md | 67 +++++++- RUNBOOK_KERNEL_CAMPAIGN2.md | 8 +- build/dsv4/glm53_megakernel.cu | 150 +++++++++--------- build/glm53/glm53_megakernel.cu | 150 +++++++++--------- overlay/modules/glm53_megakernel/README.md | 20 ++- .../glm53_megakernel/glm53_megakernel.cu | 150 +++++++++--------- probes/megakernel_glm53_bench.py | 109 +++++++------ probes/mk_gemm_moe_overlap_probe.py | 36 ++--- tests/test_logic.py | 19 ++- tools/trace_launch_positions.py | 44 +++-- 10 files changed, 434 insertions(+), 319 deletions(-) diff --git a/MEASUREMENTS.md b/MEASUREMENTS.md index 19edc990..a6125a97 100644 --- a/MEASUREMENTS.md +++ b/MEASUREMENTS.md @@ -2020,9 +2020,70 @@ partial 은 m × n × ksr ≤ 32 × 6528 × 4. 벤치 `--gemm2 both`(v1 열 옆 기대(측정 전): 메인 99발 6.4 → ~4.9 ms(고정비 ~3 µs, 균형 꼬리 제거), aux 쌍 노출 28 → ~5 µs/층(down 이 MoE 꼬리 안에서 끝남) → 스텝 −2.5 ms, 트레이스 mk_gemm 합 14.1 → ~6.5~7. -### 4. 측정 (채워질 자리) - -_srv2 `~/mkab/g2.bench.out` / `g2.overlap.out` — 플릿이 lever-chain4 로 점유돼 GPU 창 대기._ +### 4. 측정 1차 — 단독 벤치 (2026-09-05 00:48, srv2, chain4 와 chain5 사이의 3분 창, `g2.bench.out`, iters 10) + +수치 게이트: exact grid v1 3.9e-8 / v2 ksr 0·1·2·3·5·8 전부 ≤ 3.9e-8, over-ulp 0, replay 안정 — +**PASS**. v2 출력은 비분할 형상(k=512, 4096², k=2048, k=3072, m=16 4096²)에서 v1 과 **비트 +동일**, 분할 형상은 슬라이스 합 순서 차이 1e-10~5e-6. + +단독(콜드 L2)/연속 2발(x2) µs, m=8 이 C=1: + +| 형상 | stock | v1 | v1 x2 | v2(규칙) | v2 x2 | 규칙 ksr/유닛 | 스윕 최선 (ksr: 단독/x2) | +|---|---|---|---|---|---|---|---| +| [6416×4096] in_proj | 133 | 84.0 | 80.9 | 93.2 | 86.0 | 4 / 204 | 8: 86.0/79.4 | +| [4096×4096] MLA o | 100 | 61.4 | 58.4 | 60.4 | 54.8 | 6 / 192 | 2: 55.3/55.3 | +| [4096×2048] KDA o | 49 | 34.8 | 31.7 | 41.0 | 34.8 | 4 / 128 | 2: 34.8/32.8 | +| [2048×4096] MLA qkv | 49 | 36.9 | 32.8 | 39.9 | 34.8 | 8 / 128 | 6: 35.8/29.7 | +| [6144×4096] dense gu | 142 | 78.7 | 74.8 | 80.9 | 73.6 | 4 / 192 | 2: 74.8/70.7 | +| [4096×3072] dense dn | 73 | 50.2 | 46.6 | 50.2 | 44.0 | 6 / 192 | 2: 45.1/44.0 | +| [1024×4096] 공유 gu | 34 | 26.6 | 22.5 | **24.6** | **21.5** | 8 / 64 | 8 | +| [4096×512] 공유 dn | 20.5 | 18.4 | 15.4 | **16.4** | **14.8** | 1 / 32 | 1: 16.4/13.9 | +| m=16 [4096×4096] | 97 | 63.5 | 60.4 | 82.9 | 67.9 | 6 / 192 | 2: 63.5/62.5 | +| m=32 [2048×4096] | 54 | 41.9 | 37.9 | 53.2 | 49.6 | 8 / 128 | 6: 45.0/40.9 | +| m=32 [1024×4096] | 35 | 30.7 | 28.6 | 33.8 | 31.2 | 8 / 64 | 8: 33.0/30.7 | + +읽기: (1) 공유 전문가 두 형상은 단독으로도 v1 보다 빠르다(−8~11%) — 서빙에서 보는 것은 +이 값이 아니라 MoE 옆의 노출(§2)이고 그건 노출 프로브가 잰다. (2) 큰 형상은 **ksr 규칙이 +틀렸다**: "2 웨이브 + 슬라이스 ≥ 4" 가 고른 값(4·6·8)보다 스윕의 최선은 유닛 수가 상주 슬롯 +96 에 딱 맞거나(6144 × 2 = 96, 2048 × 6 = 96) 그 아래로 긴 슬라이스(4096: ksr 2 = 64 유닛) +인 쪽이고, 51 타일(in_proj)만 잘게(ksr 8 = 408) 썰어야 꼬리가 줄었다(ksr 2 = 102 유닛은 +6개짜리 두 번째 웨이브가 +30 µs). 최선을 쓰면 v2 는 m=8 전 형상에서 v1 과 같거나 빠르다 +(6416 86/79 vs 84/81, 4096² 55 vs 61, 6144 75/71 vs 79/75, 3072 45/44 vs 50/47). +(3) **m ≥ 16 은 v2 가 느리다**(m=16 4096² 83 vs 63.5, m=32 2048 53 vs 42): 블록마다 k-블록당 +x 를 다시 읽고 양자화하는 비용이 m 에 비례하고(warp 당 4행), 그 L2 적재 지연이 mma 뒤에 +드러난다 — v1 은 발사당 한 번이었다. C=1(m=8)에는 없는 비용이지만 C≥2 브래킷 전에 고쳐야 +한다(후보: m ≤ 16 특화 인스턴스 + x 2단 선적재, 또는 k-블록별 first-arriver 양자화 공유). + +### 5. 측정 2차 — ksr 3 스윕과 규칙 확정 (2026-09-05 01:11, srv2, chain5 뒤 창, `g2.ksr3.out`) + +| 형상 (m=8) | ksr 2 | ksr 3 | ksr 4 | ksr 8 | v1 x2 | +|---|---|---|---|---|---| +| [6416×4096] 51타일 | 112.6/98.8 | 87.0/82.4 | 93.2/86.6 | **86.0/79.4** | 80.4 | +| [4096×4096] 32 | 56.3/55.3 | **55.3/49.7** | 66.6/61.9 | 61.4/56.3 | 56.3 | +| [4096×2048] 32 | 34.8/32.8 | **35.8/29.7** | 40.9/34.8 | 37.9/34.8 | 32.8 | +| [2048×4096] 16 | 43.0/41.0 | 34.8/33.8 | **34.8/32.8** | 41.0/34.8 | 33.2 | +| [6144×4096] 48 | **76.0/70.7** | 85.0/82.6 | 79.0/73.2 | 81.9/75.8 | 72.8 | +| [4096×3072] 32 | 45.1/44.5 | **44.1/39.9** | 55.3/48.1 | 51.2/45.6 | 46.1 | +| m=16 [4096×4096] | 62.5/62.0 | **56.3/51.7** | 70.7/65.1 | 68.6/62.0 | 59.4 | +| m=32 [2048×4096] | 61.4/58.4 | **47.1/45.1** | 47.1/45.6 | 53.3/48.6 | 37.9 | + +(단독/x2 µs.) **유닛 수 = 상주 슬롯 96 딱 한 웨이브**(48×2, 32×3, 16×6)가 k 와 무관하게 +최선이고, 51 타일은 짧은 두 번째 웨이브(102)가 최악이라 잘게(8) 썰어야 하며, 8 타일은 +슬라이스 ≥ 4 k-블록 안에서 최대(8 = 64 유닛). 이 규칙을 `mk_choose_ksr2` 에 넣었다(SM 수는 +장치에서). 규칙 적용 시 v2 x2 vs v1 x2: in_proj 79.4/80.4, 4096² 49.7/56.3, k2048 29.7/32.8, +n2048 32.8/33.2, 6144 70.7/72.8, k3072 39.9/46.1, 1024 21.5/22.5, k512 14.8/15.4, m=16 4096² +51.7/59.4 — **m ≤ 16 전 형상에서 v2 ≥ v1**. m=32 는 45.1 vs 37.9 로 여전히 뒤진다(§4 (3)). + +노출 프로브 1차는 스톡 flashinfer 래퍼의 `run()` 에 `out=` 이 없어 실패(벤치 러너는 +메가커널 파일만 마운트하고, `out=` 은 배포 b12x 오버레이의 확장) — 반환값을 쓰도록 고쳐 +재대기. 유닛 스탬프도 다음 창. + +**리뷰(PR #305 머지 뒤, 13건)에서 나온 커널 쪽 사실**(SASS 판독): k-블록당 warp 명령 ~670 +중 LUT 쌍 재구성(프래그먼트마다 16쌍) 216, 스테이지 복사 루프 미언롤 80, 행당 양자화 70 +(`1.0f/sc` 가 IEEE 나눗셈 슬로패스 호출), 에필로그의 런타임 바운드 때문에 acc 가 로컬 +메모리(64 B 스택). 후속 PR 에서 복사 언롤·`__frcp_rn`(2의 거듭제곱이라 정확)·정적 인덱스를 +넣었고, LUT 쌍 공유(k 순열로 프래그먼트 두 개를 한 그룹에)는 v1 과의 비트 동일을 깨므로 +스탬프로 발행 한계가 확인되면 다음 라운드. ## ★★★28차 — 드래프터 W4 는 서빙된 적이 없었다(컴파일 캐시), MK-MLA 서빙 사망의 원인은 스크래치 재할당, 그리고 드래프터 W4 판정 (2026-09-04) diff --git a/RUNBOOK_KERNEL_CAMPAIGN2.md b/RUNBOOK_KERNEL_CAMPAIGN2.md index 29a4b52f..8f2777b0 100644 --- a/RUNBOOK_KERNEL_CAMPAIGN2.md +++ b/RUNBOOK_KERNEL_CAMPAIGN2.md @@ -595,9 +595,11 @@ deep_gemm 은 독립 블록이라 같은 GEMM 을 MoE 꼬리 안에서 끝냈다 두 레인 × ksr 5종). **오프라인 게이트(부팅 전, srv2 GPU 창)**: `run_megakernel_bench.sh --segments exact,gemm ---gemm2 both --ksr2-sweep 1,2,4,6,8` 전 행 PASS + v2 열이 v1 열보다 형상마다 빠름(특히 -[6144×4096]·[4096×512]·[1024×4096]) + `mk_gemm_moe_overlap_probe.py` 의 v2 노출이 v1 의 -28 µs/층보다 작음. 상한: 스텝 −2.5 ms(메인 −1.5, aux 노출 −1.0), 트레이스 mk 합 14.1 → ~7. +--gemm2 both`(기본 `--gemm2 env` 는 서빙 레인만 판정한다) 전 행 PASS + v2 열이 v1 열보다 +형상마다 같거나 빠름 + `mk_gemm_moe_overlap_probe.py` 의 v2 노출이 v1 의 28 µs/층보다 작음. +상태(09-05 01:11, 원장 30차 §4~5): exact PASS(두 레인 × ksr 5종), 규칙 확정 뒤 m ≤ 16 전 +형상 v2 ≥ v1(x2), m = 32 는 뒤짐(C=1 전용 팔) — 노출 프로브·스탬프 대기. 상한: 스텝 −2.5 ms +(메인 −1.5, aux 노출 −1.0), 트레이스 mk 합 14.1 → ~7. **부팅 게이트**: base(기본값) → cand(`VLLM_GLM53_MK_GEMM2=1`) 브래킷, step/s(acc 정규화) + 프리필 동반 + 품질 9/9 + 한국어 0/16 + pos-1 ±2 pct. 통과하면 프로필 기본 1 로 올리고 diff --git a/build/dsv4/glm53_megakernel.cu b/build/dsv4/glm53_megakernel.cu index d8aea064..6cb67e27 100644 --- a/build/dsv4/glm53_megakernel.cu +++ b/build/dsv4/glm53_megakernel.cu @@ -1017,6 +1017,14 @@ __global__ void mk_gemm_kernel(const MKGemmCtx c) { // expanded tiles. ksr > 1 slices go to an fp32 partial and the last slice // to arrive folds the tile in fixed slice order (the persistent kernel's // leftover fold, applied to every tile): deterministic, no zero pass. +// +// Shared state: the partials and the per-tile arrival counters below are +// one set for the whole device, like the persistent lane's A tiles and +// unit counter -- so, like every MK GEMM launch, a v2 launch must never +// overlap another MK GEMM launch on a different stream (two launches +// folding the same tile index would sum each other's slices). Serving +// keeps that by stream order: the side-stream pair is joined before the +// next main-stream GEMM. // =========================================================================== // v2 raw-record ring depth (VLLM_GLM53_MK_NBUF2; 2..4 keep two blocks per SM) #ifndef MK_NBUF2_DEF @@ -1112,7 +1120,12 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { const uint8_t* ssrc = (const uint8_t*)c.ws4 + ((size_t)nt * kblk + kb) * (SMEM_W_ROWS * 8); uint8_t* d = sraw + buf * W4_RAW_BYTES; - for (int t = (int)threadIdx.x; t < SMEM_W_ROWS * 4; t += MK_THREADS) { + // two chunks per thread, unrolled: the swizzled destinations are + // per-thread constants (a runtime loop recomputed them every k-block) + static_assert((SMEM_W_ROWS * 4) % MK_THREADS == 0, "chunks per thread"); +#pragma unroll + for (int u = 0; u < (SMEM_W_ROWS * 4) / MK_THREADS; ++u) { + const int t = (int)threadIdx.x + u * MK_THREADS; const int r = t >> 2, ch = t & 3; mk_cp_async16(d + r * W4_RAW_PITCH + ((ch ^ ((r >> 1) & 3)) << 4), nsrc + (size_t)t * 16); @@ -1153,7 +1166,10 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { for (int off = 16; off; off >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, off)); const float sc = mk_pow2_scale(mx); - const float rsc = 1.0f / sc; // exact: sc is a power of two + // sc is a power of two, so rcp.rn is exact -- the same rsc as the + // shared prologue's 1.0f / sc, without the IEEE divide's slow-path + // call that the SASS showed on every row of every k-block + const float rsc = __frcp_rn(sc); uint32_t pack = 0; #pragma unroll for (int q = 0; q < 4; ++q) @@ -1289,28 +1305,24 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { } MK2_TS(2); - // ---- epilogue - if (ksr == 1) { // whole tile: bf16 out, real rows / cols only + // ---- epilogue: one walk over the fragment's real rows / cols, two stores + auto store_tile = [&](auto&& put) { // put(row, col, value) #pragma unroll - for (int i = 0; i < mtiles; ++i) { + for (int i = 0; i < 2; ++i) { // static index: a runtime bound put acc in local memory + if (i >= mtiles) break; #pragma unroll for (int j = 0; j < 2; ++j) { const int r0 = i * 16 + g, r1 = r0 + 8; const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb] = __float2bfloat16(acc[i][j][0]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][1]); - } - if (r1 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb] = __float2bfloat16(acc[i][j][2]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][3]); - } + if (r0 < c.m) { put(r0, cb, acc[i][j][0]); put(r0, cb + 1, acc[i][j][1]); } + if (r1 < c.m) { put(r1, cb, acc[i][j][2]); put(r1, cb + 1, acc[i][j][3]); } } } + }; + if (ksr == 1) { // whole tile: bf16 out + store_tile([&](int r, int col, float v) { + if (col < c.n_orig) c.out[(size_t)r * c.n_orig + col] = __float2bfloat16(v); + }); MK2_TS(3); return; } @@ -1318,22 +1330,7 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { // arrival, and let the last slice fold the tile in slice order. { float* pb = g_mk2_partial + (size_t)sp * c.m * c.n; -#pragma unroll - for (int i = 0; i < mtiles; ++i) { -#pragma unroll - for (int j = 0; j < 2; ++j) { - const int r0 = i * 16 + g, r1 = r0 + 8; - const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - pb[(size_t)r0 * c.n + cb] = acc[i][j][0]; - pb[(size_t)r0 * c.n + cb + 1] = acc[i][j][1]; - } - if (r1 < c.m) { - pb[(size_t)r1 * c.n + cb] = acc[i][j][2]; - pb[(size_t)r1 * c.n + cb + 1] = acc[i][j][3]; - } - } - } + store_tile([&](int r, int col, float v) { pb[(size_t)r * c.n + col] = v; }); } __syncthreads(); __threadfence(); // release: the slice is visible device-wide first @@ -2664,6 +2661,7 @@ bool g_attrs_set = false; // resident blocks per SM the device reports for mk_gemm2_kernel (2 by // construction of GEMM2_SMEM; the v2 unit rule sizes its grid from it) int g_gemm2_bps = 0; +int g_mk_sms = 0; // multiprocessors, from the device (48 on GB10) void set_kernel_attrs() { if (g_attrs_set) return; @@ -2680,6 +2678,8 @@ void set_kernel_attrs() { MK_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &g_gemm2_bps, mk_gemm2_kernel, MK_THREADS, GEMM2_SMEM)); if (g_gemm2_bps < 1) g_gemm2_bps = 1; + MK_CHECK_CUDA(cudaDeviceGetAttribute( + &g_mk_sms, cudaDevAttrMultiProcessorCount, 0)); MK_CHECK_CUDA(cudaFuncSetAttribute( mk_mla_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, MLA_SMEM)); g_attrs_set = true; @@ -2797,12 +2797,18 @@ bool mk_gemm2_on() { } return g_probe_gemm2 == 1; } -// k-slices per tile for one v2 launch. Enough units for two waves of the -// resident slots (blocks/SM x 48 SMs), so the scheduler has something to -// balance the tail with, but no slice shorter than 4 k-blocks (a unit's -// ring fill and quant are paid once per unit). Then the contract clamps: -// ksr <= kblk (every slice non-empty), ksr <= MK2_KSR_MAX, and the fp32 -// partial must fit (m x n x ksr floats). +// k-slices per tile for one v2 launch, from the 30차 sweeps (srv2, ksr 1/2/3/ +// 4/6/8 on every production shape, single and back-to-back): what wins is +// ONE exact wave of the resident slots (blocks/SM x SMs = 96 on GB10) -- +// 48 tiles x 2, 32 x 3, 16 x 6 measured best on every k -- or, when the +// tile count does not divide the slots, the longest slices that still fit +// one wave (8 tiles x 8 = 64 units); a small SECOND wave is the worst case +// (51 tiles x 2 = 102 units: +30 us for the six stragglers), so a tile +// count above half the slots takes the finest slices instead (51 x 8 = 408, +// 86 us; x 3 = 153, 87; x 1 = 51, 97). No slice shorter than 4 k-blocks (a +// unit's ring fill and first quant are paid per unit; k = 512 stays whole). +// Then the contract clamps: ksr <= kblk (every slice non-empty), ksr <= +// MK2_KSR_MAX, and the fp32 partial must fit (m x n x ksr floats). int mk_choose_ksr2(int m, int n, int k) { const int nblk = n / SMEM_W_ROWS, kblk = k / KSTEP; if (g_probe_ksr2 < 0) { @@ -2813,10 +2819,17 @@ int mk_choose_ksr2(int m, int n, int k) { if (g_probe_ksr2 > 0) { ksr = g_probe_ksr2; } else { - const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * 48; - ksr = (2 * slots + nblk - 1) / nblk; - const int kmax = kblk / 4; - if (ksr > kmax) ksr = kmax; + const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * + (g_mk_sms > 0 ? g_mk_sms : 48); + const int kmax = kblk / 4 > 1 ? kblk / 4 : 1; + if (slots % nblk == 0 && slots / nblk <= kmax) { + ksr = slots / nblk; // one exact wave + } else if (nblk * 2 > slots) { + ksr = kmax; // would leave a short second wave: slice fine + } else { + ksr = slots / nblk; // under one wave, the longest slices + if (ksr > kmax) ksr = kmax; + } } if (ksr < 1) ksr = 1; if (ksr > kblk) ksr = kblk; @@ -2840,18 +2853,27 @@ std::vector mk_probe_device() { (int64_t)prop.sharedMemPerBlockOptin}; } -// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then -// cleared. Empty unless built with -DMK_PHASE_TS=1. -std::vector mk_read_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 8); +// Read a device stamp array into a host vector and clear it (every probe +// reader below): synchronize first -- the stamps are written by the launch +// the caller just issued. +template +std::vector mk_read_and_clear(unsigned long long (&sym)[N]) { + std::vector h(N); MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_ts, - sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), sym, sizeof(unsigned long long) * N)); void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaGetSymbolAddress(&p, sym)); + MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * N)); return std::vector(h.begin(), h.end()); +} +#endif + +// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then +// cleared. Empty unless built with -DMK_PHASE_TS=1. +std::vector mk_read_ts() { +#ifdef MK_PHASE_TS + return mk_read_and_clear(g_mk_ts); #else return {}; #endif @@ -2861,13 +2883,7 @@ std::vector mk_read_ts() { // p2 done, barrier, p3 done, barrier, p4 done. std::vector mk_read_kda_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 16); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_kda_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_kda_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_kda_ts); #else return {}; #endif @@ -2875,14 +2891,7 @@ std::vector mk_read_kda_ts() { std::vector mk_read_mhc_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_MHC_GRID_CAP * 8); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_mhc_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_mhc_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_mhc_ts); #else return {}; #endif @@ -3119,14 +3128,7 @@ void mk_set_gemm2(int64_t on, int64_t ksr) { // v2 unit timestamps of the last launch, [MK2_UNITS_MAX][4] ns, then cleared. std::vector mk_read_ts2() { #ifdef MK_PHASE_TS - std::vector h(MK2_UNITS_MAX * 4); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk2_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk2_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk2_ts); #else return {}; #endif diff --git a/build/glm53/glm53_megakernel.cu b/build/glm53/glm53_megakernel.cu index d8aea064..6cb67e27 100644 --- a/build/glm53/glm53_megakernel.cu +++ b/build/glm53/glm53_megakernel.cu @@ -1017,6 +1017,14 @@ __global__ void mk_gemm_kernel(const MKGemmCtx c) { // expanded tiles. ksr > 1 slices go to an fp32 partial and the last slice // to arrive folds the tile in fixed slice order (the persistent kernel's // leftover fold, applied to every tile): deterministic, no zero pass. +// +// Shared state: the partials and the per-tile arrival counters below are +// one set for the whole device, like the persistent lane's A tiles and +// unit counter -- so, like every MK GEMM launch, a v2 launch must never +// overlap another MK GEMM launch on a different stream (two launches +// folding the same tile index would sum each other's slices). Serving +// keeps that by stream order: the side-stream pair is joined before the +// next main-stream GEMM. // =========================================================================== // v2 raw-record ring depth (VLLM_GLM53_MK_NBUF2; 2..4 keep two blocks per SM) #ifndef MK_NBUF2_DEF @@ -1112,7 +1120,12 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { const uint8_t* ssrc = (const uint8_t*)c.ws4 + ((size_t)nt * kblk + kb) * (SMEM_W_ROWS * 8); uint8_t* d = sraw + buf * W4_RAW_BYTES; - for (int t = (int)threadIdx.x; t < SMEM_W_ROWS * 4; t += MK_THREADS) { + // two chunks per thread, unrolled: the swizzled destinations are + // per-thread constants (a runtime loop recomputed them every k-block) + static_assert((SMEM_W_ROWS * 4) % MK_THREADS == 0, "chunks per thread"); +#pragma unroll + for (int u = 0; u < (SMEM_W_ROWS * 4) / MK_THREADS; ++u) { + const int t = (int)threadIdx.x + u * MK_THREADS; const int r = t >> 2, ch = t & 3; mk_cp_async16(d + r * W4_RAW_PITCH + ((ch ^ ((r >> 1) & 3)) << 4), nsrc + (size_t)t * 16); @@ -1153,7 +1166,10 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { for (int off = 16; off; off >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, off)); const float sc = mk_pow2_scale(mx); - const float rsc = 1.0f / sc; // exact: sc is a power of two + // sc is a power of two, so rcp.rn is exact -- the same rsc as the + // shared prologue's 1.0f / sc, without the IEEE divide's slow-path + // call that the SASS showed on every row of every k-block + const float rsc = __frcp_rn(sc); uint32_t pack = 0; #pragma unroll for (int q = 0; q < 4; ++q) @@ -1289,28 +1305,24 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { } MK2_TS(2); - // ---- epilogue - if (ksr == 1) { // whole tile: bf16 out, real rows / cols only + // ---- epilogue: one walk over the fragment's real rows / cols, two stores + auto store_tile = [&](auto&& put) { // put(row, col, value) #pragma unroll - for (int i = 0; i < mtiles; ++i) { + for (int i = 0; i < 2; ++i) { // static index: a runtime bound put acc in local memory + if (i >= mtiles) break; #pragma unroll for (int j = 0; j < 2; ++j) { const int r0 = i * 16 + g, r1 = r0 + 8; const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb] = __float2bfloat16(acc[i][j][0]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][1]); - } - if (r1 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb] = __float2bfloat16(acc[i][j][2]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][3]); - } + if (r0 < c.m) { put(r0, cb, acc[i][j][0]); put(r0, cb + 1, acc[i][j][1]); } + if (r1 < c.m) { put(r1, cb, acc[i][j][2]); put(r1, cb + 1, acc[i][j][3]); } } } + }; + if (ksr == 1) { // whole tile: bf16 out + store_tile([&](int r, int col, float v) { + if (col < c.n_orig) c.out[(size_t)r * c.n_orig + col] = __float2bfloat16(v); + }); MK2_TS(3); return; } @@ -1318,22 +1330,7 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { // arrival, and let the last slice fold the tile in slice order. { float* pb = g_mk2_partial + (size_t)sp * c.m * c.n; -#pragma unroll - for (int i = 0; i < mtiles; ++i) { -#pragma unroll - for (int j = 0; j < 2; ++j) { - const int r0 = i * 16 + g, r1 = r0 + 8; - const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - pb[(size_t)r0 * c.n + cb] = acc[i][j][0]; - pb[(size_t)r0 * c.n + cb + 1] = acc[i][j][1]; - } - if (r1 < c.m) { - pb[(size_t)r1 * c.n + cb] = acc[i][j][2]; - pb[(size_t)r1 * c.n + cb + 1] = acc[i][j][3]; - } - } - } + store_tile([&](int r, int col, float v) { pb[(size_t)r * c.n + col] = v; }); } __syncthreads(); __threadfence(); // release: the slice is visible device-wide first @@ -2664,6 +2661,7 @@ bool g_attrs_set = false; // resident blocks per SM the device reports for mk_gemm2_kernel (2 by // construction of GEMM2_SMEM; the v2 unit rule sizes its grid from it) int g_gemm2_bps = 0; +int g_mk_sms = 0; // multiprocessors, from the device (48 on GB10) void set_kernel_attrs() { if (g_attrs_set) return; @@ -2680,6 +2678,8 @@ void set_kernel_attrs() { MK_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &g_gemm2_bps, mk_gemm2_kernel, MK_THREADS, GEMM2_SMEM)); if (g_gemm2_bps < 1) g_gemm2_bps = 1; + MK_CHECK_CUDA(cudaDeviceGetAttribute( + &g_mk_sms, cudaDevAttrMultiProcessorCount, 0)); MK_CHECK_CUDA(cudaFuncSetAttribute( mk_mla_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, MLA_SMEM)); g_attrs_set = true; @@ -2797,12 +2797,18 @@ bool mk_gemm2_on() { } return g_probe_gemm2 == 1; } -// k-slices per tile for one v2 launch. Enough units for two waves of the -// resident slots (blocks/SM x 48 SMs), so the scheduler has something to -// balance the tail with, but no slice shorter than 4 k-blocks (a unit's -// ring fill and quant are paid once per unit). Then the contract clamps: -// ksr <= kblk (every slice non-empty), ksr <= MK2_KSR_MAX, and the fp32 -// partial must fit (m x n x ksr floats). +// k-slices per tile for one v2 launch, from the 30차 sweeps (srv2, ksr 1/2/3/ +// 4/6/8 on every production shape, single and back-to-back): what wins is +// ONE exact wave of the resident slots (blocks/SM x SMs = 96 on GB10) -- +// 48 tiles x 2, 32 x 3, 16 x 6 measured best on every k -- or, when the +// tile count does not divide the slots, the longest slices that still fit +// one wave (8 tiles x 8 = 64 units); a small SECOND wave is the worst case +// (51 tiles x 2 = 102 units: +30 us for the six stragglers), so a tile +// count above half the slots takes the finest slices instead (51 x 8 = 408, +// 86 us; x 3 = 153, 87; x 1 = 51, 97). No slice shorter than 4 k-blocks (a +// unit's ring fill and first quant are paid per unit; k = 512 stays whole). +// Then the contract clamps: ksr <= kblk (every slice non-empty), ksr <= +// MK2_KSR_MAX, and the fp32 partial must fit (m x n x ksr floats). int mk_choose_ksr2(int m, int n, int k) { const int nblk = n / SMEM_W_ROWS, kblk = k / KSTEP; if (g_probe_ksr2 < 0) { @@ -2813,10 +2819,17 @@ int mk_choose_ksr2(int m, int n, int k) { if (g_probe_ksr2 > 0) { ksr = g_probe_ksr2; } else { - const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * 48; - ksr = (2 * slots + nblk - 1) / nblk; - const int kmax = kblk / 4; - if (ksr > kmax) ksr = kmax; + const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * + (g_mk_sms > 0 ? g_mk_sms : 48); + const int kmax = kblk / 4 > 1 ? kblk / 4 : 1; + if (slots % nblk == 0 && slots / nblk <= kmax) { + ksr = slots / nblk; // one exact wave + } else if (nblk * 2 > slots) { + ksr = kmax; // would leave a short second wave: slice fine + } else { + ksr = slots / nblk; // under one wave, the longest slices + if (ksr > kmax) ksr = kmax; + } } if (ksr < 1) ksr = 1; if (ksr > kblk) ksr = kblk; @@ -2840,18 +2853,27 @@ std::vector mk_probe_device() { (int64_t)prop.sharedMemPerBlockOptin}; } -// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then -// cleared. Empty unless built with -DMK_PHASE_TS=1. -std::vector mk_read_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 8); +// Read a device stamp array into a host vector and clear it (every probe +// reader below): synchronize first -- the stamps are written by the launch +// the caller just issued. +template +std::vector mk_read_and_clear(unsigned long long (&sym)[N]) { + std::vector h(N); MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_ts, - sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), sym, sizeof(unsigned long long) * N)); void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaGetSymbolAddress(&p, sym)); + MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * N)); return std::vector(h.begin(), h.end()); +} +#endif + +// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then +// cleared. Empty unless built with -DMK_PHASE_TS=1. +std::vector mk_read_ts() { +#ifdef MK_PHASE_TS + return mk_read_and_clear(g_mk_ts); #else return {}; #endif @@ -2861,13 +2883,7 @@ std::vector mk_read_ts() { // p2 done, barrier, p3 done, barrier, p4 done. std::vector mk_read_kda_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 16); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_kda_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_kda_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_kda_ts); #else return {}; #endif @@ -2875,14 +2891,7 @@ std::vector mk_read_kda_ts() { std::vector mk_read_mhc_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_MHC_GRID_CAP * 8); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_mhc_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_mhc_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_mhc_ts); #else return {}; #endif @@ -3119,14 +3128,7 @@ void mk_set_gemm2(int64_t on, int64_t ksr) { // v2 unit timestamps of the last launch, [MK2_UNITS_MAX][4] ns, then cleared. std::vector mk_read_ts2() { #ifdef MK_PHASE_TS - std::vector h(MK2_UNITS_MAX * 4); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk2_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk2_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk2_ts); #else return {}; #endif diff --git a/overlay/modules/glm53_megakernel/README.md b/overlay/modules/glm53_megakernel/README.md index cf05c7d9..d5e0acbe 100644 --- a/overlay/modules/glm53_megakernel/README.md +++ b/overlay/modules/glm53_megakernel/README.md @@ -82,8 +82,24 @@ which frees the 32 KB of expanded tiles: **37 KB of smem, two blocks per SM** (`__launch_bounds__(256, 2)`). One `__syncthreads` per k-block. ksr > 1 slices assign an fp32 partial and the last slice to arrive folds the tile in fixed order (deterministic, no zero pass). The slice rule -(`mk_choose_ksr2`) wants two waves of the resident slots and no slice -shorter than 4 k-blocks; `VLLM_GLM53_MK_KSR2` forces it for sweeps. +(`mk_choose_ksr2`, from the 30차 sweeps) takes ONE exact wave of the +device's resident slots (blocks/SM x SMs = 96: 48 tiles x 2, 32 x 3, +16 x 6), the longest slices that fit one wave when the tile count does +not divide it (8 x 8), and the finest slices (>= 4 k-blocks) above half +the slots, where a short second wave is the worst case (51 tiles x 2 = +102 units measured +30 us); `VLLM_GLM53_MK_KSR2` forces it for sweeps. +With that rule v2 matches or beats the persistent lane on every m <= 16 +shape back to back (in_proj 79.4 vs 80.4 us, [4096x4096] 49.7 vs 56.3, +[4096x3072] 39.9 vs 46.1, shared expert 21.5/14.8 vs 22.5/15.4); m = 32 +still trails (45.1 vs 37.9 on [2048x4096]: the per-block x quant scales +with m), so the lane is for C=1 until that path is fixed. + +Contract, inherited from the persistent lane: v2's partials and per-tile +arrival counters are one device-wide set, so a v2 launch must never +overlap another MK GEMM launch on a different stream (two launches +folding the same tile index would sum each other's slices). Serving +keeps it by stream order -- the side-stream pair is joined before the +next main-stream GEMM. Gates and numbers: the boot self-test's exact-grid gate runs on whichever lane the boot serves and the fingerprint names it (`lane v2, in_proj plan diff --git a/overlay/modules/glm53_megakernel/glm53_megakernel.cu b/overlay/modules/glm53_megakernel/glm53_megakernel.cu index d8aea064..6cb67e27 100644 --- a/overlay/modules/glm53_megakernel/glm53_megakernel.cu +++ b/overlay/modules/glm53_megakernel/glm53_megakernel.cu @@ -1017,6 +1017,14 @@ __global__ void mk_gemm_kernel(const MKGemmCtx c) { // expanded tiles. ksr > 1 slices go to an fp32 partial and the last slice // to arrive folds the tile in fixed slice order (the persistent kernel's // leftover fold, applied to every tile): deterministic, no zero pass. +// +// Shared state: the partials and the per-tile arrival counters below are +// one set for the whole device, like the persistent lane's A tiles and +// unit counter -- so, like every MK GEMM launch, a v2 launch must never +// overlap another MK GEMM launch on a different stream (two launches +// folding the same tile index would sum each other's slices). Serving +// keeps that by stream order: the side-stream pair is joined before the +// next main-stream GEMM. // =========================================================================== // v2 raw-record ring depth (VLLM_GLM53_MK_NBUF2; 2..4 keep two blocks per SM) #ifndef MK_NBUF2_DEF @@ -1112,7 +1120,12 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { const uint8_t* ssrc = (const uint8_t*)c.ws4 + ((size_t)nt * kblk + kb) * (SMEM_W_ROWS * 8); uint8_t* d = sraw + buf * W4_RAW_BYTES; - for (int t = (int)threadIdx.x; t < SMEM_W_ROWS * 4; t += MK_THREADS) { + // two chunks per thread, unrolled: the swizzled destinations are + // per-thread constants (a runtime loop recomputed them every k-block) + static_assert((SMEM_W_ROWS * 4) % MK_THREADS == 0, "chunks per thread"); +#pragma unroll + for (int u = 0; u < (SMEM_W_ROWS * 4) / MK_THREADS; ++u) { + const int t = (int)threadIdx.x + u * MK_THREADS; const int r = t >> 2, ch = t & 3; mk_cp_async16(d + r * W4_RAW_PITCH + ((ch ^ ((r >> 1) & 3)) << 4), nsrc + (size_t)t * 16); @@ -1153,7 +1166,10 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { for (int off = 16; off; off >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, off)); const float sc = mk_pow2_scale(mx); - const float rsc = 1.0f / sc; // exact: sc is a power of two + // sc is a power of two, so rcp.rn is exact -- the same rsc as the + // shared prologue's 1.0f / sc, without the IEEE divide's slow-path + // call that the SASS showed on every row of every k-block + const float rsc = __frcp_rn(sc); uint32_t pack = 0; #pragma unroll for (int q = 0; q < 4; ++q) @@ -1289,28 +1305,24 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { } MK2_TS(2); - // ---- epilogue - if (ksr == 1) { // whole tile: bf16 out, real rows / cols only + // ---- epilogue: one walk over the fragment's real rows / cols, two stores + auto store_tile = [&](auto&& put) { // put(row, col, value) #pragma unroll - for (int i = 0; i < mtiles; ++i) { + for (int i = 0; i < 2; ++i) { // static index: a runtime bound put acc in local memory + if (i >= mtiles) break; #pragma unroll for (int j = 0; j < 2; ++j) { const int r0 = i * 16 + g, r1 = r0 + 8; const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb] = __float2bfloat16(acc[i][j][0]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r0 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][1]); - } - if (r1 < c.m) { - if (cb < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb] = __float2bfloat16(acc[i][j][2]); - if (cb + 1 < c.n_orig) - c.out[(size_t)r1 * c.n_orig + cb + 1] = __float2bfloat16(acc[i][j][3]); - } + if (r0 < c.m) { put(r0, cb, acc[i][j][0]); put(r0, cb + 1, acc[i][j][1]); } + if (r1 < c.m) { put(r1, cb, acc[i][j][2]); put(r1, cb + 1, acc[i][j][3]); } } } + }; + if (ksr == 1) { // whole tile: bf16 out + store_tile([&](int r, int col, float v) { + if (col < c.n_orig) c.out[(size_t)r * c.n_orig + col] = __float2bfloat16(v); + }); MK2_TS(3); return; } @@ -1318,22 +1330,7 @@ mk_gemm2_kernel(const MKGemm2Ctx c) { // arrival, and let the last slice fold the tile in slice order. { float* pb = g_mk2_partial + (size_t)sp * c.m * c.n; -#pragma unroll - for (int i = 0; i < mtiles; ++i) { -#pragma unroll - for (int j = 0; j < 2; ++j) { - const int r0 = i * 16 + g, r1 = r0 + 8; - const int cb = nt * 128 + warp * 16 + j * 8 + (lane & 3) * 2; - if (r0 < c.m) { - pb[(size_t)r0 * c.n + cb] = acc[i][j][0]; - pb[(size_t)r0 * c.n + cb + 1] = acc[i][j][1]; - } - if (r1 < c.m) { - pb[(size_t)r1 * c.n + cb] = acc[i][j][2]; - pb[(size_t)r1 * c.n + cb + 1] = acc[i][j][3]; - } - } - } + store_tile([&](int r, int col, float v) { pb[(size_t)r * c.n + col] = v; }); } __syncthreads(); __threadfence(); // release: the slice is visible device-wide first @@ -2664,6 +2661,7 @@ bool g_attrs_set = false; // resident blocks per SM the device reports for mk_gemm2_kernel (2 by // construction of GEMM2_SMEM; the v2 unit rule sizes its grid from it) int g_gemm2_bps = 0; +int g_mk_sms = 0; // multiprocessors, from the device (48 on GB10) void set_kernel_attrs() { if (g_attrs_set) return; @@ -2680,6 +2678,8 @@ void set_kernel_attrs() { MK_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &g_gemm2_bps, mk_gemm2_kernel, MK_THREADS, GEMM2_SMEM)); if (g_gemm2_bps < 1) g_gemm2_bps = 1; + MK_CHECK_CUDA(cudaDeviceGetAttribute( + &g_mk_sms, cudaDevAttrMultiProcessorCount, 0)); MK_CHECK_CUDA(cudaFuncSetAttribute( mk_mla_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, MLA_SMEM)); g_attrs_set = true; @@ -2797,12 +2797,18 @@ bool mk_gemm2_on() { } return g_probe_gemm2 == 1; } -// k-slices per tile for one v2 launch. Enough units for two waves of the -// resident slots (blocks/SM x 48 SMs), so the scheduler has something to -// balance the tail with, but no slice shorter than 4 k-blocks (a unit's -// ring fill and quant are paid once per unit). Then the contract clamps: -// ksr <= kblk (every slice non-empty), ksr <= MK2_KSR_MAX, and the fp32 -// partial must fit (m x n x ksr floats). +// k-slices per tile for one v2 launch, from the 30차 sweeps (srv2, ksr 1/2/3/ +// 4/6/8 on every production shape, single and back-to-back): what wins is +// ONE exact wave of the resident slots (blocks/SM x SMs = 96 on GB10) -- +// 48 tiles x 2, 32 x 3, 16 x 6 measured best on every k -- or, when the +// tile count does not divide the slots, the longest slices that still fit +// one wave (8 tiles x 8 = 64 units); a small SECOND wave is the worst case +// (51 tiles x 2 = 102 units: +30 us for the six stragglers), so a tile +// count above half the slots takes the finest slices instead (51 x 8 = 408, +// 86 us; x 3 = 153, 87; x 1 = 51, 97). No slice shorter than 4 k-blocks (a +// unit's ring fill and first quant are paid per unit; k = 512 stays whole). +// Then the contract clamps: ksr <= kblk (every slice non-empty), ksr <= +// MK2_KSR_MAX, and the fp32 partial must fit (m x n x ksr floats). int mk_choose_ksr2(int m, int n, int k) { const int nblk = n / SMEM_W_ROWS, kblk = k / KSTEP; if (g_probe_ksr2 < 0) { @@ -2813,10 +2819,17 @@ int mk_choose_ksr2(int m, int n, int k) { if (g_probe_ksr2 > 0) { ksr = g_probe_ksr2; } else { - const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * 48; - ksr = (2 * slots + nblk - 1) / nblk; - const int kmax = kblk / 4; - if (ksr > kmax) ksr = kmax; + const int slots = (g_gemm2_bps > 0 ? g_gemm2_bps : 2) * + (g_mk_sms > 0 ? g_mk_sms : 48); + const int kmax = kblk / 4 > 1 ? kblk / 4 : 1; + if (slots % nblk == 0 && slots / nblk <= kmax) { + ksr = slots / nblk; // one exact wave + } else if (nblk * 2 > slots) { + ksr = kmax; // would leave a short second wave: slice fine + } else { + ksr = slots / nblk; // under one wave, the longest slices + if (ksr > kmax) ksr = kmax; + } } if (ksr < 1) ksr = 1; if (ksr > kblk) ksr = kblk; @@ -2840,18 +2853,27 @@ std::vector mk_probe_device() { (int64_t)prop.sharedMemPerBlockOptin}; } -// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then -// cleared. Empty unless built with -DMK_PHASE_TS=1. -std::vector mk_read_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 8); +// Read a device stamp array into a host vector and clear it (every probe +// reader below): synchronize first -- the stamps are written by the launch +// the caller just issued. +template +std::vector mk_read_and_clear(unsigned long long (&sym)[N]) { + std::vector h(N); MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_ts, - sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), sym, sizeof(unsigned long long) * N)); void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); + MK_CHECK_CUDA(cudaGetSymbolAddress(&p, sym)); + MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * N)); return std::vector(h.begin(), h.end()); +} +#endif + +// Phase timestamps of the last gemm launch, [MK_GRID_CAP][8] ns, then +// cleared. Empty unless built with -DMK_PHASE_TS=1. +std::vector mk_read_ts() { +#ifdef MK_PHASE_TS + return mk_read_and_clear(g_mk_ts); #else return {}; #endif @@ -2861,13 +2883,7 @@ std::vector mk_read_ts() { // p2 done, barrier, p3 done, barrier, p4 done. std::vector mk_read_kda_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_GRID_CAP * 16); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_kda_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_kda_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_kda_ts); #else return {}; #endif @@ -2875,14 +2891,7 @@ std::vector mk_read_kda_ts() { std::vector mk_read_mhc_ts() { #ifdef MK_PHASE_TS - std::vector h(MK_MHC_GRID_CAP * 8); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk_mhc_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk_mhc_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk_mhc_ts); #else return {}; #endif @@ -3119,14 +3128,7 @@ void mk_set_gemm2(int64_t on, int64_t ksr) { // v2 unit timestamps of the last launch, [MK2_UNITS_MAX][4] ns, then cleared. std::vector mk_read_ts2() { #ifdef MK_PHASE_TS - std::vector h(MK2_UNITS_MAX * 4); - MK_CHECK_CUDA(cudaDeviceSynchronize()); - MK_CHECK_CUDA(cudaMemcpyFromSymbol(h.data(), g_mk2_ts, - sizeof(unsigned long long) * h.size())); - void* p = nullptr; - MK_CHECK_CUDA(cudaGetSymbolAddress(&p, g_mk2_ts)); - MK_CHECK_CUDA(cudaMemset(p, 0, sizeof(unsigned long long) * h.size())); - return std::vector(h.begin(), h.end()); + return mk_read_and_clear(g_mk2_ts); #else return {}; #endif diff --git a/probes/megakernel_glm53_bench.py b/probes/megakernel_glm53_bench.py index 34b36622..cdd7939c 100644 --- a/probes/megakernel_glm53_bench.py +++ b/probes/megakernel_glm53_bench.py @@ -215,6 +215,7 @@ def probe_gemm(iters: int, gemm2: str = "both", sweep=None) -> bool: # (86/step, 3.5 ms, 2.4 ms of it exposed on the critical path -- 28차): # the shared expert's gate_up [1024 x 4096] and down [4096 x 512], and # o_proj [4096 x 4096]. Override with --gemm-shapes m:n:k,... + kept = {} # (m, n, k) -> (x, p4, p4b): the sweep times the same tensors for m, n, k in GEMM_SHAPES: _lane(ext, 1 if gemm2 == "1" else 0) torch.manual_seed(0) @@ -230,6 +231,8 @@ def probe_gemm(iters: int, gemm2: str = "both", sweep=None) -> bool: sq2, sws2, srows2, scols2 = mk._stock_fp8_pair(w2) p4b = mk.build_mk_weight_w4(w2) del w2 + if sweep: + kept[(m, n, k)] = (x, p4, p4b) ref = _fp8_dense_gemm(x, sq, sws, srows, scols) got = mk._gemm_call(x, p4, n) torch.cuda.synchronize() @@ -298,18 +301,16 @@ def probe_gemm(iters: int, gemm2: str = "both", sweep=None) -> bool: # shape, single launch and back-to-back on two weights print(f"{'v2 ksr sweep':<24}" + "".join(f"{'ksr=' + str(r):>14}" for r in sweep)) for m, n, k in GEMM_SHAPES: - torch.manual_seed(0) - w = torch.randn(n, k, dtype=torch.bfloat16, device=DEV) * 0.05 - w2 = torch.randn(n, k, dtype=torch.bfloat16, device=DEV) * 0.05 - x = torch.randn(m, k, dtype=torch.bfloat16, device=DEV) - p4, p4b = mk.build_mk_weight_w4(w), mk.build_mk_weight_w4(w2) - del w, w2 + x, p4, p4b = kept[(m, n, k)] # built once above cells = [] for r in sweep: - if r > k // 128: + _lane(ext, 1, r) + # the host clamps a forced ksr (ksr <= kblk, <= MK2_KSR_MAX, + # the partial-buffer bound): a cell is only honest at the + # ksr the launch will use + if int(ext.gemm2_plan(m, n, k)[1]) != r: cells.append(f"{'-':>14}") continue - _lane(ext, 1, r) t1 = _time(lambda: mk._gemm_call(x, p4, n), iters, hot=(x,)) t2 = _time(lambda: (mk._gemm_call(x, p4, n), mk._gemm_call(x, p4b, n)), iters, @@ -324,9 +325,9 @@ def probe_gemm(iters: int, gemm2: str = "both", sweep=None) -> bool: def probe_stamps2(iters: int) -> bool: """v2 unit timeline per shape (needs VLLM_GLM53_MK_PHASE_TS=1 so the build stamps g_mk2_ts): entry skew across units, first-record latency, - loop time, tail (last exit - median exit) and the event span, from one - cold-L2 launch after warm-up. Tells whether a shape's ksr leaves a - balance tail or pays too many ring fills.""" + loop time, tail (last exit - median exit) and the event span -- medians + over `iters` cold-L2 launches after warm-up. Tells whether a shape's + ksr leaves a balance tail or pays too many ring fills.""" from vllm.model_executor.layers import glm53_megakernel as mk ext = mk._build() @@ -345,33 +346,38 @@ def probe_stamps2(iters: int) -> bool: mk._gemm_call(x, p4, n) torch.cuda.synchronize() ext.read_ts2() # clear - _l2_flush(hot=(x,)) s = torch.cuda.Event(enable_timing=True) e = torch.cuda.Event(enable_timing=True) - s.record() - mk._gemm_call(x, p4, n) - e.record() - torch.cuda.synchronize() - ev = s.elapsed_time(e) * 1e3 - ts = list(ext.read_ts2()) - if not ts: - print(" (no stamps: build with VLLM_GLM53_MK_PHASE_TS=1)") - return True - rows = [ts[u * 4:u * 4 + 4] for u in range(units)] - rows = [r for r in rows if all(v > 0 for v in r)] - if not rows: + samples = [] # (skew, first, loop, tail, span, ev) per launch + for _ in range(max(1, iters)): + _l2_flush(hot=(x,)) + s.record() + mk._gemm_call(x, p4, n) + e.record() + torch.cuda.synchronize() + ev = s.elapsed_time(e) * 1e3 + ts = list(ext.read_ts2()) + if not ts: + print(" (no stamps: build with VLLM_GLM53_MK_PHASE_TS=1)") + return True + rows = [ts[u * 4:u * 4 + 4] for u in range(units)] + rows = [r for r in rows if all(v > 0 for v in r)] + if not rows: + continue + t0 = min(r[0] for r in rows) + ent = sorted(r[0] - t0 for r in rows) + first = sorted(r[1] - r[0] for r in rows) + loop = sorted(r[2] - r[1] for r in rows) + exits = sorted(r[3] - t0 for r in rows) + samples.append((ent[-1], first[len(first) // 2], loop[len(loop) // 2], + exits[-1] - exits[len(exits) // 2], exits[-1], ev * 1e3)) + if not samples: print(f" gemm m={m:<3}n={n:<5}k={k:<5} no complete unit stamps") continue - t0 = min(r[0] for r in rows) - ent = sorted(r[0] - t0 for r in rows) - first = sorted(r[1] - r[0] for r in rows) - loop = sorted(r[2] - r[1] for r in rows) - exits = sorted(r[3] - t0 for r in rows) - med = lambda a: a[len(a) // 2] / 1e3 + med = [sorted(c)[len(c) // 2] / 1e3 for c in zip(*samples)] print(f" gemm m={m:<3}n={n:<5}k={k:<5}{units:>6}{plan[1]:>4}" - f"{ent[-1] / 1e3:>7.1f}{med(first):>7.1f}{med(loop):>7.1f}" - f"{(exits[-1] - exits[len(exits) // 2]) / 1e3:>7.1f}" - f"{exits[-1] / 1e3:>7.1f}{ev:>7.1f}") + f"{med[0]:>7.1f}{med[1]:>7.1f}{med[2]:>7.1f}{med[3]:>7.1f}" + f"{med[4]:>7.1f}{med[5]:>7.1f}") _lane_restore(ext) return True @@ -393,14 +399,6 @@ def probe_exact(gemm2: str = "both") -> bool: "both": [(0, 0)] + [(1, r) for r in (0, 1, 2, 3, 5, 8)]}[gemm2] ok = True print(f"{'case':<24}{'rel_err':>10}{'gate':>8}") - for on, ksr in lanes: - _lane(ext, on, ksr) - ok &= _probe_exact_lane(mk, f"v{on + 1}" + (f" ksr={ksr}" if ksr else "")) - _lane_restore(ext) - return ok - - -def _probe_exact_lane(mk, tag: str) -> bool: torch.manual_seed(0) n, k, m = 1024, 4096, 8 code = torch.randint(0, 8, (n, k // 16, 16), device=DEV) @@ -419,19 +417,31 @@ def _probe_exact_lane(mk, tag: str) -> bool: torch.randn_like(code.float()) < 0, -1.0, 1.0) w_exact = w_exact.view(n, k).to(torch.bfloat16) x = torch.randn(m, k, dtype=torch.bfloat16, device=DEV) - p4 = mk.build_mk_weight_w4(w_exact) + p4 = mk.build_mk_weight_w4(w_exact) # once: the 16-candidate search w_back = mk.mk_w4_dequant(p4[0], p4[1], n, p4[2]) # p4[2]: 2^-shift e_pack = _rel(w_back, w_exact) # the pack itself must round-trip - got = mk._gemm_call(x, p4, n) + mark = "!" if e_pack != 0.0 else " " + ok &= e_pack == 0.0 + print(f"{mark}w4 pack roundtrip{e_pack:>17.2e}{0:>8.0e}") ref = mk._mk_quant_x_ref(x) @ w_back.float().T + for on, ksr in lanes: + _lane(ext, on, ksr) + ok &= _probe_exact_lane(mk, f"v{on + 1}" + (f" ksr={ksr}" if ksr else ""), + x, p4, ref, n) + _lane_restore(ext) + return ok + + +def _probe_exact_lane(mk, tag: str, x, p4, ref, n: int) -> bool: + """One lane's launch against the shared reference (see probe_exact).""" + got = mk._gemm_call(x, p4, n) torch.cuda.synchronize() # the kernel writes bf16: judge against the bf16-rounded reference, no # element more than one bf16 ulp off (a different fp32 summation order # flips a few by one ulp; a layout bug moves whole rows) e_exact, n_ulp = mk._exact_gate(got, ref) - ok = e_pack == 0.0 and e_exact <= 1e-3 and n_ulp == 0 + ok = e_exact <= 1e-3 and n_ulp == 0 mark = "!" if not ok else " " - print(f"{mark}w4 pack roundtrip{e_pack:>17.2e}{0:>8.0e}") print(f"{mark}{('w4 exact grid ' + tag):<24}{e_exact:>10.2e}{1e-3:>8.0e} over-ulp={n_ulp}") return ok @@ -648,13 +658,18 @@ def main() -> int: ap.add_argument("--sinkhorn", type=int, default=None) # the GEMM lane under test: 0 = the persistent kernel, 1 = the v2 # non-persistent kernel, both = persistent in the mk columns and v2 in - # the mk2 columns (diffed against each other) - ap.add_argument("--gemm2", choices=("0", "1", "both"), default="both") + # the mk2 columns (diffed against each other), env = whichever lane the + # process's VLLM_GLM53_MK_GEMM2 serves. The default is env: this + # probe's VERDICT authorises boot brackets, so it must judge the lane a + # boot will run, not one it will not. + ap.add_argument("--gemm2", choices=("0", "1", "both", "env"), default="env") ap.add_argument("--ksr2-sweep", default=None, help="comma list of v2 slice counts to force per shape, e.g. 1,2,4,6,8") ap.add_argument("--stamps2", action="store_true", help="v2 per-unit timeline (VLLM_GLM53_MK_PHASE_TS=1 build)") args = ap.parse_args() + if args.gemm2 == "env": + args.gemm2 = "1" if os.environ.get("VLLM_GLM53_MK_GEMM2") == "1" else "0" if args.gemm_shapes: GEMM_SHAPES[:] = [tuple(int(v) for v in t.split(":")) for t in args.gemm_shapes.split(",")] if args.sinkhorn is not None and args.sinkhorn < 1: diff --git a/probes/mk_gemm_moe_overlap_probe.py b/probes/mk_gemm_moe_overlap_probe.py index da599c09..13501600 100644 --- a/probes/mk_gemm_moe_overlap_probe.py +++ b/probes/mk_gemm_moe_overlap_probe.py @@ -39,36 +39,17 @@ import torch # noqa: E402 -DEV = "cuda" -E, TOPK, HID, INTER = 288, 8, 4096, 512 # per-rank GLM-5.3 TP=4 geometry -T = 8 # C=1 verify batch (k=7 + 1) +# the served MoE fixture -- routing, weight sets, geometry -- is the go/no-go +# probe's; one definition of the C=1 decode MoE call for both probes +from moe_decode_stream_probe import ( # noqa: E402 + DEV, E, HID, INTER, T, TOPK, _routing, _weight_set) + U = 40 # unique experts a layer (27차) SHARED_N = 1024 # shared expert gate_up rows / rank SHARED_K = 512 # down-proj K / rank REPS = 20 -def _routing(U: int): - 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 _capture(fn_a, fn_b, order: str): """One graph: fork, A on stream a, B on stream b, join. `order` is the issue order inside the capture (A first or B first).""" @@ -136,11 +117,14 @@ def main() -> int: ones = torch.ones(E, dtype=torch.float32, device=DEV) ids, w = _routing(U) x = torch.randn(T, HID, dtype=torch.bfloat16, device=DEV) * 0.5 - out = torch.empty(T, HID, dtype=torch.bfloat16, device=DEV) + # the wrapper returns its output; the `out=` keyword exists only on the + # deployed b12x overlay, and this probe runs against whichever wrapper + # the container has (the bench runner mounts the megakernel files only) + print(f"wrapper {type(wrapper).__module__}") def moe(): wrapper.run(x, w13, sf13, w2, sf2, ids, w, w1_alpha=ones, - w2_alpha=ones, fc2_input_scale=ones, out=out) + w2_alpha=ones, fc2_input_scale=ones) mk.maybe_arm() assert mk._ARMED["gemm"], "MK-GEMM did not arm" diff --git a/tests/test_logic.py b/tests/test_logic.py index 6b3cc884..40339f81 100644 --- a/tests/test_logic.py +++ b/tests/test_logic.py @@ -8534,7 +8534,8 @@ def walk(n): "the raw nibble chunks land XOR-swizzled at copy time and the " "fragment loads read through the same swizzle (eight rows on a 64 B " "pitch would otherwise share two bank groups)") - check("pb[(size_t)r0 * c.n + cb] = acc[i][j][0];" in v2 + check("store_tile([&](int r, int col, float v) { pb[(size_t)r * c.n + col] = v; });" in v2 + and "if (r0 < c.m) { put(r0, cb, acc[i][j][0]); put(r0, cb + 1, acc[i][j][1]); }" in v2 and "atomicAdd(&g_mk2_tile_arrive[nt], 1u)" in v2 and "if (s_last) g_mk2_tile_arrive[nt] = 0u;" in v2 and "for (int s = 0; s < ksr; ++s) { // fixed order -> reproducible" in v2 @@ -8548,16 +8549,22 @@ def walk(n): and cu.index("if (mk_gemm2_on()) { // the non-persistent lane") < cu.index("c.grid = mk_resident_grid(mk_gemm_kernel, g_gemm_grid, GEMM_SMEM);") and "int mk_choose_ksr2(int m, int n, int k)" in cu + and "(g_mk_sms > 0 ? g_mk_sms : 48);" in cu + and "&g_mk_sms, cudaDevAttrMultiProcessorCount, 0));" in cu + and "if (slots % nblk == 0 && slots / nblk <= kmax) {" in cu + and "} else if (nblk * 2 > slots) {" in cu and "if (ksr > kblk) ksr = kblk;" in cu and "while (ksr > 1 && (size_t)m * n * ksr > (size_t)MK2_PART_ELEMS) --ksr;" in cu and consts["MK2_PART_ELEMS"] >= 32 * consts["KDA_INPROJ_N_PAD"] * 4, "the v2 lane is a kill-switched dispatch ahead of the persistent " - "launch (VLLM_GLM53_MK_GEMM2, default off); its slice rule keeps " - "every slice non-empty and the fp32 partial inside its buffer at " - "the widest per-rank linear") + "launch (VLLM_GLM53_MK_GEMM2, default off); its slice rule takes one " + "exact wave of the device's resident slots (SM count from the device, " + "not a constant), fine slices above half the slots, and keeps every " + "slice non-empty and the fp32 partial inside its buffer") check('"-DMK_NBUF2_DEF=" in pysrc_full' if False else "-DMK_NBUF2_DEF=" in pysrc_full and "_EXT.gemm2_plan(8, KDA_INPROJ_N, HIDDEN)" in pysrc_full - and 'ap.add_argument("--gemm2", choices=("0", "1", "both"), default="both")' in bench + and 'ap.add_argument("--gemm2", choices=("0", "1", "both", "env"), default="env")' in bench + and 'args.gemm2 = "1" if os.environ.get("VLLM_GLM53_MK_GEMM2") == "1" else "0"' in bench and "ext.set_gemm2(on, ksr)" in bench and "same = bool(torch.equal(got2, got))" in bench, "the driver builds v2's ring depth in, the boot fingerprint names " @@ -8746,7 +8753,7 @@ def walk(n): "the build attaches the W4 pack next to the deepgemm pair on every " "eligible linear, no arm knob") check("def probe_exact(gemm2: str = \"both\") -> bool:" in bench - and "def _probe_exact_lane(mk, tag: str) -> bool:" in bench + and "def _probe_exact_lane(mk, tag: str, x, p4, ref, n: int) -> bool:" in bench and "probe_w4" not in bench and "run_gemm_w4" not in bench and "build_mk_weight(" not in bench and "VLLM_GLM53_MK_W4" not in bench diff --git a/tools/trace_launch_positions.py b/tools/trace_launch_positions.py index ffce0b21..a555d42e 100644 --- a/tools/trace_launch_positions.py +++ b/tools/trace_launch_positions.py @@ -75,12 +75,31 @@ def main() -> int: keep = {k: a[k] for k in ("stream", "registers per thread", "shared memory", "grid", "block") if k in a} print(f"ATTR {e['name'][:56]:<56} {json.dumps(keep)}") - main_stream = collections.Counter(stream_of(e) for e in ev).most_common(1)[0][0] + # the main stream is the one carrying the model's forward: the stream + # with the most kernel TIME (the routed MoE kernels alone are ~half the + # step). A kernel-count heuristic picks the glue stream, and the step + # anchor's stream is not it either -- the 09-01 stock trace has the prep + # kernels and the shared-expert pair on 17 and the forward on 210, the + # armed 09-04 trace the other way round. + busy = collections.Counter() + for e in ev: + busy[stream_of(e)] += e["dur"] + main_stream = busy.most_common(1)[0][0] + print("streams (kernels, busy ms over the trace): " + + ", ".join(f"{k}: {sum(1 for e in ev if stream_of(e) == k)}, {v / 1e3:.0f}" + for k, v in busy.most_common())) rows = [] - exposed_by = collections.defaultdict(list) - for si, (a, b) in enumerate(zip(starts[:-1], starts[1:])): + # per step, per kernel family: exposed us; a family absent from a step + # contributes 0 to that step (a median over occurrences only would let a + # single prefill step dominate, as it did for the MoE kernels) + exposed_by = collections.defaultdict(dict) + nsteps = 0 + # step 0 is the profiler's warm-up (a prefill and the first decode): skip + for si, (a, b) in enumerate(zip(starts[1:-1], starts[2:]), start=1): + nsteps += 1 seg = ev[a:b] seg_ts = [e["ts"] for e in seg] + lookback = max(e["dur"] for e in seg) # any kernel still running at s0 by_stream = collections.defaultdict(list) for e in seg: by_stream[stream_of(e)].append(e) @@ -95,13 +114,13 @@ def main() -> int: if stream_of(e) != main_stream: expo[e["name"][:30]] += _exposed(main_ivs, e["ts"], e["ts"] + e["dur"]) for k, v in expo.items(): - exposed_by[k].append(v) + exposed_by[k][si] = v gemms = [e for e in seg if pat in e["name"]] for gi, e in enumerate(gemms): p = prev[id(e)] - gap = e["ts"] - (p["ts"] + p["dur"]) if p else float("nan") + gap = e["ts"] - (p["ts"] + p["dur"]) if p else None # first on its stream s0, e0 = e["ts"], e["ts"] + e["dur"] - lo = bisect.bisect_left(seg_ts, s0 - 5000) + lo = bisect.bisect_left(seg_ts, s0 - lookback) hi = bisect.bisect_right(seg_ts, e0) ov = ov_moe = 0.0 for o in seg[lo:hi]: @@ -113,7 +132,8 @@ def main() -> int: if "moecute" in o["name"] or "moe_static" in o["name"]: ov_moe += x rows.append({"step": si, "idx": gi, "dur": e["dur"], "stream": stream_of(e), - "gap": round(gap, 1), "ov_moe": round(ov_moe, 1), + "gap": round(gap, 1) if gap is not None else None, + "ov_moe": round(ov_moe, 1), "ov_other": round(ov - ov_moe, 1), "prev": p["name"][:40] if p else ""}) if not rows: @@ -141,16 +161,20 @@ def main() -> int: d = sorted(r["dur"] for r in rs) md = statistics.median(d) tot += md + gaps = [r["gap"] for r in rs if r["gap"] is not None] + mg = statistics.median(gaps) if gaps else float("nan") print(f"{i:3d} {md:8.1f} {d[len(d) // 10]:7.1f} {d[len(d) * 9 // 10]:7.1f} " - f"{statistics.median(r['gap'] for r in rs):8.1f} " + f"{mg:8.1f} " f"{statistics.median(r['ov_moe'] for r in rs):7.1f} " f"{statistics.median(r['ov_other'] for r in rs):7.1f} " f"{collections.Counter(r['stream'] for r in rs).most_common(1)[0][0]:>4} " f"{collections.Counter(r['prev'] for r in rs).most_common(1)[0][0]}") print(f"sum of per-position medians = {tot / 1e3:.2f} ms/step") print(f"\nnon-main-stream kernels (main = stream {main_stream}): exposed us/step, median over steps") - for k, v in sorted(exposed_by.items(), key=lambda kv: -statistics.median(kv[1]))[:12]: - print(f" {statistics.median(v):8.1f} {k}") + def _med(d): + return statistics.median(list(d.values()) + [0.0] * (nsteps - len(d))) + for k, v in sorted(exposed_by.items(), key=lambda kv: -_med(kv[1]))[:12]: + print(f" {_med(v):8.1f} {k}") return 0