Module model_a 

2026-09-25: Basic store accessors: raw pointers, scalars, KV scales, ModelOpt NVFP4, native MXFP4, and dtype-converting dense loads.

Owner: model-layers (weight loading). Invariants: none beyond the types.

Structs

ModelWeights

2026-09-25: Whole-model weights built by ModelWeights::from_store.

Functions

dense

dense_auto_fp8_or_bf16

2026-09-25: Load {prefix}.weight as BF16: a BF16 tensor is the store pointer, FP8E4M3 goes through dequant_fp8_blockscaled_to_bf16, and any other dtype is an error.

dense_f32_as_bf16

2026-09-25: Load an FP32 tensor into a new BF16 device buffer, converted on the host with f32_to_bf16; any other dtype is an error.

dense_f32_safe

2026-09-25: Load name as BF16: an FP32 tensor is truncated into a new buffer, any other dtype is returned as the store pointer unchanged.

dense_keep_f32

2026-09-25: Load name as FP32: an FP32 tensor is the store pointer, BF16 is widened on the host into a new buffer, and any other dtype is an error. load_ssm uses it for A_log and dt_bias.

load_checkpoint_kv_scales

2026-09-28: The checkpoint’s FP8 KV-cache scales (k_scale, v_scale) for the attention layer at attn_prefix, or None when it ships none. The keys come from WeightStore::kv_scale_keys, the resolver serve’s census also uses.

load_kv_scales

2026-09-28: load_checkpoint_kv_scales, with (1.0, 1.0) for a layer the checkpoint ships no scales for. Serve logs how many layers that is (resolve_kv_cache_config), from the same resolver.

quantized

2026-09-25: Build a QuantizedWeight from ModelOpt NVFP4 keys: weight, weight_scale, the FP32 scalar weight_scale_2 (read to the host), and an optional input_scale.

quantized_mxfp4_e8m0

2026-09-25: quantized_mxfp4_e8m0_pair on {prefix}.weight and {prefix}.scale.

quantized_mxfp4_e8m0_pair

2026-09-25: Hand out a native MXFP4 weight’s store pointers unchanged: E2M1 nibbles [n, k/2] in weight_key and one E8M0 scale per 32 weights in scale_key. Any other group size inferred from the element counts is an error. weight_scale_2 is 1.0 because the format has no per-tensor scale.

scalar_bf16_or_f32

2026-09-28: Read a one-element BF16 or FP32 tensor to the host as f32; any other dtype or element count is an error. Checkpoints ship FP8 KV scales in both widths (unsloth: BF16 [1]; ModelOpt: FP32 []), and an F16 export is already BF16 in the store.