nemo_rl.modelopt.models.generation.vllm_modelopt#

Narrow vLLM extensions for ModelOpt NVFP4 rollout checkpoints.

vLLM owns checkpoint-layout restoration, layerwise post-load processing, CUDA-graph-stable tensor placement, and KV-cache scale reload. This module only supplies the vLLM 0.25 gaps needed here: ModelOpt W4A16 NVFP4 methods, per-projection ModelOpt MoE input-scale loading, materialization of FlashInfer’s global-scale views, and retention of method-owned MoE kernel references across layerwise reload. (Rank-local Marlin padding is gone: 0.25’s prepare_nvfp4_moe_layer_for_marlin pads natively.)

Module Contents#

Functions#

quantization_method_for_mode

Return the registered vLLM quantization method for a rollout mode.

_load_modelopt_moe_input_scale

Load a ModelOpt input scale without losing the gate/up shard.

_validated_w4a16_config

_canonicalize_nvfp4_scale_

Remove the E4M3 sign bit before Marlin’s unsigned scale conversion.

register_nemo_modelopt_nvfp4

Register NeMo’s two ModelOpt NVFP4 configs through vLLM’s public API.

Data#

NEMO_MODELOPT_W4A4

NEMO_MODELOPT_W4A16

_W4A4_ALGO

_W4A16_ALGO

_registered

API#

nemo_rl.modelopt.models.generation.vllm_modelopt.NEMO_MODELOPT_W4A4#

‘nemo_modelopt_nvfp4’

nemo_rl.modelopt.models.generation.vllm_modelopt.NEMO_MODELOPT_W4A16#

‘nemo_modelopt_w4a16_nvfp4’

nemo_rl.modelopt.models.generation.vllm_modelopt._W4A4_ALGO#

‘NVFP4’

nemo_rl.modelopt.models.generation.vllm_modelopt._W4A16_ALGO#

‘W4A16_NVFP4’

nemo_rl.modelopt.models.generation.vllm_modelopt._registered#

False

nemo_rl.modelopt.models.generation.vllm_modelopt.quantization_method_for_mode(mode: str) → str#

Return the registered vLLM quantization method for a rollout mode.

nemo_rl.modelopt.models.generation.vllm_modelopt._load_modelopt_moe_input_scale(
moe_layer: Any,
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: int,
return_success: bool = False,
) → bool | None#

Load a ModelOpt input scale without losing the gate/up shard.

Replaces the input_scale branch of vLLM’s MoE weight loader, which drops the gate/up (w1/w3) shard index. vLLM 0.25 handles this natively at https://github.com/vllm-project/vllm/blob/v0.25.1/vllm/model_executor/layers/fused_moe/routed_experts.py#L699-L712 but this override is still required: it uses shard_index = 0 if w1 else min(1, param.shape[-1] - 1) where upstream hardcodes 1, so it also handles the single-scale (non-gated) layout, and it copies with reshape_as rather than _to_scalar(). Delete once upstream covers the single-column case too.

nemo_rl.modelopt.models.generation.vllm_modelopt._validated_w4a16_config(
config: dict[str, Any],
) → dict[str, Any]#
nemo_rl.modelopt.models.generation.vllm_modelopt._canonicalize_nvfp4_scale_(scale: torch.Tensor) → None#

Remove the E4M3 sign bit before Marlin’s unsigned scale conversion.

nemo_rl.modelopt.models.generation.vllm_modelopt.register_nemo_modelopt_nvfp4() → None#

Register NeMo’s two ModelOpt NVFP4 configs through vLLM’s public API.