feat(rust): Laguna batched prefill, 156 to 559 pp tok/s, plus weight cache and dtype-cache fix#70
Merged
Conversation
… 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.
…, 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. Reference C++ implementation on identical weights and box: 663 and 660 stock.
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.
TheTom
marked this pull request as ready for review
July 23, 2026 20:44
TheTom
added a commit
that referenced
this pull request
Jul 23, 2026
…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.
TheTom
added a commit
that referenced
this pull request
Jul 23, 2026
…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).
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
Three commits on top of the Laguna decode PR:
Numbers (GB10, single stream, identical weights and box as the reference C++ implementation)
Prefill went 156 to 559 through profiled lever iterations (Marlin projections, MoE fusion, scheduling, chunk tuning); each lever gated on byte-identical greedy continuations against the reference oracle. Levers measured and REJECTED along the way (kept out of the tree): fp8-activation grouped MoE GEMM (slower at these tiny-expert shapes), native-SASS NVRTC for prefill (no delta; PTX stays default and measured slightly better for decode).
Tests
Kernel unit tests: batched rope (yarn + plain), batched softplus gate, windowed varlen skip vs mask-only, fused swiglu gather vs unfused compose, scheduling A/B on skewed synthetic groups, non-64-aligned column copy. Integration: env-gated greedy-continuation identity for decode-only vs prefill-then-decode.
Notes