Repository navigation
b200_diffusion_residual_gate_add__multi_shape: fused CUDA kernel (~2.19x vs eager, 67/67 correct) - #134
Merged
BBuf merged 12 commits intoJun 25, 2026
Conversation
…t scaffold Foundation for the B200 residual-gate-add / 4D broadcast-add kernel task. - Recover upstream lineage + semantics from sgl-project/sglang main @ 67b2a9ed0cfba8ec625d3f26548e502646fd914d: residual + update*gate (LTX-2/Ideogram4/FLUX.2) and the LTX 4D cross-attn-gate broadcast add. Provenance in docs/baseline_source.md + baseline/upstream_residual_gate_snippets.md. - Faithful PyTorch-eager baseline (two-launch mul+add with cached scratch; single eager add for the 4D row) in baseline/binding.py. - Fused candidate CUDA kernels (residual_gate_add full+broadcast gate; separate broadcast_add_4d) in solution/kernel.cu: 128-bit vectorized fast path + scalar tail/fallback, fp32 one-round, out-of-place contiguous with alias/shape/dtype rejection. Built via tvm-ffi in solution/build.py. - Symmetric destination-passing ABI; bench/adapter.py shared dispatch, no sglang import at runtime; bench/correctness.py fp32 oracle + candidate-vs- baseline + poison + regression grid + rejection tests; bench/benchmark.py verbatim from the standalone template; bench/workloads.json frozen 8 rows. - config.toml reconciled to plural binding-based entry points. Build/correctness/benchmark verification on remote B200 (ion-b200) is the immediate next step (not yet run; no local CUDA). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…(R1) Addresses Codex Round-0 review blocking findings: - Tighten row-broadcast gate validation in solution/kernel.cu: require same-rank [.,1,D] with all leading dims == 1; reject malformed [B,D]/[D] and batch>1 broadcast (the i%D indexing assumes exactly D contiguous elements). - Reject batch>1 in broadcast_add_4d (a is reused per row; production is B=1). - Add a shared adapter-level validator (bench/adapter.py::_validate) invoked in _dispatch so baseline (eager) and candidate reject the same malformed inputs with identical overhead; update bench/correctness.py rejection tests to run on BOTH sides and add leading-dim!=1 and batch>1 cases. - Add an optional extra_provenance() adapter hook + minimal additive merge in bench/benchmark.py (baseline commit, candidate sha256, tvm-ffi/nvcc versions, REMOTE_GPU_ID); timing harness unchanged. - Refresh upstream lineage pin to latest main 8314247d (residual-gate / 4D broadcast expressions verified byte-identical to the earlier 67b2a9ed). - Remove plan-workflow markers (DEC-*) from code comments per the code-style note. Build/correctness still pending remote B200 (no local CUDA). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ker cleanup (R2) Addresses Codex R1-review blocking finding + queued cleanups: - Full-shape gate must be contiguous: add the check to bench/adapter.py::_validate (full-gate branch) and solution/kernel.cu (before GateMode::kFull). A non-contiguous full gate was previously accepted and silently mis-read by the scalar path; now rejected on both sides. Add rga-full-gate-noncontig both-side rejection test in bench/correctness.py. - Correct stale baseline SHA in solution/kernel.cu header (67b2a9ed -> 8314247d). - Remove plan DEC-1 marker from config.toml comment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Addresses Codex R1-review blocking #2 (AC-6). bench/adapter.assert_pinned_gpu(), called from make_case, fails closed when KDA_REQUIRE_PINNED_GPU=1 unless REMOTE_GPU_ID is set and is the sole visible CUDA device (cuda:0), so benchmark numbers are always attributable to one pinned idle GPU. No-op for local runs. Remote evidence (ion-b200, GPU 2 idle): candidate builds via tvm-ffi and bench/correctness.py --impl both --rows all PASSES 56/56 (AC-3/AC-4 verified). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… 2.199x (R2) docs/run_log.md records the ion-b200 run (container sglang_bbuf, GPU 2 pinned via REMOTE_GPU_ID/CUDA_VISIBLE_DEVICES, idle before/after 0%/0MiB): - Build: solution/kernel.cu compiles via tvm-ffi (both exports). - Correctness (AC-3/AC-4): bench/correctness.py --impl both --rows all = 56/56 PASS (production + grid + poison + both-side rejection incl. full-gate-noncontig). - Benchmark (AC-5): all 8 rows PASS, candidate vs faithful eager baseline geomean 2.199x (broadcast-gate rows 2.8-3.3x; full-gate 1.6-1.8x; 4D add 1.20x). - Provenance: baseline 8314247d, candidate sha256 a450f863, tvm-ffi 0.1.9, nvcc 13.0, B200, torch 2.11.0+cu130. - Preliminary roofline: largest bcast row ~5.67 TB/s (~71% of HBM3e peak). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…/results.md (R3) Completes the final result (AC-4/AC-6/AC-7) and analyze pass: - AC-4: add deterministic zero/sign rows + repeated-seed rows to bench/correctness.py; strict-pinned rerun on B200 (REMOTE_GPU_ID=0) = 67/67 PASS. - AC-7: NCU (--set basic) on 3 representative rows + per-row roofline table in docs/results.md. Bounds: large full-gate DRAM-bound (NCU 72.1%, near roofline); large broadcast-gate SM/occupancy-leaning (SM 69.8% > DRAM 59.5%); small rows launch/grid-bound (grid 126 < 148 SMs); 4D launch/low-fusion (inferred). - Analyze (task3/task7/task10 via Codex): baseline faithfulness + fp32 one-round numerics + out-of-place/contiguous policy all ACCEPT; ranked optimization = broadcast-gate modulo removal (modest/conditional); GO decision accepts the generic fused candidate as final (geomean 2.199x; residual-gate-only 2.397x; call-count-weighted 1.699x). - Doc-drift fixes: benchmark.py "verbatim" wording corrected (additive provenance hook noted); removed DEC-* markers from docs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… doc drift (R4) Addresses Codex R3-review AC-8 hygiene: - Remove AC-4 marker text from bench/correctness.py comments (coverage unchanged). - benchmark_method.md: secondary diagnostics + roofline are reported in docs/results.md (not "planned"); drop AC-7/AC-6 markers from the doc. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Addresses Codex R3-review AC-6 (split-GPU evidence). Reran correctness, full benchmark, and the 3 NCU profiles all on one idle pinned GPU (KDA_REQUIRE_PINNED_GPU=1 REMOTE_GPU_ID=7 CUDA_VISIBLE_DEVICES=7), idle before/ after 0%/0MiB: - correctness 67/67 PASS; benchmark 8/8 PASS geomean 2.193x (results.jsonl remote_gpu_id=7); NCU full_big 72.5% DRAM, bcast_big 59.5%/SM 69.7%, full_small 3.8%/grid-bound; raw .ncu-rep saved under ignored /tmp. - docs/results.md + docs/run_log.md updated to the unified GPU-7 numbers; corrected the NCU-retention note (R4 saves .ncu-rep). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…rank) (R5) - [P2] bench/benchmark.py: call torch.cuda.set_device(device) before importing the adapter in _run_one_workload, so a non-default --device builds/reports the candidate for the intended GPU (adapter builds the candidate at import). - [P3] solution/kernel.cu broadcast_add_4d: require out.ndim()==4 before the per-dim size check, rejecting a >4D out with a matching 4-dim prefix. - Both fixes are host-side only; the device kernels + timed path for valid inputs are unchanged, so the Round-4 geomean 2.193x stands. Rebuilt on B200; correctness re-confirmed 67/67 (new candidate sha256 27f67c5f). Benchmark re-run deferred (no idle GPU at fix time); docs/results.md + docs/run_log.md updated with the hash change + compute-neutrality note. benchmark_method.md note updated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Addresses Codex R6 review (3x P2): - [P2] device-before-import: make the candidate build lazy in bench/adapter.py (_candidate_fns(), lru-cached) so importing the adapter is cheap/device- independent and the device-querying tvm-ffi build runs on first candidate use, after the runner's torch.cuda.set_device. Fixes correctness.py (and hardens benchmark.py) import ordering at the root. - [P2] batched [B,1,D] gate + [P2] batched 4D add: narrow the documented contract to batch-1 (every production row is B=1). B>1 broadcast stays rejected symmetrically on both candidate and eager baseline (shared _validate) and is covered by rga-gate-leaddim-not1 / bcast-batch-gt1 tests; documented as out-of-scope in baseline_source.md / results.md / benchmark_method.md. Avoids adding per-vector batch arithmetic to the bandwidth-critical path for a shape absent from the frozen workload. solution/kernel.cu unchanged (hash 27f67c5f); Round-4 geomean 2.193x stands. Correctness re-confirmed 67/67 on B200 with the lazy-build adapter. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… (R7) Addresses Codex R7 review [P2]: residual_gate_add / broadcast_add_4d grabbed the process-current device's stream and launched there even if the tensors lived on another CUDA device (wrong-GPU launch / illegal access on multi-GPU hosts). Both exported functions now validate all operands share one CUDA device id and install c10::cuda::CUDAGuard(dev) + take that device's stream (at::cuda::getCurrentCUDAStream(dev)) before launching. Single-visible-GPU runs are unaffected (guard is a no-op when the tensor device is already current), so the Round-4 geomean 2.193x stands. Rebuilt on B200; correctness re-confirmed 67/67 (new candidate sha256 a6ab9c86). docs/results.md + run_log.md updated with the hash + compute-neutrality note. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Addresses Codex R8 review [P2]: bench/benchmark.py main() called _provenance() before any set_device. _provenance() reads _gpu_name() (get_device_name on the process-current device) and the adapter's extra_provenance()/candidate build, so a non-default --device recorded cuda:0's GPU name / SM capability / compile flags in results.jsonl while the workloads later ran on args.device. main() now calls torch.cuda.set_device(torch.device(args.device)) before _provenance(), mirroring the R5 _run_one_workload device-first pattern. Harness-only; solution/kernel.cu is unchanged (hash a6ab9c86). Re-validated fresh on idle B200 GPU 7 (strict pin): correctness 67/67 PASS, benchmark 8/8 PASS geomean 2.186x (reproduces the Round-4 2.193x within run-to-run noise), and results.jsonl provenance now records gpu=NVIDIA B200, remote_gpu_id=7. docs/run_log.md + docs/results.md updated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
7 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements the standalone B200 diffusion kernel task
b200_diffusion_residual_gate_add__multi_shape, fusing the memory-boundresidual-gate update into a single CUDA pass:
residual_gate_add(residual, update, gate, out)→out = residual + update * gate,supporting both full
[.,L,D]and row-broadcast[.,1,D]gate modes.broadcast_add_4d(a, b, out)→ 4D broadcast add (a=[1,1,P,D]overb=[1,S,P,D]).The candidate replaces the faithful PyTorch-eager production path (two launches,
mul+add, plus an intermediate tensor) with one fused, 128-bit-vectorizedgrid-stride kernel (scalar tail; fp32 one-round accumulation; no fast-math),
exposed through a symmetric destination-passing ABI built via tvm-ffi.
Results (NVIDIA B200, one pinned idle GPU)
faithful eager baseline (bf16 atol=rtol=5e-2); plus a regression grid (gate modes,
bf16/fp16/fp32, odd/non-vec-aligned/small tails, zero/sign rows), a poison
self-test, and both-side ABI rejection tests.
frozen production rows (min 1.18x, max 3.33x; residual-gate-only 2.396x). The win
is removing one kernel launch, the intermediate temp's full write+read, and one
Python dispatch.
rows reach 71–81% of HBM3e peak (DRAM-bound, near roofline); small rows are
launch/grid-bound (grid < 148 SMs).
Full evidence:
docs/results.md,docs/run_log.md,docs/benchmark_method.md,docs/baseline_source.md.Scope & safety
diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/.inputs); multi-GPU device guard (launches on the input tensors' device); documented
batch-1 contract (B>1 broadcast is out of scope and rejected + tested).
🤖 Generated with Claude Code