perf(osar): all-reduce 5커널 → 3커널 (가드·웨이트 프롤로그 융합) - #89
Conversation
디코드 스텝의 고정비를 실측으로 분해했더니 커널 개수가 지배적이었다.
step = 42.0 ms (고정) + 0.706 ms x distinct_experts
가변부는 실측 대역폭 223.4 GB/s 의 91% 라 짜낼 게 없다. 고정 42.0 중 dense
가중치 읽기가 9.9 ms 이고 나머지 32.1 ms 는 메모리가 아니다. 같은 스텝에서
커널이 2,210 개 뜨니 커널당 15.4 us — 32.1 ms 를 정확히 설명한다(전부 커널
개수에 귀속시킨 상한이다).
그 2,210 중 우리 all-reduce 가 522 개(23.6%)로 단일 최대 항목이다. 한 콜렉티브가
다섯 번 뜬다: k_guard, k_copy_in, k_signal, k_wait, k_reduce. 그런데 guard 와
wait 는 각각 <<<1,32>>> 로 전역 플래그만 읽고 도는 스핀이다.
- k_guard 의 조건은 tx_seq/ack_seq — 모든 블록이 같은 값을 보고, 같은 커널의
다른 블록을 기다리지 않는다. k_copy_in 프롤로그로 접어도 미상주 블록이
교착을 만들 수 없다. ack_seq 는 단조 증가라 한 번 참이면 계속 참이다.
- k_wait 의 조건은 피어의 rxf 플래그 — 원격 RDMA 쓰기로만 바뀐다. 같은 이유로
k_reduce 프롤로그로 접힌다. 펜스는 있던 자리(대기 뒤, 피어 데이터 읽기 앞)에
그대로 둔다.
k_signal 은 남긴다. k_copy_in 의 **모든** 블록이 끝난 뒤에만 돌아야 하는데,
협조적 런치나 원자 카운터 없이는 커널 안에서 그리드 전체 완료를 알 수 없고,
카운터는 cudagraph 재생이 못 보는 곳에서 리셋해야 한다.
콜렉티브당 5 -> 3, 스텝당 -209 커널(9.4%). 15.4 us 상한이면 3.2 ms/step.
안전: 부팅마다 도는 _self_test 가 one-shot 을 NCCL 과 비교(div<=0.5)하고
실패하면 NCCL 로 물러난다 — 융합이 틀려도 서빙이 오염되지 않고 로그에 남는다.
다만 self-test 는 콜렉티브 한 번이라 링 공간 버그(RING 번 뒤에야 발화)는 못
잡는다. 그건 긴 생성이 멀쩡한지로 판정한다.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: serverGenReqId_d13831a6-e1bd-4422-a7c7-1e9601878c2f) |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Left a non-blocking comment: Cursor Bugbot was present but skipped (usage limit), so this run is not approving. Human review is needed; requesting a reviewer on the changed one-shot all-reduce path.
Sent by Cursor Approval Agent: Pull Request Router and Approver
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0107e4f986
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // the same argument: it polls the peers' inbound flags, which are globals, so | ||
| // every block can wait on them by itself. The fence stays where it was -- after | ||
| // the wait, before anything reads peer data. | ||
| __global__ void k_reduce(Ctrl *c, const bf16 *src, bf16 *dst, int n) { |
There was a problem hiding this comment.
For any multi-block input, folding k_wait into this kernel also moves its unconditional __threadfence_system() into every reduction block: the self-test now executes it from 24,576 threads and a MAXEL input from 65,536 threads, instead of the 32 threads in the former k_wait<<<1,32>>>. This turns a fixed one-warp synchronization step into grid-sized system-fence work on every collective and can erase the launch-overhead reduction this change is intended to provide; retain a bounded waiter/fence mechanism rather than running it from every element-processing thread.
Useful? React with 👍 / 👎.
| // a block of this same kernel, so a block that is not resident yet cannot | ||
| // deadlock us. Folding it in removes one launch per collective; at ~104 | ||
| // collectives a decode step that is ~104 kernels of the 2,210 a step runs. | ||
| __global__ void k_copy_in(Ctrl *c, const bf16 *src, int n) { |
There was a problem hiding this comment.
When an eligible tensor has numel() == 0, py_oneshot computes grid == 0, so both kernels containing the newly fused protocol checks launch with no blocks and neither the ring-space guard nor the peer wait can run; the fixed one-block k_signal launch can nevertheless advance tx_seq. The former dedicated <<<1,32>>> guard and wait did not disappear for this input, so an empty collective can now bypass flow control, overwrite a still-active ring slot under backlog, or let later collectives advance before peers have completed the same sequence. Reject empty inputs or run the protocol prologues from a nonzero launch.
Useful? React with 👍 / 👎.
#89 가 k_signal 을 분리해 둔 이유는 "원자 카운터는 cudagraph 재생이 못 보는 곳에서 리셋해야 한다"였다. 그러나 카운터가 단조 증가하고 모든 런치가 정확히 ARGRID 번 증가시키면(그리드 고정) 리셋이 필요 없다: 스트림 순서 덕에 커널 진입 시점마다 done_ctr 은 ARGRID 의 배수이고, atomicAdd 반환값 % ARGRID == ARGRID-1 인 블록이 바로 이번 런치의 256번째 완료 블록이다. 그 블록은 이번 런치의 모든 블록이 copy 를 마친 뒤에만 존재한다 — 모든 블록은 자기 copy 전에 tx_seq 를 읽으므로, 마지막 블록이 tx_seq = s+1 을 게시해도 늦게 상주한 블록이 잘못된 슬롯을 읽을 수 없다. 가드(tx_seq/ack_seq)와 피어 대기(rxf)는 #89 와 같은 논증으로 커널 안에서 블록별로 돌린다. 커널 구조 (k_oneshot, <<<ARGRID=256, 256>>> 고정): 가드 스핀 → copy → fence → atomicAdd → 마지막 블록이 nbytes+tx_seq 게시 → rxf 대기 → fence → reduce 콜렉티브당 3 -> 1, 스텝당 -208 커널(9.4%), 15.4 us/커널 상한이면 -3.2 ms/step. 빈 블록 비용은 가드 읽기 1회·시스템 원자 1회·스핀뿐이다. 안전장치는 #89 와 동일: 부팅 _self_test 가 one-shot vs NCCL 을 div<=0.5 로 채점하고 실패 시 NCCL 폴백. 판정 프리셋(부팅 로그)은 [osar] self-test PASS + 트레이스에 k_oneshot 만 존재(k_guard/k_wait/k_copy_in/k_signal/k_reduce 부재). 링 공간 버그는 self-test 가 못 잡으므로 긴 생성 일관성으로 보강.
셀프테스트 게이트(self-test PASS + real=True)는 NCCL이냐 osar냐만 가르고 어느 판본이냐는 안 가린다. 런처가 마운트하는 건 레포가 아니라 PROFILE_OVERLAY_DIR 이고 deploy-overlays.sh 만 그걸 채우므로, 배포를 우회한 검증 부팅은 융합 안 된 5커널 베이스라인인 채로 판정 전제를 통과해 버렸다 (#73·#83·#84·#69·#88 에 이은 여섯 번째 같은 부류). _build() 가 마운트된 .cu 를 컴파일하기 직전에 md5 8자리 + __global__ 개수를 WARNING 으로 남긴다: kernels=5 → #89 이전, 3 → #89, 1 → #90. 부팅 로그 한 줄로 판본 판정이 끝나고, 배포 검증(md5 정합)의 부팅 후 확인 수단이 된다.
C=1 스텝의 커널 1,886 개 중 우리 소유는 186 개(9.9%)뿐이다. "우리 커널만 최적화 대상" 이라는 전제가 잘못됐다는 것을 코드로 보이고, 나머지 1,700 개의 출처·결정 지점·수정 가능성을 적는다. 이 레포의 오버레이 체계 자체가 벤더 코드를 접수해 고치는 물건이다. - 아키텍처가 커널 개수를 설명한다: 하이브리드(KDA 34 + 희소어텐션 11), MoE 42 층, MHC 전 층. gated_delta_rule 34 · causal_conv1d 34 · mhc 2x90 · MoE 42 · 게이트 42 · MLA 33 이 전부 이 구조와 맞아떨어진다. 부수로 k_oneshot 102 가 콜렉티브/스텝을 실측 확정한다(앞서 추정한 73 은 틀렸다) - elementwise 483(25.6%)이 최대 덩어리이고 원인은 custom_ops 융합 장벽으로 보인다 -- 이 스택에서 A/B 된 적이 없다. 부팅 1 회·코드 변경 0 - 잔여 bf16 GEMM 145 개의 정체를 _INCLUDE 정규식과 대조해 특정: 저랭크 사영의 b 절반(f_b/g_b 68, q_b/kv_b 22)과 인덱서(~24). 로더가 q/k/v/b/f_a/g_a 를 in_proj_qkvbfg_a 로 병합하므로 앞쪽은 이미 잡힌다 - W8A8 의 대가가 보인다: per_token_group_quant 179/스텝 - b12x 의 제약 두 개를 코드로 인용: 가중치 NVFP4 고정(W4A8 불가), EP 미지원 (288 전문가 전 랭크 복제 = 스텝 62% 가 샤딩 없는 읽기) - 닫힌 축 4 종과 그 사유. 교훈: pass_config 확정값이 True 라도 런타임 지원 게이트가 한 겹 더 있다 -- 판정 전제는 확정값이 아니라 경고 부재다 - 커널 축 천장 15.4%. 다만 근거가 약해지고 있음을 명시: #89 이후 324 개를 더 없앴는데 스텝이 안 변했다. 5.4us/커널은 상한이다 census.py 동봉 -- 스텝 분모를 추정하지 않고 센다(앞선 인구조사가 추정해서 틀렸다). tests/test_logic.py 671 검사 통과. 문서·분석기만, 코드 무변경. Co-authored-by: choiceoh <astra7471@gmail.com>
Summary
디코드 고정비 42.0 ms 중 32.1 ms 가 비-메모리이고, 커널 2,210개 × 15.4 µs 가 그걸 정확히 설명한다. 그 중 우리 all-reduce 가 522개(23.6%)로 단일 최대 항목이라 콜렉티브당 커널 수를 5→3 으로 줄인다.
Changes
k_guard→k_copy_in프롤로그. 조건이tx_seq/ack_seq전역 읽기뿐이라 블록마다 독립적으로 스핀 가능하고, 같은 커널의 다른 블록을 기다리지 않으므로 미상주 블록이 교착을 못 만든다.k_wait→k_reduce프롤로그. 조건이 피어의rxf플래그(원격 RDMA 쓰기)라 같은 논증이 적용된다.__threadfence_system()은 원래 자리 유지.k_signal은 분리 유지 —k_copy_in의 전 블록 완료를 커널 안에서 알 수 없고, 원자 카운터는 cudagraph 재생이 못 보는 곳에서 리셋해야 한다.스텝당 −209 커널(9.4%), 15.4 µs/커널 상한이면 3.2 ms/step.
Verification
_self_test가 one-shot vs NCCL 을div<=0.5로 채점하고 실패 시 NCCL 폴백 → 융합 오류가 서빙을 오염시킬 수 없다.[osar] self-test PASS, 그리고 커널 트레이스에k_guard/k_wait부재 +k_copy_in/k_signal/k_reduce존재.🤖 Generated with Claude Code
Note
Medium Risk
Changes GPU synchronization ordering for RDMA all-reduce; correctness depends on per-block spinning and fences, and ring-depth bugs may not show in single-collective self-tests.
Overview
Cuts one-shot all-reduce from five CUDA kernels per collective to three by merging synchronization into the copy and reduce kernels, targeting decode fixed overhead where this path dominates kernel count.
k_guard(ring-space wait ontx_seq/ack_seq) is now a block-0 spin +__syncthreads()at the start ofk_copy_in, so every block can poll the same globals without a separate<<<1,32>>>launch.k_wait(peerrxfflags) is folded the same way intok_reduce, with__threadfence_system()still after the wait and before reading peerrxdata.k_signalstays a standalone launch: signaling must happen only after allk_copy_inblocks finish, which this design does not synchronize inside one kernel without cooperative groups or atomics—and atomics would need resets invisible to CUDA graph replay.oneshot_arnow enqueuesk_copy_in→k_signal→k_reduceonly (~−2 kernels per collective, ~−209 kernels/step at ~104 collectives).Reviewed by Cursor Bugbot for commit 0107e4f. Bugbot is set up for automated code reviews on this repo. Configure here.