shard_cast_utils#

Shared primitives for streaming checkpoint casts and export.

Functions

`build_w13_amax_overrides`	Return one shared weight amax for each fused w1/w3 pair.
`build_w13_kmax_overrides`	Return one shared E8M0 maximum exponent for each fused w1/w3 pair.
`dequantize_mxfp4_to_bf16`	Dequantize packed MXFP4 weights and E8M0 scales to BF16.
`mxfp4_kmax`	Return the largest non-zero unbiased exponent in an E8M0 scale tensor.
`quantize_mxfp4_to_nvfp4`	Dequantize MXFP4 and requantize it to NVFP4 using an optional global amax.
`quantize_mxfp4_to_nvfp4_lossless`	Closed-form MXFP4-to-NVFP4 cast with lossless-block accounting.

build_w13_amax_overrides(expert_bases, get_amax)#

Return one shared weight amax for each fused w1/w3 pair.

Parameters:

Return type:

dict[str, torch.Tensor]

build_w13_kmax_overrides(expert_bases, get_scale, device)#

Return one shared E8M0 maximum exponent for each fused w1/w3 pair.

Parameters:

Return type:

dict[str, int]

dequantize_mxfp4_to_bf16(mxfp4_weight, mxfp4_scale, device)#

Dequantize packed MXFP4 weights and E8M0 scales to BF16.

Parameters:

Return type:

Tensor

mxfp4_kmax(mxfp4_scale, device='cpu')#

Return the largest non-zero unbiased exponent in an E8M0 scale tensor.

Parameters:

Return type:

int

quantize_mxfp4_to_nvfp4(mxfp4_weight, mxfp4_scale, weight_amax, device)#

Dequantize MXFP4 and requantize it to NVFP4 using an optional global amax.

Parameters:

Return type:

tuple[Tensor, Tensor, Tensor, bool]

quantize_mxfp4_to_nvfp4_lossless(mxfp4_weight, mxfp4_scale, k_max, device)#

Closed-form MXFP4-to-NVFP4 cast with lossless-block accounting.

Parameters:

Return type:

tuple[Tensor, Tensor, Tensor, int, int]