Module fp8_residency

Module fp8_residency 

Source
Expand description

The native-FP8 dense loader’s derived-copy decision table and the shape arithmetic that prices it (#915). Which derived weight copies the native-FP8 dense route actually needs, and a tally of the ones it still builds.

WHY (#915 root cause; O9 of the 2026-09-05 rental; refs #916, #917, #736). Measured on 1xH100, 2026-09-11, Qwen/Qwen3.8-27B-FP8 at 5f78270dc, native-FP8 profile (ATLAS_DENSE_FP8=1, --lm-head-dtype bf16):

  • the checkpoint itself is 28.75 GBWeightStore after prune: 1606 tensors, 28.747 GiB still resident, one ledger site (fast_weights/mod.rs:434, 29,436.7 MB x1606);
  • the ledger nevertheless reported 58.1 GB live before the KV cache was sized, and ATLAS_MEM_PROFILE recorded GPU-free falling 437 MB per layer over all 64 layers (28.0 GB) after the checkpoint was resident;
  • the teardown sweep reclaimed 28.01 GB across 1,980 allocations that no ModelResource owned.

The six sites the ledger named, and what each holds per layer:

sitebytes x counttensor
weight_map/loaders_fp8.rs:2298,960 MB x256quantize_to_nvfp4 packed [N,K/2]
weight_map/quantized.rs:2618,960 MB x256the transposed twin of the same
weight_loader/qwen35_dense.rs:1353,840 MB x48SSM fused [QKV|Z] FP8 concat
weight_map/quantized.rs:6431,600 MB x64attention Fp8WeightTransposed::weight_t
weight_map/loaders_fp8.rs:2301,120 MB x256the NVFP4 per-16 group scales
weight_map/quantized.rs:2621,120 MB x256the transposed twin of those scales

The 256 counts are 64 layers x 3 dense-FFN projections (gate/up/down, 42.5 MiB packed each) plus 16 attention layers x 4 (q/k/v/o, 25/6.25/ 6.25/12.5 MiB) — 8,160 + 800 MB, which is exactly the 8,960 reported.

None of the NVFP4 copies is reachable once the native FP8 overlay is installed. DenseFfnLayer::forward (dense_ffn.rs:1044) and forward_prefill_inner (dense_ffn.rs:1985) both return from inside their if let Some(ref fp8w) = self.fp8_weights arm, forward_k2/forward_k3/ forward_km redirect to forward_prefill via native_small_batch_uses_prefill, and the w8_gemm! macro binds gate_t/up_t/down_t to a literal None, so even the W8A16 fallback rungs read the original [N,K] E4M3 bytes. The NVFP4 gate/up/down and their transposed twins are pure load-time waste: 18.4 GiB of the 28.

The loader runs before dispatch exists, so the route has to be derived from the same resolvers the forward pass uses rather than from a ForwardContext:

  • crate::layers::ops::GemmDispatch::from_env — the exact constructor model/impl_a1.rs:836 calls to build the context. It is a pure function of the environment and a serve never rewrites its own environment, so the loader’s answer and the context’s answer cannot disagree.

ATLAS_FFN_W8A16_ONLY is deliberately NOT an input: it steers the dense FFN from the W8A8 arm onto rung 5 of the same w8_gemm! match, whose transposed operand is that literal None — so it selects between two kernels that both read the original FP8 bytes, and cannot resurrect an NVFP4 reader.

That is the contract: any new lever that can route a native-FP8 layer back onto an NVFP4 kernel must be added to DenseFp8Plan::resolve as well as to the dispatch site, or the loader will have freed the weight that site wants. ATLAS_DENSE_FP8_KEEP_NVFP4 is the escape hatch that restores the pre-#915 behaviour wholesale while such a gap is diagnosed.

Structs§

DenseFp8Inputs
The inputs DenseFp8Plan::resolve decides from. Taken as a struct so the CPU decision-table test can pin every clause without touching the process environment (the OnceLock resolvers cannot be toggled per test).
DenseFp8Plan
What a native-FP8 dense layer must materialise beyond the checkpoint bytes.
DerivedResidency
Running tally of the derived (non-checkpoint) device bytes this loader allocated, and of the ones it decided not to build.
RouteEnv
The environment-resolved half of DenseFp8Inputs, read ONCE per load.
TwinsBuilt
Which derived twin families the plan built. Reported by name so the serve log says which copies are resident rather than only how many bytes.

Functions§

attn_fp8_twin_bytes
The FP8 prefill twins want selects, for one full-attention layer.
attn_nvfp4_bytes
What the pre-#915 loader spent per full-attention layer on NVFP4: q/k/v/o, each with a transposed twin, plus the fused [q|k|v] transposed twin (transpose_concat_for_gemm).
dense_ffn_nvfp4_bytes
What the pre-#915 loader spent per dense-FFN layer on NVFP4: gate, up and down, each with a transposed twin of identical size.
ffn_gateup_fused_bytes
What ONE fused dense-FFN gate+up weight costs: the two [inter, hidden] E4M3 blocks appended along N, plus their two [inter/128, hidden/128] FP32 block-scale grids appended the same way (#927).
ffn_gateup_fused_parts
ffn_gateup_fused_bytes split into (weight bytes, scale-grid bytes) — the loader adopts the two buffers separately, so it needs the terms rather than the sum, and taking them from here is what keeps the prediction and the tally one arithmetic.
fp8_twin_bytes
Bytes one FP8 [K, N] transposed twin costs: the E4M3 bytes plus the transposed [K/128, N/128] FP32 block-scale grid. Mirrors Fp8Weight::transpose_for_gemm (weight_map/quantized.rs:643/:660).
keep_nvfp4_fallback
The per-layer kill switch that restores the pre-#915 behaviour: build every NVFP4 fallback copy even where the dispatch cannot reach it.
nvfp4_bytes
Bytes one NVFP4 QuantizedWeight costs: packed [N, K/2] E2M1 nibbles plus the [N, K/16] per-group scale byte. Mirrors quantize_to_nvfp4 (weight_map/loaders_fp8.rs:229-230) and transpose_for_gemm_gs (weight_map/quantized.rs:261-262), which allocate the same two sizes.