feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints#250
Open
rsafier wants to merge 1 commit into
Open
feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints#250rsafier wants to merge 1 commit into
rsafier wants to merge 1 commit into
Conversation
08b1e63 to
31c1b3e
Compare
This was referenced Jul 4, 2026
rsafier
added a commit
that referenced
this pull request
Jul 5, 2026
Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used
by the MoE loader) into the dense loader SSM arm so dense Holo models
(0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense
SSM arm loaded the GDN weights UNSHARDED on every rank while the forward
ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine).
Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs
divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/
concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows,
conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous
head range. norm.weight [vd] is REPLICATED (shared across value heads).
The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer
forward. Byte-identical at tp_size==1 (gated else pass-through).
Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were
garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope:
BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not
yet sharded (follow-up). Stacks on #250 + #254.
SeedSource
pushed a commit
that referenced
this pull request
Jul 9, 2026
…acks on #254+#250] (#259) * feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints Load Hcompany Holo-3.1-0.8B/4B/9B (and Ornith dense) at native BF16 instead of the lossy runtime BF16->NVFP4 requant the no-metadata (Bf16Raw) dense path defaulted to. Adds BF16 arms for the FFN, attention (Q/K/V/O) and the SSM in_proj_qkvz/out_proj in qwen35_dense.rs, and gates the batched-decode qkvz path (ssm_batched.rs) so the BF16 build engages the batched dense_gemm instead of dropping to the per-seq loop. The NVFP4 weights stay installed as the spec-decode/batched fallback. Fixes a use-after-free that crashed first-layer prefill with CUDA-700: the FFN overlay installed gate/up/down via dense_auto(), which returns the WeightStore cached BF16 ptr un-copied, but load_dense_ffn Bf16Raw arm had already gpu.free() that buffer while building the NVFP4 fallback. Snapshot fresh D2D copies before the free. Validated on Holo-3.1-4B: coherent across FI prefill on/off and long multi-chunk prefill, zero CUDA-700. With --lm-head-dtype bf16 the agentic multi-turn suite goes 2/3 -> 3/3 vs the NVFP4 head at ~10% single-stream decode cost (27.8 -> 25.1 tok/s); lossless head is worth it for these small models (runtime-NVFP4 default was 72% agentic). * feat(gdn-tp): TP-shard the dense-model GDN/SSM heads (dense pure-TP) Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used by the MoE loader) into the dense loader SSM arm so dense Holo models (0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense SSM arm loaded the GDN weights UNSHARDED on every rank while the forward ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine). Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/ concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows, conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous head range. norm.weight [vd] is REPLICATED (shared across value heads). The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer forward. Byte-identical at tp_size==1 (gated else pass-through). Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope: BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not yet sharded (follow-up). Stacks on #250 + #254.
SeedSource
added a commit
that referenced
this pull request
Jul 9, 2026
rsafier
added a commit
that referenced
this pull request
Jul 10, 2026
Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used
by the MoE loader) into the dense loader SSM arm so dense Holo models
(0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense
SSM arm loaded the GDN weights UNSHARDED on every rank while the forward
ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine).
Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs
divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/
concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows,
conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous
head range. norm.weight [vd] is REPLICATED (shared across value heads).
The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer
forward. Byte-identical at tp_size==1 (gated else pass-through).
Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were
garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope:
BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not
yet sharded (follow-up). Stacks on #250 + #254.
31c1b3e to
e017018
Compare
…eckpoints Load Hcompany Holo-3.1-0.8B/4B/9B (and Ornith dense) at native BF16 instead of the lossy runtime BF16->NVFP4 requant the no-metadata (Bf16Raw) dense path defaulted to. Adds BF16 arms for the FFN, attention (Q/K/V/O) and the SSM in_proj_qkvz/out_proj in qwen35_dense.rs, and gates the batched-decode qkvz path (ssm_batched.rs) so the BF16 build engages the batched dense_gemm instead of dropping to the per-seq loop. The NVFP4 weights stay installed as the spec-decode/batched fallback. Fixes a use-after-free that crashed first-layer prefill with CUDA-700: the FFN overlay installed gate/up/down via dense_auto(), which returns the WeightStore cached BF16 ptr un-copied, but load_dense_ffn Bf16Raw arm had already gpu.free() that buffer while building the NVFP4 fallback. Snapshot fresh D2D copies before the free. Validated on Holo-3.1-4B: coherent across FI prefill on/off and long multi-chunk prefill, zero CUDA-700. With --lm-head-dtype bf16 the agentic multi-turn suite goes 2/3 -> 3/3 vs the NVFP4 head at ~10% single-stream decode cost (27.8 -> 25.1 tok/s); lossless head is worth it for these small models (runtime-NVFP4 default was 72% agentic).
e017018 to
6059f7c
Compare
rsafier
added a commit
that referenced
this pull request
Jul 13, 2026
Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used
by the MoE loader) into the dense loader SSM arm so dense Holo models
(0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense
SSM arm loaded the GDN weights UNSHARDED on every rank while the forward
ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine).
Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs
divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/
concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows,
conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous
head range. norm.weight [vd] is REPLICATED (shared across value heads).
The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer
forward. Byte-identical at tp_size==1 (gated else pass-through).
Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were
garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope:
BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not
yet sharded (follow-up). Stacks on #250 + #254.
rsafier
added a commit
that referenced
this pull request
Jul 13, 2026
Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used
by the MoE loader) into the dense loader SSM arm so dense Holo models
(0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense
SSM arm loaded the GDN weights UNSHARDED on every rank while the forward
ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine).
Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs
divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/
concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows,
conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous
head range. norm.weight [vd] is REPLICATED (shared across value heads).
The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer
forward. Byte-identical at tp_size==1 (gated else pass-through).
Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were
garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope:
BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not
yet sharded (follow-up). Stacks on #250 + #254.
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
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
Adds a native-BF16 load path for no-metadata dense Qwen3.5-family checkpoints (Hcompany Holo-3.1-0.8B/4B/9B, Ornith dense). These ship plain BF16
.weighttensors with no quantization metadata, so today the loader hardwires them through the runtime BF16→NVFP4 requant path (Nvfp4Variant::Bf16Raw) — an uncalibrated, lossy quantization that degrades small-model quality (the BF16-source models exist precisely because no calibrated NVFP4/FP8 release does).This keeps FFN, attention (Q/K/V/O), and the SSM
in_proj_qkvz/out_projin native BF16, dispatching thedense_gemv_bf16/dense_gemm_bf16kernels already in the standard bundle (same path Gemma-4 dense uses). Loader-only; no new kernels.qwen35_dense.rs: BF16 arms for FFN (snapshot overlay), attention (BF16 dense q/k/v +set_o_dense_bf16), SSM (BF16in_proj_qkvz+out_proj_dense), all gated onBf16Raw. NVFP4 weights are still built + kept as the spec-decode/batched fallback (a null NVFP4 weight under w4a16 dispatch is the CUDA-700-at-concurrency footgun).ssm_batched.rs: extend theqkvz_okgate so the pure-BF16 build engages the batcheddense_gemmdecode path instead of the per-seq loop.Use-after-free fix (would-be CUDA-700)
The naive FFN overlay installs gate/up/down via
dense_auto(), which returns theWeightStorecached BF16 pointer un-copied. Butload_dense_ffn's Bf16Raw arm runtime-quantizes to NVFP4 first andgpu.free()s that same buffer (quantized_any,nvfp4_detect.rs), and the store never evicts — so the overlay would handset_bf16_weightsfreed memory →dense_gemm_tcillegal access on the first prefill. Fixed by snapshotting fresh D2D copies beforeload_dense_ffnfrees the originals.Recommended serve flag
Pair with
--lm-head-dtype bf16to keep the 248K-vocab LM head lossless too (else it defaults to runtime-NVFP4).Validation
Built + served Holo-3.1-4B on GB10 from this branch (cherry-picked onto current
main, post-#229):dense weight loader: 32 layers,LM head kept as BF16).converttask is a flaky exact-value reasoning check; not a loader signal).Notes
ATLAS_TARGET_QUANT=bf16build (that strips the w4a16 stubs the dense loader still constructs) — use the standard bundle.--target-kv-tokens) lives in the separate Holo-GB10 stack; this PR is just the loader capability, whichmainalready routes via theqwen35_densedispatch.Related PRs — independent, all target
mainNow that #229 is merged, these three each cherry-pick cleanly onto current main with zero conflicts and are independent (no stacking) — review/merge in any order:
Follow-ups (not yet opened): Tier-2 FlashInfer-prefill + decode-concurrency reconciliation vs #229 FI (needs a merge, not a clean pick); and closing older stacked PRs folded into #229 (e.g. #220 FlashInfer-GDN AOT) after a rebase-diff confirms no outstanding delta.