feat(rust): Laguna decode + batched prefill on dev (re-landed from #69/#70)#73
Merged
Conversation
#69) (renamed) * sync(rust): land in-flight CUDA engine work from the GB10 working tree Carries the uncommitted engine state the GB10 box has been running: CUDA-graph capture trait plumbing (begin/end capture, graph_launch), fp8 projection microbench and MoE grouped-MMA test updates, and the device trait additions they depend on. Precedes the Laguna port commits that build on these interfaces. * feat(laguna): Laguna-S-2.1 decode on the CUDA engine Full single-stream decode for the 117.55B/8.14B-active hybrid MoE: sigmoid top-10 router with score-correction bias and shared expert, per-head softplus attention gate, per-head QK RMSNorm, period-4 full/sliding-window attention with a 512-slot ring KV cache, YaRN partial-rotary rope on full layers (scaling read from GGUF metadata, this export is a 256K/factor-32 checkpoint), Q8 LM head, and a Q4-requantized weight load straight from GGUF. New ops: rope_yarn_partial (+posbuf variant), gate_softplus_mul_perhead, argmax_f32_device, sdpa_decode_nbuf (n_kv from a device buffer), write_u32, moe_gather_q4 and moe_gather_q4_swiglu wrappers (kept as always-error stubs in this rebrand — the iron_moe_gather_q4 / iron_moe_gather_q4_swiglu kernels aren't in thewafflehaus/iron@dev yet). Decode runs eagerly or as a captured CUDA graph (BUTTER_LAGUNA_GRAPH=1, warmup-capture-replay via LagunaDecodeCtx with a persistent workspace), with optional fused MoE gate+up (BUTTER_LAGUNA_FUSE=1) and micro-fusions (BUTTER_LAGUNA_MICRO=1: concatenated QKVG projection, shared expert via the fused gather, o_proj residual accumulate). GB10 single-stream greedy: 21.2 tok/s reference baseline, 36.5 tok/s here (graph+fuse+micro), greedy tokens verified identical to the reference implementation. Env-gated integration smoke: BUTTER_LAGUNA_GGUF plus laguna tests in wh-butter-cuda (unit kernels run without weights).
…cache and dtype-cache fix (#70) (renamed) * fix(cuda): key compiled-module and shared-size caches by kernel dtype signature The ops layer caches kernel IR by (name, dtype) but the backend cached compiled modules by bare kernel name, so the first dtype to touch a name won and every other dtype silently ran the wrong binary: wrong element stride, out-of-bounds reads, and in the shrinking-stride direction silent corruption. Found when the prefill path's first f16 gather inherited the decode path's f32 module. Shared-memory sizing had the same hazard. * feat(laguna): batched prefill with tensor-core projections, fused MoE, CTA scheduling Chunked multi-token prefill (default chunk 2048): batched YaRN/plain rope, per-query windowed varlen attention with a window-aware KV-block skip, linear sliding-window scratch compacted into the decode ring, grouped-GEMM MoE with on-device descriptors, Marlin W4A16 tensor-core dense projections (concatenated QKV, o_proj, dense FFN, shared expert), fused gate+up expert stacks, and grouped-GEMM CTA scheduling (descending-size expert order plus N-banded CTA order for weight L2 reuse, default on). Correctness gates: prefill-then-decode greedy continuation byte-identical to decode-only; last-token argmax matches the reference oracle; kernel unit tests for the batched rope, batched gate, windowed varlen skip, fused swiglu gather, and scheduling A/B on skewed synthetic groups. GB10 single-stream: prefill 559 tok/s at 2048 (was 156 at first light), 504 at 8192; decode unchanged at 36.3 via graph replay. The reference C++ engine on identical weights and box: 663 and 660 stock. * feat(laguna): on-disk weight cache and windowed parallel conversion Content-keyed cache of every converted engine-format weight blob (per tensor artifact, keyed by format version, source GGUF identity, and the load-shaping env flags), written on first conversion and mmap-read on later loads. Conversion itself runs rayon-parallel over a bounded window of layers (full parallelism held tens of GB of transients and got OOM-killed on the shared 128GB). Warm reload: 56s, down from ~7.5 minutes; cache hits and misses are reported at load end. Also carries the comment hygiene sweep across the Laguna files (dash style, neutral phrasing for external references) and the stale decode-only module doc fix.
…gates Both cherry-picked commits (#69, #70) brought every op they need with them — no separate ops-porting commit was required. These are gate-driven fixups only: - wh-butter-cuda/src/imp.rs: drop the moe_marlin_gemm/marlin_repack/ marlin_build_routing pass-throughs PR #70 added — they call through to methods on wh_iron_runtime::CudaDevice that don't exist in thewafflehaus/iron@dev (confirmed: no `marlin` symbol anywhere in that repo). Falls back to the wh-butter-core default "unsupported on this backend" stub (already present, unaffected) instead of failing to compile. Same root cause as the moe_gather_q4/moe_gather_q4_swiglu stubs already documented in wh-butter-ops. - wh-butter-cuda/tests/all_models.rs: dropped a `smallmodel_fuse_slicecast_ab` test calling `wh_butter_modeltests::bench_smallmodel_fuse_slicecast` — confirmed that function doesn't exist anywhere in wh-butter-modeltests on the original tom/feat/cuda-hip-vulkan-backends branch either (pre-existing dead reference predating this cherry-pick, not something it introduced). - wh-butter-cuda/tests/f16norm_f32in.rs: dropped — same pre-existing-bug class, references a `add_rms_norm_f16norm` op that was never implemented on the source branch. - rust/Cargo.lock: regenerated via `cargo build` (not hand-edited).
BUTTER_MARLIN_FIXTURE_DIR selects the reference-fixture dir; unset gives a quiet [skipped] instead of unwrapping a hardcoded /tmp path (repo convention: fixture/model paths come from env vars, absent -> skip).
iron@dev's iron_strided_col_copy takes 5 bindings and has no internal idx guard (the guarded 6-binding variant only exists on an un-landed kernels feature branch). Drop the extra total binding and cover exactly s*width threads, using the largest power-of-two block (<=64) dividing the total so Laguna's non-64-aligned gate-column shapes stay correct.
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.
What
Re-lands Laguna-S-2.1 decode (#69) and Laguna batched prefill + weight cache + dtype-cache fix (#70) onto dev — they were merged into
tom/feat/cuda-hip-vulkan-backendsby mistake. Cherry-picked, swept to the Butter/Iron names.Stacked on #71 (base =
rename/butter-branding). Merge #71 first; this then retargets to dev. Split out of #71 per Tom to keep the rename PR pure.c7e6a84e— feat(rust): Laguna-S-2.1 decode on the CUDA engine, 36.5 tok/s on GB10 #69 Laguna-S-2.1 decode, 36.5 tok/s on GB10 (renamed)b57c5415— feat(rust): Laguna batched prefill, 156 to 559 pp tok/s, plus weight cache and dtype-cache fix #70 Laguna batched prefill 156→559 pp tok/s + weight cache + dtype-cache fix (renamed)a073b097— fixups surfaced by gates (see notes)16947ab2— env-gate the marlin weight-prep fixture test (BUTTER_MARLIN_FIXTURE_DIR, unset → quiet skip)FYI @ekryski.
Notes for review
moe_marlin_gemm/marlin_repack/marlin_build_routing) fall back to the default "unsupported" stub — theirCudaDeviceruntime methods aren't inthewafflehaus/iron@devyet (implementations recovered, being validated on the 5090; iron-side PR to follow, then these get un-stubbed). Laguna itself never calls them — its W4A16 grouped GEMM dispatches raw kernels directly; only the NemotronH path is affected and it errors cleanly.bench_smallmodel_fuse_slicecast,add_rms_norm_f16norm) — they broke the CUDA test-crate build on the source branch too.Verification (M5 Max + agents)
cargo build --workspaceclean;--features cuda/--features vulkantypecheck cleancargo test(ops/models/loader/core/modeltests/metal): all pass; skips = pre-existingdsv4_layerMetal hang (tracked separately) + the now-env-gated fixture testswift build --build-testsclean; full Swift unit suite 1671 tests green on the combined tree (one Metal microbench flaked under 3-way parallel load, passes in isolation)