NVIDIA-NeMo/Speech · #16341
SpeechLM vLLM: serve FP8 and NVFP4 ModelOpt checkpoints (quantized decoder + prequantized FP8 audio encoder)
nemo/collections/asr/modules/transformer_encoder.py3 + / 1 −
@@ -590,7 +590,9 @@ def _forward_sequence_packed( def _project_sequence_packed_qkv(self, x, *, position_ids, fused_qkv): """Project packed states to Q/K/V and apply any Q/K preparation.""" total_tokens = x.shape[0]- if fused_qkv:+ # Slicing .weight for F.linear is valid only for a plain nn.Linear; any other projection+ # (quantized, adapted) applies its own forward.+ if fused_qkv or type(self.w_qkv) is not nn.Linear: qkv = self.w_qkv(x).view(total_tokens, 3, self.n_heads, self.head_dim) q, k, v = (projection.contiguous() for projection in qkv.unbind(dim=1)) else:nemo/collections/speechlm2/vllm/salm/backends.py6 + / 0 −
@@ -311,6 +311,12 @@ def nemo_to_hf_llm_weights( ) elif hf_name in ( "backbone.embed_tokens.weight",+ # NemotronH's embedding is `embeddings`, and a ModelOpt export+ # already carries that HF name, so it never hits the+ # `embed_tokens` rename above. Omitting it here leaves the+ # embedding unpadded and every quantized checkpoint dies on+ # vLLM's `loaded_weight.shape[output_dim] == org_vocab_size`.+ "backbone.embeddings.weight", "lm_head.weight", ): if target_vocab:nemo/collections/speechlm2/vllm/salm/fp8_encoder.pyadded270 + / 0 −
@@ -0,0 +1,270 @@+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+# http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++"""FP8 audio encoder for SpeechLM checkpoints quantized with ModelOpt FP8.++A checkpoint describes its audio encoder's quantization the way it describes its+decoder's, in the ModelOpt ``quantization_config`` of ``config.json``+(``quant_algo: FP8``): every Linear of the perception module that no ``ignore``+entry matches is quantized. Such a layer stores an FP8 E4M3 ``weight``, a float32+scalar ``weight_scale`` and a float32 scalar ``input_scale``, the static per-tensor+activation scale, both in the divisor convention ``x_fp8 = x / scale``. A+speaker-aware checkpoint keeps its diarizer, the ASR encoder's input projection and+the connector unquantized by listing them::++ "ignore": [..., "perception.encoder.diarization_model*",+ "perception.encoder.asr_encoder.pre_encode*", "perception.proj"]++The audio encoder runs as plain PyTorch inside the plugin rather than through vLLM's+quantized layers, so ``build_fp8_encoder`` applies vLLM's exclusion rule to the+perception Linears itself and builds ``FP8Linear`` layers for the quantized ones+when the model is constructed; the checkpoint then loads into them directly.++A ``quantization_config`` without any ``perception`` entry predates encoder+quantization (its decoder was quantized on its own), so its encoder stays unquantized.++A ModelOpt ``MIXED_PRECISION`` checkpoint (for example an NVFP4 decoder with FP8 layers)+has no such rule: it names every quantized layer in ``quantized_layers``. There an+encoder Linear is FP8 when listed with ``quant_algo: FP8`` and unquantized when not+listed::++ "quantized_layers": {..., "perception.encoder.asr_encoder.layers.0.attn.w_qkv": {"quant_algo": "FP8"}}+"""++from collections.abc import Callable, Mapping+from typing import Any, Optional++import torch+from torch import nn++from nemo.utils import logging++FP8_DTYPE = torch.float8_e4m3fn+FP8_MAX = 448.0+_PERCEPTION_PREFIX = "perception."+++class FP8Linear(nn.Linear):+ """Linear layer with FP8 E4M3 weights and static per-tensor scales, run with vLLM's CUTLASS FP8 GEMM.++ The module is built empty and filled by ``load_state_dict``: ``weight``,+ ``weight_scale`` and ``input_scale`` are persistent buffers named and shaped like+ ModelOpt's FP8 checkpoint tensors. Dtype casts of an enclosing module convert only+ the bias; the FP8 weight and the float32 scales keep their dtypes and follow+ device moves. It subclasses ``nn.Linear`` because encoders choose how to call an+ input projection with ``isinstance(module, nn.Linear)``.++ Args:+ in_features: Input width. CUTLASS needs a multiple of 16.+ out_features: Output width.+ bias: Whether the layer has a bias.+ bias_dtype: Dtype of the bias parameter.+ device: Device for the empty buffers.+ """++ def __init__(+ self,+ in_features: int,+ out_features: int,+ *,+ bias: bool,+ bias_dtype: torch.dtype = torch.bfloat16,+ device: Optional[torch.device] = None,+ ) -> None:+ # nn.Linear.__init__ would allocate full-precision parameters that the FP8 buffers replace.+ nn.Module.__init__(self)+ if in_features % 16 != 0:+ raise ValueError(f"in_features={in_features} must be a multiple of 16 for the CUTLASS FP8 GEMM")+ self.in_features = in_features+ self.out_features = out_features+ self.register_buffer("weight", torch.empty(out_features, in_features, dtype=FP8_DTYPE, device=device))+ self.register_buffer("weight_scale", torch.empty((), dtype=torch.float32, device=device))+ self.register_buffer("input_scale", torch.empty((), dtype=torch.float32, device=device))+ self.bias = None+ if bias:+ self.bias = nn.Parameter(torch.empty(out_features, dtype=bias_dtype, device=device), requires_grad=False)++ def _apply(self, fn: Callable[[torch.Tensor], torch.Tensor], recurse: bool = True) -> "FP8Linear":+ """Apply ``fn`` to the bias as usual but only its device change to the buffers.++ The perception module is cast to bfloat16 as a whole; applied to the buffers,+ that cast would turn the FP8 weight into bfloat16 values and round the scales.+ """+ buffers = dict(self._buffers)+ self._buffers.clear()+ try:+ super()._apply(fn, recurse)+ finally:+ for name, buf in buffers.items():+ if buf is not None:+ device = fn(torch.empty(0, dtype=buf.dtype, device=buf.device)).device+ buf = buf.to(device=device)+ self._buffers[name] = buf+ return self++ def forward(self, x: torch.Tensor) -> torch.Tensor:+ """Quantize ``x`` to FP8 with the static input scale and multiply by the FP8 weight, returning ``x.dtype``."""+ import vllm._custom_ops as ops++ out_shape = (*x.shape[:-1], self.out_features)+ # The quantization kernel needs a contiguous last dimension; a transposed input reshapes to a strided view.+ xq, x_scale = ops.scaled_fp8_quant(x.reshape(-1, self.in_features).contiguous(), self.input_scale)+ # CUTLASS wants a column-major B, which the transpose gives as a view.+ out = ops.cutlass_scaled_mm(xq, self.weight.t(), x_scale, self.weight_scale, x.dtype, self.bias)+ return out.reshape(out_shape)++ def extra_repr(self) -> str:+ """Summarize the layer shape and quantization."""+ return f"in_features={self.in_features}, out_features={self.out_features}, fp8_e4m3, static per-tensor scales"+++def build_fp8_encoder(perception: nn.Module, quant_config: Optional[Any]) -> int:+ """Build empty ``FP8Linear`` layers in place of the perception Linears that ``quant_config`` quantizes.++ Call this when the model is constructed, before any weights load: the+ checkpoint's FP8 tensors then load into these layers directly, and later dtype+ casts of the perception module leave them intact.++ Args:+ perception: The SpeechLM perception module, whose names start with ``perception.`` in the checkpoint.+ quant_config: vLLM's quantization config for the checkpoint, or None.++ Returns:+ The number of Linears replaced; 0 when ``quant_config`` does not quantize the encoder.++ Raises:+ NotImplementedError: If the encoder is quantized with anything but ModelOpt FP8.+ ValueError: If a quantized Linear's input width does not suit the CUTLASS FP8 GEMM.+ """+ if hasattr(quant_config, "get_name") and quant_config.get_name() == "modelopt_mixed":+ targets = _listed_fp8_targets(perception, quant_config)+ else:+ targets = _unexcluded_fp8_targets(perception, quant_config)+ if not targets:+ return 0++ replacements = [+ (+ name,+ FP8Linear(+ module.in_features,+ module.out_features,+ bias=module.bias is not None,+ bias_dtype=module.bias.dtype if module.bias is not None else torch.bfloat16,+ device=module.weight.device,+ ),+ )+ for name, module in targets+ ]+ for name, layer in replacements:+ parent_name, _, child_name = name.rpartition(".")+ parent = perception.get_submodule(parent_name) if parent_name else perception+ setattr(parent, child_name, layer)++ logging.info(f"FP8 audio encoder: {len(replacements)} Linears built from quantization_config")+ return len(replacements)+++def _unexcluded_fp8_targets(perception: nn.Module, quant_config: Optional[Any]) -> list[tuple[str, nn.Linear]]:+ """Perception Linears that a ModelOpt FP8 config quantizes: those no exclusion entry matches."""+ exclude_modules = getattr(quant_config, "exclude_modules", None) or ()+ if not hasattr(quant_config, "is_layer_excluded") or not any(+ str(entry).startswith("perception") for entry in exclude_modules+ ):+ return []+ targets = [+ (name, module)+ for name, module in perception.named_modules()+ if isinstance(module, nn.Linear)+ and not isinstance(module, FP8Linear)+ and not quant_config.is_layer_excluded(_PERCEPTION_PREFIX + name)+ ]+ if not targets:+ return []+ scheme = (quant_config.get_name(), getattr(quant_config, "quant_method", None))+ if scheme != ("modelopt", "FP8"):+ raise NotImplementedError(+ f"quantization_config quantizes {len(targets)} audio encoder Linear(s), e.g. "+ f"{_PERCEPTION_PREFIX + targets[0][0]!r}, as {scheme[0]} {scheme[1]}; the audio encoder supports only "+ "ModelOpt FP8 with static per-tensor scales"+ )+ return targets+++def _listed_fp8_targets(perception: nn.Module, quant_config: Any) -> list[tuple[str, nn.Linear]]:+ """Perception Linears that a ModelOpt MIXED_PRECISION config lists in ``quantized_layers``."""+ quantized_layers = quant_config.quantized_layers+ targets = []+ others = {}+ for name, module in perception.named_modules():+ prefix = _PERCEPTION_PREFIX + name+ if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear) or prefix not in quantized_layers:+ continue+ # vLLM's mixed-precision config checks exclusions before the per-layer table.+ if quant_config.is_layer_excluded(prefix):+ continue+ algo = str(quantized_layers[prefix].get("quant_algo", "")).upper()+ if algo == "FP8":+ targets.append((name, module))+ else:+ others[prefix] = algo+ if others:+ raise NotImplementedError(+ f"quantized_layers lists {len(others)} audio encoder Linear(s) as {sorted(set(others.values()))}, "+ f"e.g. {next(iter(others))!r}; the audio encoder supports only FP8 with static per-tensor scales"+ )+ return targets+++def check_fp8_encoder_weights(perception: nn.Module, weights: Mapping[str, torch.Tensor]) -> None:+ """Fail if the checkpoint's FP8 encoder tensors do not match the layers ``build_fp8_encoder`` built.++ ``load_state_dict`` converts dtypes silently, so FP8 data copied into a bfloat16+ Linear, or bfloat16 data into an ``FP8Linear``, would load without an error.++ Args:+ perception: The perception module after ``build_fp8_encoder``.+ weights: The checkpoint's perception tensors, named relative to the perception module.++ Raises:+ ValueError: If an ``FP8Linear`` lacks an FP8 weight or a per-tensor scale, or the+ checkpoint stores an FP8 weight for a layer that is not quantized.+ """+ fp8_layers = {name for name, module in perception.named_modules() if isinstance(module, FP8Linear)}+ problems = []+ for name in sorted(fp8_layers):+ weight = weights.get(f"{name}.weight")+ if weight is None or weight.dtype != FP8_DTYPE:+ problems.append(f"{name}.weight is {'missing' if weight is None else weight.dtype}, not {FP8_DTYPE}")+ for scale_name in ("weight_scale", "input_scale"):+ scale = weights.get(f"{name}.{scale_name}")+ if scale is None or scale.numel() != 1:+ problems.append(+ f"{name}.{scale_name} is {'missing' if scale is None else tuple(scale.shape)}, not one scale"+ )+ unquantized = sorted(+ key+ for key, tensor in weights.items()+ if tensor.dtype == FP8_DTYPE and key.rpartition(".")[0] not in fp8_layers+ )+ problems.extend(f"{key} is FP8 but quantization_config leaves its layer unquantized" for key in unquantized)+ if problems:+ shown = "; ".join(problems[:5]) + (f"; and {len(problems) - 5} more" if len(problems) > 5 else "")+ hint = (+ " (config.json's quantization_config must describe the encoder's FP8 layers; an "+ "encoder_quantization block is not read)"+ if unquantized+ else ""+ )+ raise ValueError(f"FP8 audio encoder weights do not match quantization_config: {shown}{hint}")nemo/collections/speechlm2/vllm/salm/model.py4 + / 0 −
@@ -65,6 +65,7 @@ ) from nemo.collections.speechlm2.vllm.salm.backends import HybridBackend, make_backend from nemo.collections.speechlm2.vllm.salm.config import _AUDIO_PLACEHOLDER+from nemo.collections.speechlm2.vllm.salm.fp8_encoder import build_fp8_encoder, check_fp8_encoder_weights _AUDIO_INPUT_DTYPE = torch.float32 _PERCEPTION_DTYPE = torch.bfloat16@@ -148,6 +149,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): getattr(config, "pe_encoder_overrides", None), ) self._uses_pe_encoder = _is_parallel_expert_encoder(getattr(self.perception, "encoder", None))+ # Built here, as vLLM builds a quantized decoder, so FP8 weights load directly.+ build_fp8_encoder(self.perception, vllm_config.quant_config) self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors @@ -288,6 +291,7 @@ def get_mm_mapping(self) -> MultiModelKeys: # ── weight loading ── def _load_perception_weights(self, perception_weights: dict[str, torch.Tensor]) -> set[str]:+ check_fp8_encoder_weights(self.perception, perception_weights) self.perception = self.perception.to(_PERCEPTION_DTYPE) incompatible = self.perception.load_state_dict(perception_weights, strict=False) nemo/collections/speechlm2/vllm/salm/mtp.py23 + / 5 −
@@ -170,10 +170,28 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: pattern = self.config.mtp_hybrid_override_pattern if not isinstance(pattern, str) or not pattern: raise ValueError(f"mtp_hybrid_override_pattern must be a non-empty string, got {pattern!r}.")- return super().load_weights(- _remap_nemo_mtp_weights(- weights,- target_vocab,- expected_layer_modules=len(pattern),+ head_tensors = 0++ def _count_head_tensors(items: Iterable[tuple[str, torch.Tensor]]) -> Iterable[tuple[str, torch.Tensor]]:+ nonlocal head_tensors+ for name, tensor in items:+ if name.startswith("mtp."):+ head_tensors += 1+ yield name, tensor++ loaded = super().load_weights(+ _count_head_tensors(+ _remap_nemo_mtp_weights(+ weights,+ target_vocab,+ expected_layer_modules=len(pattern),+ ) ) )+ if head_tensors == 0:+ raise ValueError(+ "Speculative MTP is enabled, but the checkpoint has no 'llm.mtp.*' draft-head tensors, so the "+ "draft head would run on uninitialized weights (vLLM skips its missing-weight check for quantized "+ "checkpoints). Export the MTP head with the checkpoint or disable speculative decoding."+ )+ return loadedtests/collections/asr/test_packed_transformer_encoder.py26 + / 0 −
@@ -323,6 +323,32 @@ def test_sequence_packed_fused_qkv_matches_default_and_preserves_checkpoint_keys assert set(encoder.state_dict()) == state_keys +def test_sequence_packed_qkv_runs_a_specialized_projection_through_its_forward():+ torch.manual_seed(0)+ encoder = _make_encoder(position="rope")+ inputs = torch.randn(2, 6, encoder.d_model)+ lengths = torch.tensor([6, 3])+ with torch.no_grad():+ plain = encoder.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True)+ calls = []++ class _RecordingLinear(torch.nn.Linear):+ def forward(self, x):+ calls.append(tuple(x.shape))+ return super().forward(x)++ for layer in encoder.layers:+ projection = _RecordingLinear(layer.attn.w_qkv.in_features, layer.attn.w_qkv.out_features)+ projection.load_state_dict(layer.attn.w_qkv.state_dict())+ layer.attn.w_qkv = projection++ with torch.no_grad():+ specialized = encoder.forward_sequence_packed(inputs, lengths, bypass_pre_encode=True)++ assert calls == [(9, encoder.d_model)] * len(encoder.layers)+ torch.testing.assert_close(specialized.data, plain.data, rtol=1e-5, atol=1e-6)++ def test_sequence_packed_layers_receive_only_valid_tokens(monkeypatch): encoder = _make_encoder(position="rope") encoded = torch.randn(3, 7, encoder.d_model)tests/collections/speechlm2/test_vllm_fp8_encoder.pyadded512 + / 0 −
@@ -0,0 +1,512 @@+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+# http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++"""Unit tests for serving SpeechLM checkpoints whose audio encoder is quantized with ModelOpt FP8."""++from fnmatch import fnmatch++import pytest+import torch+from torch import nn++try:+ from nemo.collections.speechlm2.vllm.salm.fp8_encoder import (+ FP8_DTYPE,+ FP8_MAX,+ FP8Linear,+ build_fp8_encoder,+ check_fp8_encoder_weights,+ )++ _HAS_PLUGIN = True+except (ImportError, RuntimeError):+ _HAS_PLUGIN = False++D_MODEL = 32+N_LAYERS = 2+N_QUANTIZED = N_LAYERS * 4+DECODER_IGNORE = ["backbone.embeddings", "lm_head", "mtp*"]+PERCEPTION_IGNORE = [+ "perception.encoder.diarization_model*",+ "perception.encoder.asr_encoder.pre_encode*",+ "perception.proj",+]+++class _QuantConfig:+ """The parts of vLLM's ModelOpt quantization config that the FP8 encoder reads."""++ def __init__(self, ignore, name="modelopt", quant_method="FP8"):+ self.exclude_modules = list(ignore)+ self.quant_method = quant_method+ self._name = name++ def get_name(self):+ return self._name++ def is_layer_excluded(self, prefix):+ return any(prefix == entry or fnmatch(prefix, entry) for entry in self.exclude_modules)+++DECODER_LAYERS = {+ "language_model.model.layers.0.mixer.in_proj": {"quant_algo": "FP8"},+ "language_model.model.layers.1.mixer.experts": {"quant_algo": "NVFP4"},+}+++class _MixedQuantConfig(_QuantConfig):+ """The parts of vLLM's ModelOpt MIXED_PRECISION config that the FP8 encoder reads."""++ def __init__(self, quantized_layers, ignore=DECODER_IGNORE):+ super().__init__(ignore, name="modelopt_mixed", quant_method=None)+ self.quantized_layers = dict(quantized_layers)+++def _asr_layer_rows(quant_algo: str = "FP8") -> dict[str, dict[str, str]]:+ return {+ f"perception.encoder.asr_encoder.layers.{i}.{leaf}": {"quant_algo": quant_algo}+ for i in range(N_LAYERS)+ for leaf in ("attn.w_qkv", "attn.out_proj", "ffn.net.0", "ffn.net.3")+ }+++class _Block(nn.Module):+ """Mirrors the Linear names of a SpeechLM transformer encoder block."""++ def __init__(self):+ super().__init__()+ self.attn = nn.Module()+ self.attn.w_qkv = nn.Linear(D_MODEL, 3 * D_MODEL)+ self.attn.out_proj = nn.Linear(D_MODEL, D_MODEL)+ self.ffn = nn.Module()+ self.ffn.net = nn.Sequential(+ nn.Linear(D_MODEL, 4 * D_MODEL), nn.ReLU(), nn.Identity(), nn.Linear(4 * D_MODEL, D_MODEL)+ )+++class _Encoder(nn.Module):+ def __init__(self):+ super().__init__()+ self.pre_encode = nn.Module()+ self.pre_encode.proj = nn.Linear(2 * D_MODEL, D_MODEL)+ self.layers = nn.ModuleList([_Block() for _ in range(N_LAYERS)])+++def _speaker_aware_perception() -> nn.Module:+ """A freshly constructed (float32) perception module laid out like a speaker-aware SpeechLM's: an ASR+ branch and a diarizer with identical Linear names, plus the connector."""+ perception = nn.Module()+ perception.encoder = nn.Module()+ perception.encoder.asr_encoder = _Encoder()+ perception.encoder.diarization_model = nn.Module()+ perception.encoder.diarization_model.encoder = _Encoder()+ perception.proj = nn.Linear(D_MODEL, D_MODEL)+ return perception+++def _asr_linears(perception: nn.Module) -> list[nn.Module]:+ return [+ linear+ for layer in perception.encoder.asr_encoder.layers+ for linear in (layer.attn.w_qkv, layer.attn.out_proj, layer.ffn.net[0], layer.ffn.net[3])+ ]+++def _fp8_bits(tensor: torch.Tensor) -> torch.Tensor:+ return tensor.view(torch.uint8)+++def _checkpoint(perception: nn.Module) -> dict[str, torch.Tensor]:+ """A random perception state dict in the checkpoint's dtypes: FP8 weights, float32 scalar scales, bfloat16 rest."""+ state = {}+ for name, tensor in perception.state_dict().items():+ if tensor.dtype == FP8_DTYPE:+ state[name] = torch.randn(tensor.shape).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)+ elif name.endswith(("weight_scale", "input_scale")):+ state[name] = torch.rand(()) * 1e-3 + 1e-4 # not representable in bfloat16+ else:+ state[name] = torch.randn(tensor.shape).to(torch.bfloat16)+ return state+++def _quantized_perception() -> nn.Module:+ perception = _speaker_aware_perception()+ build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + PERCEPTION_IGNORE))+ return perception+++@pytest.mark.skipif(not _HAS_PLUGIN, reason="SpeechLM vLLM plugin not available")+class TestBuildFP8Encoder:+ def test_builds_fp8_linears_for_the_layers_not_ignored(self):+ perception = _speaker_aware_perception()++ replaced = build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + PERCEPTION_IGNORE))++ assert replaced == N_QUANTIZED+ assert all(isinstance(linear, FP8Linear) for linear in _asr_linears(perception))+ assert sum(isinstance(m, FP8Linear) for m in perception.modules()) == N_QUANTIZED+ assert type(perception.encoder.asr_encoder.pre_encode.proj) is nn.Linear+ assert type(perception.proj) is nn.Linear++ def test_decoder_only_quant_config_leaves_the_encoder_unquantized(self):+ perception = _speaker_aware_perception()++ assert build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE)) == 0+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ @pytest.mark.parametrize("quant_config", [None, object(), _QuantConfig(DECODER_IGNORE + ["perception*"])])+ def test_no_encoder_quantization(self, quant_config):+ perception = _speaker_aware_perception()++ assert build_fp8_encoder(perception, quant_config) == 0+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ @pytest.mark.parametrize(+ ("name", "quant_method"), [("modelopt", "FP8_PER_CHANNEL_PER_TOKEN"), ("modelopt_fp4", "NVFP4")]+ )+ def test_rejects_other_schemes_before_changing_the_encoder(self, name, quant_method):+ perception = _speaker_aware_perception()++ with pytest.raises(NotImplementedError, match="ModelOpt FP8"):+ build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + PERCEPTION_IGNORE, name, quant_method))+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ def test_rejects_widths_cutlass_cannot_run_before_changing_the_encoder(self):+ perception = _speaker_aware_perception()+ perception.encoder.asr_encoder.layers[1].ffn.net[3] = nn.Linear(24, D_MODEL)++ with pytest.raises(ValueError, match="multiple of 16"):+ build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + PERCEPTION_IGNORE))+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ def test_fp8_layers_count_as_linear_and_are_not_rebuilt(self):+ perception = _quantized_perception()+ w_qkv = perception.encoder.asr_encoder.layers[0].attn.w_qkv++ assert isinstance(w_qkv, nn.Linear)+ assert build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + PERCEPTION_IGNORE)) == 0+ assert perception.encoder.asr_encoder.layers[0].attn.w_qkv is w_qkv++ @pytest.mark.parametrize("encoder_name", ["TransformerEncoder", "ConformerEncoder"])+ def test_quantized_input_projection_is_called_like_a_linear(self, encoder_name, monkeypatch):+ from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder+ from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder++ encoder_cls, kwargs = {+ "TransformerEncoder": (+ TransformerEncoder,+ {"self_attention_model": "no_pos", "sync_max_audio_length": False},+ ),+ "ConformerEncoder": (ConformerEncoder, {}),+ }[encoder_name]+ # Other test modules set CUDA as the default device at import, but ConformerEncoder+ # always creates its relative-position biases on the CPU.+ with torch.device("cpu"):+ encoder = encoder_cls(+ feat_in=32, d_model=64, n_heads=4, n_layers=1, subsampling=None, subsampling_factor=1, **kwargs+ ).eval()+ perception = nn.Module()+ perception.encoder = encoder+ build_fp8_encoder(perception, _QuantConfig(DECODER_IGNORE + ["perception.proj"]))+ monkeypatch.setattr(FP8Linear, "forward", lambda self, x: x.new_zeros(*x.shape[:-1], self.out_features))++ with torch.no_grad():+ encoded, length = encoder(audio_signal=torch.randn(1, 32, 3), length=torch.tensor([3]))[:2]++ assert isinstance(encoder.pre_encode, FP8Linear)+ assert encoded.shape == (1, 64, 3) and length.tolist() == [3]++ def test_vllm_modelopt_config_selects_the_same_layers(self):+ pytest.importorskip("vllm")+ from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config++ quant_config = ModelOptFp8Config.from_config(+ {"quant_method": "modelopt", "quant_algo": "FP8", "ignore": DECODER_IGNORE + PERCEPTION_IGNORE}+ )+ perception = _speaker_aware_perception()++ assert build_fp8_encoder(perception, quant_config) == N_QUANTIZED+ assert all(isinstance(linear, FP8Linear) for linear in _asr_linears(perception))+ assert sum(isinstance(m, FP8Linear) for m in perception.modules()) == N_QUANTIZED++ def test_mixed_precision_config_builds_the_encoder_layers_it_lists_as_fp8(self):+ perception = _speaker_aware_perception()++ replaced = build_fp8_encoder(perception, _MixedQuantConfig({**DECODER_LAYERS, **_asr_layer_rows()}))++ assert replaced == N_QUANTIZED+ assert all(isinstance(linear, FP8Linear) for linear in _asr_linears(perception))+ assert sum(isinstance(m, FP8Linear) for m in perception.modules()) == N_QUANTIZED++ def test_mixed_precision_config_without_encoder_rows_leaves_the_encoder_unquantized(self):+ perception = _speaker_aware_perception()++ assert build_fp8_encoder(perception, _MixedQuantConfig(DECODER_LAYERS)) == 0+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ def test_mixed_precision_config_rejects_other_encoder_formats_before_changing_the_encoder(self):+ perception = _speaker_aware_perception()+ rows = {**_asr_layer_rows(), "perception.encoder.asr_encoder.layers.0.ffn.net.0": {"quant_algo": "NVFP4"}}++ with pytest.raises(NotImplementedError, match="NVFP4"):+ build_fp8_encoder(perception, _MixedQuantConfig({**DECODER_LAYERS, **rows}))+ assert not any(isinstance(m, FP8Linear) for m in perception.modules())++ def test_mixed_precision_config_leaves_excluded_rows_unquantized(self):+ perception = _speaker_aware_perception()+ diarizer_row = {"perception.encoder.diarization_model.encoder.layers.0.attn.w_qkv": {"quant_algo": "FP8"}}+ quant_config = _MixedQuantConfig(+ {**DECODER_LAYERS, **_asr_layer_rows(), **diarizer_row}, ignore=DECODER_IGNORE + PERCEPTION_IGNORE+ )++ assert build_fp8_encoder(perception, quant_config) == N_QUANTIZED+ assert type(perception.encoder.diarization_model.encoder.layers[0].attn.w_qkv) is nn.Linear++ def test_vllm_mixed_precision_config_selects_the_listed_layers(self):+ pytest.importorskip("vllm")+ from vllm.model_executor.layers.quantization.modelopt import ModelOptMixedPrecisionConfig++ quant_config = ModelOptMixedPrecisionConfig.from_config(+ {+ "quant_method": "modelopt",+ "quant_algo": "MIXED_PRECISION",+ "ignore": DECODER_IGNORE,+ "quantized_layers": {**DECODER_LAYERS, **_asr_layer_rows()},+ }+ )+ perception = _speaker_aware_perception()++ assert build_fp8_encoder(perception, quant_config) == N_QUANTIZED+ assert all(isinstance(linear, FP8Linear) for linear in _asr_linears(perception))+ assert sum(isinstance(m, FP8Linear) for m in perception.modules()) == N_QUANTIZED++ def test_bfloat16_cast_keeps_fp8_weights_and_float32_scales(self):+ perception = _quantized_perception()+ state = _checkpoint(perception)+ perception.load_state_dict(state, strict=True)++ perception.to(torch.bfloat16)++ layer = perception.encoder.asr_encoder.layers[0].attn.w_qkv+ prefix = "encoder.asr_encoder.layers.0.attn.w_qkv."+ assert layer.weight.dtype == FP8_DTYPE+ assert torch.equal(_fp8_bits(layer.weight), _fp8_bits(state[prefix + "weight"]))+ for scale_name in ("weight_scale", "input_scale"):+ scale = getattr(layer, scale_name)+ assert scale.dtype == torch.float32 and torch.equal(scale, state[prefix + scale_name])+ assert layer.bias.dtype == torch.bfloat16++ def test_checkpoint_state_dict_fills_fp8_modules(self):+ perception = _quantized_perception()+ perception.to(torch.bfloat16) # the plugin casts perception right before loading+ state = _checkpoint(perception)++ result = perception.load_state_dict(state, strict=True)++ assert not result.missing_keys and not result.unexpected_keys+ w_qkv = perception.encoder.asr_encoder.layers[1].attn.w_qkv+ assert torch.equal(_fp8_bits(w_qkv.weight), _fp8_bits(state["encoder.asr_encoder.layers.1.attn.w_qkv.weight"]))+ assert torch.equal(w_qkv.input_scale, state["encoder.asr_encoder.layers.1.attn.w_qkv.input_scale"])++ def test_weight_check_accepts_a_matching_checkpoint(self):+ perception = _quantized_perception()++ check_fp8_encoder_weights(perception, _checkpoint(perception))++ @pytest.mark.parametrize(+ ("key", "value", "match"),+ [+ ("encoder.asr_encoder.layers.0.attn.w_qkv.input_scale", None, "input_scale is missing"),+ ("encoder.asr_encoder.layers.0.attn.w_qkv.weight", torch.zeros(3 * D_MODEL, D_MODEL), "not torch.float8"),+ ("encoder.asr_encoder.layers.0.ffn.net.0.weight_scale", torch.ones(4 * D_MODEL), "not one scale"),+ (+ "encoder.diarization_model.encoder.layers.0.attn.w_qkv.weight",+ torch.zeros(3 * D_MODEL, D_MODEL).to(torch.float8_e4m3fn),+ "leaves its layer unquantized.*encoder_quantization block is not read",+ ),+ ],+ )+ def test_weight_check_rejects_a_mismatched_checkpoint(self, key, value, match):+ perception = _quantized_perception()+ state = _checkpoint(perception)+ if value is None:+ del state[key]+ else:+ state[key] = value++ with pytest.raises(ValueError, match=match):+ check_fp8_encoder_weights(perception, state)++ def test_plugin_perception_loader_loads_fp8_checkpoint(self):+ pytest.importorskip("vllm")+ from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration++ perception = _quantized_perception()+ state = _checkpoint(perception)+ model = object.__new__(NeMoSpeechLMForConditionalGeneration)+ nn.Module.__init__(model)+ model.perception = perception+ model._uses_pe_encoder = True # turns on the exact-architecture check++ loaded = model._load_perception_weights(state)++ assert loaded == {f"perception.{name}" for name in state}+ w_qkv = model.perception.encoder.asr_encoder.layers[1].attn.w_qkv+ assert w_qkv.weight.dtype == FP8_DTYPE+ assert torch.equal(_fp8_bits(w_qkv.weight), _fp8_bits(state["encoder.asr_encoder.layers.1.attn.w_qkv.weight"]))+ assert model.perception.encoder.diarization_model.encoder.layers[0].attn.w_qkv.weight.dtype == torch.bfloat16++ def test_plugin_perception_loader_rejects_fp8_weights_for_an_unquantized_encoder(self):+ pytest.importorskip("vllm")+ from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration++ state = _checkpoint(_quantized_perception())+ model = object.__new__(NeMoSpeechLMForConditionalGeneration)+ nn.Module.__init__(model)+ model.perception = _speaker_aware_perception()+ model._uses_pe_encoder = True++ with pytest.raises(ValueError, match="leaves its layer unquantized"):+ model._load_perception_weights(state)+++@pytest.mark.skipif(not _HAS_PLUGIN or not torch.cuda.is_available(), reason="needs CUDA")+def test_device_move_keeps_fp8_dtypes():+ perception = _quantized_perception()++ perception.to("cuda", torch.bfloat16)++ layer = perception.encoder.asr_encoder.layers[0].attn.w_qkv+ assert layer.weight.device.type == "cuda" and layer.weight.dtype == FP8_DTYPE+ assert layer.weight_scale.device.type == "cuda" and layer.weight_scale.dtype == torch.float32+ assert layer.input_scale.device.type == "cuda" and layer.input_scale.dtype == torch.float32+ assert layer.bias.device.type == "cuda" and layer.bias.dtype == torch.bfloat16+++def _fp8_gemm_available() -> bool:+ if not _HAS_PLUGIN or not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9):+ return False+ try:+ import vllm._custom_ops # noqa: F401+ except ImportError:+ return False+ return True+++@pytest.mark.skipif(not _fp8_gemm_available(), reason="needs an FP8-capable GPU and vLLM's CUTLASS kernels")+def test_fp8_linear_matches_dequantized_reference():+ torch.manual_seed(0)+ weight = torch.randn(64, 48, device="cuda")+ weight_scale = weight.abs().amax() / FP8_MAX+ weight_fp8 = (weight / weight_scale).to(FP8_DTYPE)+ bias = torch.randn(64, device="cuda", dtype=torch.bfloat16)+ x = torch.randn(5, 7, 48, device="cuda", dtype=torch.bfloat16)+ input_scale = x.float().abs().amax() / FP8_MAX++ layer = FP8Linear(48, 64, bias=True, device="cuda")+ layer.load_state_dict(+ {"weight": weight_fp8, "weight_scale": weight_scale, "input_scale": input_scale, "bias": bias}+ )+ out = layer(x)++ # Quantize the activations the way vLLM's kernel does (multiplying by the reciprocal scale),+ # so the comparison isolates the GEMM's scaling.+ x_dequant = (x.float() * (1.0 / input_scale)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE).float() * input_scale+ reference = x_dequant @ (weight_fp8.float() * weight_scale).t() + bias.float()+ assert out.shape == (5, 7, 64)+ torch.testing.assert_close(out.float(), reference, rtol=0.02, atol=0.1)+++@pytest.mark.skipif(not _fp8_gemm_available(), reason="needs an FP8-capable GPU and vLLM's CUTLASS kernels")+def test_fp8_linear_accepts_a_strided_input():+ layer = _quantized_perception().encoder.asr_encoder.layers[0].attn.w_qkv+ layer.load_state_dict(_checkpoint(layer))+ layer.to("cuda", torch.bfloat16)+ x = torch.randn(1, D_MODEL, 7, device="cuda", dtype=torch.bfloat16).transpose(1, 2)++ assert not x.reshape(-1, D_MODEL).is_contiguous()+ torch.testing.assert_close(layer(x), layer(x.contiguous()), rtol=0, atol=0)+++def _transformer_encoder(**kwargs) -> nn.Module:+ """A small real TransformerEncoder whose Linear widths suit the CUTLASS FP8 GEMM."""+ from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder++ return TransformerEncoder(+ d_model=D_MODEL,+ n_heads=2,+ n_layers=N_LAYERS,+ ff_expansion=1.0,+ self_attention_model="rope",+ drop_rate=0.0,+ dropout_pre_encoder=0.0,+ sync_max_audio_length=False,+ **kwargs,+ ).eval()+++def _load_random_fp8(perception: nn.Module) -> None:+ perception.load_state_dict(_checkpoint(perception))+ perception.to("cuda", torch.bfloat16)+++@pytest.mark.skipif(not _fp8_gemm_available(), reason="needs an FP8-capable GPU and vLLM's CUTLASS kernels")+def test_sequence_packed_encoder_runs_quantized_qkv_projections():+ torch.manual_seed(0)+ perception = nn.Module()+ perception.encoder = _transformer_encoder(feat_in=8, subsampling_factor=2)+ ignore = DECODER_IGNORE + ["perception.encoder.pre_encode*"]+ assert build_fp8_encoder(perception, _QuantConfig(ignore)) == N_QUANTIZED+ _load_random_fp8(perception)+ x = torch.randn(2, 6, D_MODEL, device="cuda", dtype=torch.bfloat16)+ lengths = torch.tensor([6, 3], device="cuda")++ with torch.no_grad():+ default = perception.encoder.forward_sequence_packed(x, lengths, bypass_pre_encode=True)+ fused = perception.encoder.forward_sequence_packed(x, lengths, bypass_pre_encode=True, fused_qkv=True)++ assert default.data.shape == (9, D_MODEL)+ torch.testing.assert_close(default.data, fused.data, rtol=0, atol=0)+++@pytest.mark.skipif(not _fp8_gemm_available(), reason="needs an FP8-capable GPU and vLLM's CUTLASS kernels")+def test_independent_dual_encoder_with_a_quantized_asr_branch():+ from nemo.collections.asr.parts.packed_sequence import pack_encoder_output+ from nemo.collections.speechlm2.modules.perception import IndependentDualEncoder++ torch.manual_seed(0)+ branches = {+ name: _transformer_encoder(feat_in=4, subsampling="feature_stacking", subsampling_factor=2)+ for name in ("asr", "auxiliary")+ }+ perception = nn.Module()+ perception.encoder = IndependentDualEncoder(branches["asr"], branches["auxiliary"], frame_shift_seconds=0.01)+ ignore = DECODER_IGNORE + [+ "perception.encoder.auxiliary_encoder*",+ "perception.encoder.asr_encoder.pre_encode*",+ ]+ assert build_fp8_encoder(perception, _QuantConfig(ignore)) == N_QUANTIZED+ _load_random_fp8(perception)+ features = torch.randn(2, 4, 17, device="cuda", dtype=torch.bfloat16)+ lengths = torch.tensor([17, 10], device="cuda")++ with torch.no_grad():+ encoded, encoded_lengths = perception.encoder(features, lengths)+ packed = pack_encoder_output(features.transpose(1, 2), lengths)+ asr = branches["asr"].forward_sequence_packed(packed, packed.lengths, fused_qkv=True)++ assert encoded.shape == (2, 2 * D_MODEL, 9) and encoded_lengths.tolist() == [9, 5]+ assert isinstance(branches["asr"].layers[0].attn.w_qkv, FP8Linear)+ assert not isinstance(branches["auxiliary"].layers[0].attn.w_qkv, FP8Linear)+ torch.testing.assert_close(encoded[0, :D_MODEL, :9].transpose(0, 1), asr.data[:9], rtol=0, atol=0)tests/collections/speechlm2/test_vllm_plugin.py68 + / 0 −
@@ -668,6 +668,29 @@ def test_already_canonical_fp32_param_names_are_unchanged(self, backend, param_n assert mapped_name == canonical_name assert mapped_tensor is tensor + @pytest.mark.parametrize(+ ("source_name", "hf_name"),+ [+ ("llm.model.embed_tokens.weight", "backbone.embed_tokens.weight"),+ # ModelOpt exports already use NemotronH's own embedding name.+ ("backbone.embeddings.weight", "backbone.embeddings.weight"),+ ("llm.lm_head.weight", "lm_head.weight"),+ ],+ )+ def test_pads_vocabulary_tensors_to_served_vocab(self, source_name, hf_name):+ import torch++ from nemo.collections.speechlm2.vllm.salm.backends import HybridBackend++ backend = HybridBackend(SimpleNamespace(text_config=SimpleNamespace(vocab_size=6)))+ tensor = torch.ones(4, 3)+ [(mapped_name, mapped_tensor)] = backend.nemo_to_hf_llm_weights([(source_name, tensor)])++ assert mapped_name == hf_name+ assert mapped_tensor.shape == (6, 3)+ assert torch.equal(mapped_tensor[:4], tensor)+ assert not mapped_tensor[4:].any()+ @staticmethod def _lora_weights(): import torch@@ -2319,6 +2342,51 @@ def test_mtp_load_weights_rejects_invalid_serialized_pattern(self, monkeypatch, with pytest.raises(ValueError, match="non-empty string"): NeMoSpeechLMMTP.load_weights(model, []) + def test_mtp_load_weights_rejects_checkpoint_without_draft_head(self, monkeypatch):+ """A quantized export that dropped llm.mtp.* must fail instead of drafting from uninitialized weights."""+ import torch+ from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP++ from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP++ model = object.__new__(NeMoSpeechLMMTP)+ object.__setattr__(+ model,+ "config",+ SimpleNamespace(mtp_hybrid_override_pattern="*E", vocab_size=3),+ )+ monkeypatch.setattr(NemotronHMTP, "load_weights", lambda self, weights: {name for name, _ in weights})+ tensor = torch.ones(1, 2)++ with pytest.raises(ValueError, match="no 'llm.mtp"):+ NeMoSpeechLMMTP.load_weights(+ model,+ [("backbone.layers.0.mixer.in_proj.weight", tensor), ("lm_head.weight", tensor)],+ )++ def test_mtp_load_weights_accepts_checkpoint_with_draft_head(self, monkeypatch):+ """A checkpoint carrying llm.mtp.* passes the guard and hands the head tensors to vLLM's loader."""+ import torch+ from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP++ from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP++ model = object.__new__(NeMoSpeechLMMTP)+ object.__setattr__(+ model,+ "config",+ SimpleNamespace(mtp_hybrid_override_pattern="*E", vocab_size=3),+ )+ monkeypatch.setattr(NemotronHMTP, "load_weights", lambda self, weights: {name for name, _ in weights})+ tensor = torch.ones(1, 2)++ loaded = NeMoSpeechLMMTP.load_weights(+ model,+ [("backbone.layers.0.mixer.in_proj.weight", tensor), ("llm.mtp.layers.0.eh_proj.weight", tensor)],+ )++ assert loaded == {"mtp.layers.0.eh_proj.weight"}+ def test_mtp_weight_remap_splits_packed_experts(self): """Packed Automodel MTP experts must become vLLM per-expert weights.""" import torch