Skip to content

feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints#250

Open
rsafier wants to merge 1 commit into
mainfrom
feat/holo-native-bf16-dense-loader
Open

feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints#250
rsafier wants to merge 1 commit into
mainfrom
feat/holo-native-bf16-dense-loader

Conversation

@rsafier

@rsafier rsafier commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator

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 .weight tensors 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_proj in native BF16, dispatching the dense_gemv_bf16/dense_gemm_bf16 kernels 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 (BF16 in_proj_qkvz + out_proj_dense), all gated on Bf16Raw. 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 the qkvz_ok gate so the pure-BF16 build engages the batched dense_gemm decode 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 the WeightStore cached BF16 pointer un-copied. But load_dense_ffn's Bf16Raw arm runtime-quantizes to NVFP4 first and gpu.free()s that same buffer (quantized_any, nvfp4_detect.rs), and the store never evicts — so the overlay would hand set_bf16_weights freed memory → dense_gemm_tc illegal access on the first prefill. Fixed by snapshotting fresh D2D copies before load_dense_ffn frees the originals.

Recommended serve flag

Pair with --lm-head-dtype bf16 to 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):

Notes

  • Does not need an ATLAS_TARGET_QUANT=bf16 build (that strips the w4a16 stubs the dense loader still constructs) — use the standard bundle.
  • Full Holo serving config (e.g. --target-kv-tokens) lives in the separate Holo-GB10 stack; this PR is just the loader capability, which main already routes via the qwen35_dense dispatch.

Related PRs — independent, all target main

Now 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.

@rsafier rsafier force-pushed the feat/holo-native-bf16-dense-loader branch from 08b1e63 to 31c1b3e Compare July 4, 2026 22:01
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.
@rsafier rsafier force-pushed the feat/holo-native-bf16-dense-loader branch from 31c1b3e to e017018 Compare July 12, 2026 23:28
…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).
@rsafier rsafier force-pushed the feat/holo-native-bf16-dense-loader branch from e017018 to 6059f7c Compare July 13, 2026 00:43
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant