`nemo_rl.models.generation.trtllm.quantization.fp8`#

Routed-expert block-FP8 / MXFP8 refit support for TRT-LLM (MoE models).

Module Contents#

Functions#

`build_quant_config`	Build the TRT-LLM `quantization_config` for block-FP8 or MXFP8.
`validate_fused_expert_layout`	Validate Qwen3.5’s fused routed-expert layout before the first refit.
`validate_model_type`	Fail setup for models outside the supported list.
`validate_routed_experts`	Fail setup unless the FP8 filter will quantize something, correctly.
`configure_fp8_llm_kwargs`	Apply the block-FP8 / MXFP8 contract to TRT-LLM args.
`configure_fp8_moe_backend`	Force (or check) the MoE backend for the requested scale format.
`_has_fp8_block_scales`
`is_fp8_model`	Return whether a TRT-LLM model uses 128x128 block FP8.
`is_mxfp8_model`	Return whether a TRT-LLM model uses MXFP8 (E4M3 + UE8M0 1x32).
`is_quantized_expert_refit`	Whether refit must quantize routed experts before loading.
`cast_tensor_to_fp8_blockwise`	Quantize the last two dims into E4M3 with FP32 block scales.
`cast_tensor_to_mxfp8_blockwise`	Quantize the last dim into E4M3 with UE8M0 1x{block} scales.
`_insert_unique`
`_insert_quantized_projection`
`clone_mapper_staging_weights`	Clone tensors the Qwen3.5 mapper may retain across reload calls.
`_convert_fused_expert_weight`
`load_weights`	Convert BF16 weights to block-FP8 or MXFP8 for the TRT-LLM engine.

Data#

`FP8_BLOCK_SIZE`
`FP8_EXPERT_CHUNK_SIZE`
`MXFP8_BLOCK_SIZE`
`UE8M0_BIAS`
`E4M3_MAX`
`ROUTED_EXPERTS_ONLY_IGNORE_PATTERNS`
`FP8_BLOCK_QUANT_KWARGS`
`MXFP8_BLOCK_QUANT_KWARGS`
`_EXPERT_PREFIX`
`_FUSED_EXPERT_RE`
`_SPLIT_EXPERT_RE`
`_DENSE_MLP_RE`
`_SPLIT_LINEAR_ATTN_RE`
`SUPPORTED_MODEL_TYPES`

API#

nemo_rl.models.generation.trtllm.quantization.fp8.FP8_BLOCK_SIZE#: (128, 128)

nemo_rl.models.generation.trtllm.quantization.fp8.FP8_EXPERT_CHUNK_SIZE#: 16

nemo_rl.models.generation.trtllm.quantization.fp8.MXFP8_BLOCK_SIZE#: 32

nemo_rl.models.generation.trtllm.quantization.fp8.UE8M0_BIAS#: 127

nemo_rl.models.generation.trtllm.quantization.fp8.E4M3_MAX#: 448.0

nemo_rl.models.generation.trtllm.quantization.fp8.ROUTED_EXPERTS_ONLY_IGNORE_PATTERNS#: (‘model.layers..self_attn’, ‘model.layers..linear_attn’, ‘model.layers.*.mlp.gate’, ‘model.layer…

nemo_rl.models.generation.trtllm.quantization.fp8.build_quant_config(is_mx: bool) → dict[str, Any]#: Build the TRT-LLM quantization_config for block-FP8 or MXFP8.

nemo_rl.models.generation.trtllm.quantization.fp8.FP8_BLOCK_QUANT_KWARGS: dict[str, Any]#: ‘build_quant_config(…)’

nemo_rl.models.generation.trtllm.quantization.fp8.MXFP8_BLOCK_QUANT_KWARGS: dict[str, Any]#: ‘build_quant_config(…)’

nemo_rl.models.generation.trtllm.quantization.fp8._EXPERT_PREFIX: str#: ‘(?:(?:(?:model.)?(?:language_model.)?)layers.\d+|mtp.layers.\d+).mlp.experts’

nemo_rl.models.generation.trtllm.quantization.fp8._FUSED_EXPERT_RE#: ‘compile(…)’

nemo_rl.models.generation.trtllm.quantization.fp8._SPLIT_EXPERT_RE#: ‘compile(…)’

nemo_rl.models.generation.trtllm.quantization.fp8._DENSE_MLP_RE#: ‘compile(…)’

nemo_rl.models.generation.trtllm.quantization.fp8._SPLIT_LINEAR_ATTN_RE#: ‘compile(…)’

nemo_rl.models.generation.trtllm.quantization.fp8.validate_fused_expert_layout( state_dict_info: dict[str, Any], ) → None#: Validate Qwen3.5’s fused routed-expert layout before the first refit.

nemo_rl.models.generation.trtllm.quantization.fp8.SUPPORTED_MODEL_TYPES#: (‘qwen3_5_moe’, ‘qwen3_moe’)

nemo_rl.models.generation.trtllm.quantization.fp8.validate_model_type(model_type: str) → None#: Fail setup for models outside the supported list.

nemo_rl.models.generation.trtllm.quantization.fp8.validate_routed_experts( state_dict_info: dict[str, Any], ) → None#

Fail setup unless the FP8 filter will quantize something, correctly.

Routed experts are either fused (mlp.experts.{gate_up,down}_proj, Qwen3.5) or per-expert (mlp.experts.{i}.{gate,up,down}_proj.weight, Qwen3 MoE). Only experts are quantized, so at least one must exist and dense MLP weights are rejected.

nemo_rl.models.generation.trtllm.quantization.fp8.configure_fp8_llm_kwargs( llm_kwargs: dict[str, Any], *, is_mx: bool = False, ) → None#

Apply the block-FP8 / MXFP8 contract to TRT-LLM args.

Conflicting quantization or load-format overrides raise at setup.

nemo_rl.models.generation.trtllm.quantization.fp8.configure_fp8_moe_backend( llm_kwargs: dict[str, Any], moe_config_type: type[Any], *, is_mx: bool = False, ) → None#

Force (or check) the MoE backend for the requested scale format.

Block-FP8 needs TRTLLM (keeps FP32 scales). MXFP8 uses CUTLASS (default) or CUTEDSL (Rubin).

nemo_rl.models.generation.trtllm.quantization.fp8._has_fp8_block_scales(quant_config: Any) → bool#

nemo_rl.models.generation.trtllm.quantization.fp8.is_fp8_model(quant_config: Any) → bool#: Return whether a TRT-LLM model uses 128x128 block FP8.

nemo_rl.models.generation.trtllm.quantization.fp8.is_mxfp8_model(quant_config: Any) → bool#: Return whether a TRT-LLM model uses MXFP8 (E4M3 + UE8M0 1x32).

nemo_rl.models.generation.trtllm.quantization.fp8.is_quantized_expert_refit(quant_config: Any) → bool#: Whether refit must quantize routed experts before loading.

nemo_rl.models.generation.trtllm.quantization.fp8.cast_tensor_to_fp8_blockwise( data_hp: torch.Tensor, weight_block_size: collections.abc.Sequence[int] = FP8_BLOCK_SIZE, ) → tuple[torch.Tensor, torch.Tensor]#

Quantize the last two dims into E4M3 with FP32 block scales.

Leading dims are independent. weight_scale_inv is [..., out_block, in_block]; dequantize as fp8.float() * weight_scale_inv.

nemo_rl.models.generation.trtllm.quantization.fp8.cast_tensor_to_mxfp8_blockwise( data_hp: torch.Tensor, block_size: int = MXFP8_BLOCK_SIZE, ) → tuple[torch.Tensor, torch.Tensor]#

Quantize the last dim into E4M3 with UE8M0 1x{block} scales.

Mirrors TRT-LLM’s quant_bf16_to_mxfp8: the power-of-two exponent maps each block’s amax just under the E4M3 max. Blocks span K only.

Returns (e4m3 [..., K], ue8m0 uint8 [..., K // block]).

nemo_rl.models.generation.trtllm.quantization.fp8._insert_unique( output: dict[str, torch.Tensor], name: str, tensor: torch.Tensor, ) → None#

nemo_rl.models.generation.trtllm.quantization.fp8._insert_quantized_projection( output: dict[str, torch.Tensor], name: str, tensor: torch.Tensor, *, is_mx: bool = False, ) → None#

nemo_rl.models.generation.trtllm.quantization.fp8.clone_mapper_staging_weights( weights: dict[str, torch.Tensor], ) → dict[str, torch.Tensor]#

Clone tensors the Qwen3.5 mapper may retain across reload calls.

IPC buffers are reused after each ACK, so retained views must own storage.

nemo_rl.models.generation.trtllm.quantization.fp8._convert_fused_expert_weight( output: dict[str, torch.Tensor], *, name: str, tensor: torch.Tensor, prefix: str, projection: str, is_mx: bool = False, ) → None#

nemo_rl.models.generation.trtllm.quantization.fp8.load_weights( weight_list: collections.abc.Iterable[tuple[str, torch.Tensor]], *, is_mx: bool = False, ) → dict[str, torch.Tensor]#

Convert BF16 weights to block-FP8 or MXFP8 for the TRT-LLM engine.

Only routed experts are quantized. Fused expert tensors (Qwen3.5) are expanded to per-expert HF names; other weights pass through unchanged.

nemo_rl.models.generation.trtllm.quantization.fp8#

Module Contents#

Functions#

Data#

API#

`nemo_rl.models.generation.trtllm.quantization.fp8`#