From ae541fcd5ad771b54f151743c7d0ee528fb49f73 Mon Sep 17 00:00:00 2001 From: derek Date: Mon, 27 Jul 2026 19:43:47 -0400 Subject: [PATCH 1/4] glm52: self-describing two-level NVFP4 MLA KV records (VLLM_NVFP4_MLA_DYNAMIC_SCALE) The nvfp4_ds_mla writer's outer scale is statically 1.0, parking shallow-layer group scales in E4M3 subnormals (the deep-retrieval regression root cause; PR #145 calibrates it statically per layer). This mode removes the calibration requirement: the SparkInfer writer derives a per-token second-level scale stored inside the record ([292,296) of the existing pad; record stays 368 bytes) and the readers consume it from the record instead of a per-layer launch scalar. - VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 gate; requires KV_FP8_ROPE=1 368-byte records and a SparkInfer build with per_token_scale / latent_scale_per_token support (fail-closed inspect checks) - mutually exclusive with VLLM_NVFP4_MLA_SCALES_FILE (ValueError) - writer calls (do_kv_cache_update + DCP gather re-quantization) pass per_token_scale only when enabled, keeping older SparkInfer builds working with the mode off - kernel format kwargs pin latent_scale=1.0 and add latent_scale_per_token in mode Draft: not yet compiled/run (no CUDA on authoring host); CN4 validation runbook in glm52-opt design/nvfp4-dynamic-second-level- scale-phaseA-addendum.md. Co-authored-by: Fable Claude-Session: https://claude.ai/code/session_01YHuYXXq9krPdvcgpzY4EfE --- vllm/model_executor/layers/mla.py | 9 ++ .../attention/backends/mla/b12x_mla_sparse.py | 95 ++++++++++++++++--- 2 files changed, 90 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 1d008699d086..d084e99b7abf 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -221,8 +221,17 @@ def __init__( # The deployed NVFP4 writer accepts a scale tensor but discards it and # quantizes with outer scale 1.0. Feeding x/s_l is exactly the missing # writer-side normalization; the CuTe readers restore s_l in-kernel. + # VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 supersedes both: the writer derives a + # per-token second-level scale stored in the record, so the host-side + # divide stays identity and a scales file must not also be supplied. self._nvfp4_mla_outer_scale = 1.0 scale_file = os.getenv(_NVFP4_MLA_SCALES_ENV, "").strip() + if scale_file and os.getenv("VLLM_NVFP4_MLA_DYNAMIC_SCALE", "0") == "1": + raise ValueError( + f"{_NVFP4_MLA_SCALES_ENV} and VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 " + "are mutually exclusive: the dynamic mode derives per-token " + "scales in the writer and ignores static calibration" + ) if scale_file and ( self.mla_attn.kv_cache_dtype == "nvfp4_ds_mla" and self.mla_attn.attn_backend.get_name() == "B12X_MLA_SPARSE" diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 043e1ffccb37..94a28be56296 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -127,6 +127,16 @@ def _kv_fp8_rope_enabled() -> bool: return _KV_FP8_ROPE_REQUESTED and _is_glm_moe_dsa_model() +# Self-describing two-level NVFP4 records: the writer derives a per-token +# second-level scale (fp32 at record bytes [292, 296)) and the readers consume +# it from the record instead of a per-layer launch scalar. Requires the +# 368-byte KV_FP8_ROPE=1 record and a SparkInfer build with the matching +# writer/reader mode. Mutually exclusive with VLLM_NVFP4_MLA_SCALES_FILE. +_NVFP4_DYNAMIC_SCALE_REQUESTED = ( + os.getenv("VLLM_NVFP4_MLA_DYNAMIC_SCALE", "0") == "1" +) + + def _cdiv(x: int, y: int) -> int: return (int(x) + int(y) - 1) // int(y) @@ -1284,6 +1294,27 @@ def __init__( self._concat_and_cache_nvfp4_mla_fp8_rope = ( concat_and_cache_nvfp4_mla_fp8_rope ) + self._nvfp4_dynamic_scale = bool( + self._kv_fp8_rope and _NVFP4_DYNAMIC_SCALE_REQUESTED + ) + if _NVFP4_DYNAMIC_SCALE_REQUESTED and not self._kv_fp8_rope: + raise RuntimeError( + "VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 requires the 368-byte " + "KV_FP8_ROPE=1 nvfp4_ds_mla record (got kv_cache_dtype=" + f"{self.kv_cache_dtype!r}, KV_FP8_ROPE=" + f"{'1' if _kv_fp8_rope_enabled() else '0'})" + ) + if self._nvfp4_dynamic_scale and ( + "per_token_scale" + not in inspect.signature( + self._concat_and_cache_nvfp4_mla_fp8_rope + ).parameters + ): + raise RuntimeError( + "VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 requires a SparkInfer build " + "whose concat_and_cache_nvfp4_mla_fp8_rope supports " + "per_token_scale" + ) # MLA dims (absorbed: Q post-projection is [T, H, kv_lora_rank + rope]). self.kv_lora_rank: int = mla_args["kv_lora_rank"] @@ -1427,6 +1458,8 @@ def __init__( if self._b12x_scale_format is not None: required_kwargs = {"latent_scale", "scale_format"} + if self._nvfp4_dynamic_scale: + required_kwargs = required_kwargs | {"latent_scale_per_token"} unsupported_forwards = [ mode for mode, forward in ( @@ -1702,13 +1735,25 @@ def do_kv_cache_update( f"KV_FP8_ROPE writer reached a non-NVFP4 cache: {kv_cache_dtype!r}" ) k_pe_flat = k_pe.squeeze(1) - self._concat_and_cache_nvfp4_mla_fp8_rope( - kv_c_normed, - k_pe_flat, - kv_cache, - slot_mapping.flatten(), - k_scale, - ) + if self._nvfp4_dynamic_scale: + self._concat_and_cache_nvfp4_mla_fp8_rope( + kv_c_normed, + k_pe_flat, + kv_cache, + slot_mapping.flatten(), + k_scale, + per_token_scale=True, + ) + else: + # Keyword omitted so pre-two-level SparkInfer builds keep working + # when the mode is off. + self._concat_and_cache_nvfp4_mla_fp8_rope( + kv_c_normed, + k_pe_flat, + kv_cache, + slot_mapping.flatten(), + k_scale, + ) def _borrow_workspaces(self) -> list[torch.Tensor]: workspaces = current_workspace_manager().get_simultaneous( @@ -2284,13 +2329,23 @@ def _append_current_chunk_to_gathered( k_scale = getattr(layer, "_k_scale", None) if self._kv_fp8_rope: - self._concat_and_cache_nvfp4_mla_fp8_rope( - kv_c, - k_pe_flat, - gathered_buffer, - slots, - k_scale, - ) + if self._nvfp4_dynamic_scale: + self._concat_and_cache_nvfp4_mla_fp8_rope( + kv_c, + k_pe_flat, + gathered_buffer, + slots, + k_scale, + per_token_scale=True, + ) + else: + self._concat_and_cache_nvfp4_mla_fp8_rope( + kv_c, + k_pe_flat, + gathered_buffer, + slots, + k_scale, + ) elif self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla"): ops.concat_and_cache_mla( kv_c, @@ -2325,6 +2380,18 @@ def _sync_warmup(self) -> None: def _b12x_kernel_format_kwargs(self, latent_scale: float = 1.0) -> dict[str, Any]: if self._b12x_scale_format is None: return {} + if self._nvfp4_dynamic_scale: + if float(latent_scale) != 1.0: + raise RuntimeError( + "VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 is mutually exclusive with " + "a per-layer outer scale (got latent_scale=" + f"{latent_scale!r}); unset VLLM_NVFP4_MLA_SCALES_FILE" + ) + return { + "latent_scale": 1.0, + "scale_format": self._b12x_scale_format, + "latent_scale_per_token": True, + } return { "latent_scale": float(latent_scale), "scale_format": self._b12x_scale_format, From 72445d4c5e9667909c274bdea9f47b83ca7eebe5 Mon Sep 17 00:00:00 2001 From: derek Date: Tue, 28 Jul 2026 10:00:03 -0400 Subject: [PATCH 2/4] glm52: make NVFP4 MLA cache format immutable and cache-keyed Capture one server-static writer/reader mode, fail closed on incompatible SparkInfer APIs, and include the record ABI in persistent offload namespaces. Add focused tests for invalid modes, both writer sites, and stale-cache separation. Assisted-by: OpenAI Codex --- .../layers/test_mla_cache_format.py | 90 ++++++++++++++++ .../test_b12x_mla_fp8_rope_writer.py | 101 ++++++++++++++++++ tests/v1/kv_offload/test_factory.py | 22 ++++ tests/v1/kv_offload/test_file_mapper.py | 29 ++++- .../kv_connector/v1/offloading/config.py | 6 ++ vllm/model_executor/layers/mla.py | 51 ++++----- .../model_executor/layers/mla_cache_format.py | 71 ++++++++++++ .../attention/backends/mla/b12x_mla_sparse.py | 66 +++++++----- vllm/v1/kv_offload/config.py | 3 + vllm/v1/kv_offload/file_mapper.py | 7 ++ 10 files changed, 389 insertions(+), 57 deletions(-) create mode 100644 tests/model_executor/layers/test_mla_cache_format.py create mode 100644 vllm/model_executor/layers/mla_cache_format.py diff --git a/tests/model_executor/layers/test_mla_cache_format.py b/tests/model_executor/layers/test_mla_cache_format.py new file mode 100644 index 000000000000..861ef3350e50 --- /dev/null +++ b/tests/model_executor/layers/test_mla_cache_format.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import hashlib + +import pytest + +from vllm.model_executor.layers.mla_cache_format import ( + KV_FP8_ROPE_ENV, + NVFP4_MLA_DYNAMIC_SCALE_ENV, + NVFP4_MLA_SCALES_ENV, + Nvfp4MlaCacheFormat, +) + + +def test_from_env_captures_one_server_static_mode(monkeypatch): + monkeypatch.setenv(NVFP4_MLA_DYNAMIC_SCALE_ENV, "1") + monkeypatch.setenv(KV_FP8_ROPE_ENV, "1") + monkeypatch.delenv(NVFP4_MLA_SCALES_ENV, raising=False) + + cache_format = Nvfp4MlaCacheFormat.from_env() + monkeypatch.setenv(NVFP4_MLA_DYNAMIC_SCALE_ENV, "0") + + assert cache_format.dynamic_scale + assert cache_format.fp8_rope + assert cache_format.scales_file == "" + assert ( + cache_format.record_abi("nvfp4_ds_mla") + == "nvfp4_ds_mla:fp8-rope-368:dynamic-token-v1" + ) + + +@pytest.mark.parametrize( + "cache_format", + [ + Nvfp4MlaCacheFormat( + dynamic_scale=True, + fp8_rope=True, + scales_file="/tmp/static-scales.json", + ), + Nvfp4MlaCacheFormat( + dynamic_scale=True, + fp8_rope=False, + scales_file="", + ), + ], +) +def test_invalid_dynamic_combinations_fail_closed(cache_format): + with pytest.raises(ValueError): + cache_format.validate() + + +def test_static_scale_contents_participate_in_record_abi(tmp_path): + scales = tmp_path / "scales.json" + payload = b'{"format":"example","scales":[1.0]}' + scales.write_bytes(payload) + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=False, + fp8_rope=True, + scales_file=str(scales), + ) + + expected_digest = hashlib.sha256(payload).hexdigest() + assert cache_format.record_abi("nvfp4_ds_mla") == ( + f"nvfp4_ds_mla:fp8-rope-368:static-calibrated-v1:{expected_digest}" + ) + + scales.write_bytes(b'{"format":"example","scales":[2.0]}') + assert cache_format.record_abi("nvfp4_ds_mla") != ( + f"nvfp4_ds_mla:fp8-rope-368:static-calibrated-v1:{expected_digest}" + ) + + +def test_missing_static_scale_file_cannot_form_persistent_abi(tmp_path): + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=False, + fp8_rope=True, + scales_file=str(tmp_path / "missing.json"), + ) + with pytest.raises(ValueError, match="Cannot fingerprint"): + cache_format.record_abi("nvfp4_ds_mla") + + +def test_non_nvfp4_cache_abi_is_unaffected_by_nvfp4_mode(): + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=True, + fp8_rope=False, + scales_file="/does/not/matter", + ) + assert cache_format.record_abi("bfloat16") == "bfloat16:default-v1" diff --git a/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py b/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py index 0063050b8fa6..b0dff7f5a5fb 100644 --- a/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py +++ b/tests/v1/attention/test_b12x_mla_fp8_rope_writer.py @@ -3,6 +3,7 @@ import sys import types +from types import SimpleNamespace import pytest import torch @@ -55,6 +56,7 @@ def _construct_through_writer_binding( enabled: bool, ) -> B12xMLASparseImpl: monkeypatch.setattr(b12x_mla_sparse, "_KV_FP8_ROPE_REQUESTED", enabled) + monkeypatch.setattr(b12x_mla_sparse, "_NVFP4_DYNAMIC_SCALE_REQUESTED", False) monkeypatch.setattr(b12x_mla_sparse, "_IS_GLM_MOE_DSA_CACHE", True) def stop_after_writer_binding(): @@ -106,6 +108,7 @@ def _install_fake_writer_package( def _enabled_impl(writer) -> B12xMLASparseImpl: impl = object.__new__(B12xMLASparseImpl) impl._kv_fp8_rope = True + impl._nvfp4_dynamic_scale = False impl._concat_and_cache_nvfp4_mla_fp8_rope = writer return impl @@ -333,3 +336,101 @@ def reject_fallback_initialization(): match="compact writer initialization failed", ): _initialize_writer_seam(object.__new__(B12xMLASparseImpl)) + + +def test_dynamic_mode_rejects_writer_without_per_token_scale( + monkeypatch: pytest.MonkeyPatch, +): + def legacy_writer(kv_c, k_pe, kv_cache, slot_mapping, scale): + pass + + _install_fake_writer_package(monkeypatch, legacy_writer) + monkeypatch.setattr(b12x_mla_sparse, "_KV_FP8_ROPE_REQUESTED", True) + monkeypatch.setattr(b12x_mla_sparse, "_NVFP4_DYNAMIC_SCALE_REQUESTED", True) + monkeypatch.setattr(b12x_mla_sparse, "_IS_GLM_MOE_DSA_CACHE", True) + + with pytest.raises(RuntimeError, match="per_token_scale"): + _initialize_writer_seam(object.__new__(B12xMLASparseImpl)) + + +def test_dynamic_mode_rejects_non_368_byte_layout( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(b12x_mla_sparse, "_KV_FP8_ROPE_REQUESTED", False) + monkeypatch.setattr(b12x_mla_sparse, "_NVFP4_DYNAMIC_SCALE_REQUESTED", True) + monkeypatch.setattr(b12x_mla_sparse, "_IS_GLM_MOE_DSA_CACHE", True) + + with pytest.raises(RuntimeError, match="requires the 368-byte"): + _initialize_writer_seam(object.__new__(B12xMLASparseImpl)) + + +def test_dynamic_mode_rejects_reader_without_per_token_scale(): + def decode(*, latent_scale, scale_format): + pass + + def extend(*, latent_scale, scale_format, latent_scale_per_token): + pass + + with pytest.raises(RuntimeError, match="unsupported: decode"): + b12x_mla_sparse._require_callable_parameters( + "dynamic readers", + (("decode", decode), ("extend", extend)), + frozenset({"latent_scale", "scale_format", "latent_scale_per_token"}), + ) + + +def test_dynamic_normal_writer_call_propagates_per_token_scale(): + writer_calls = [] + + def writer(*args, **kwargs): + writer_calls.append((args, kwargs)) + + impl = _enabled_impl(writer) + impl._nvfp4_dynamic_scale = True + impl.do_kv_cache_update( + torch.zeros((2, 512), dtype=torch.bfloat16), + torch.zeros((2, 1, 64), dtype=torch.bfloat16), + torch.empty((1, 2, 368), dtype=torch.uint8), + torch.tensor([[0], [1]], dtype=torch.int64), + "nvfp4_ds_mla", + torch.tensor(1.0), + ) + + assert len(writer_calls) == 1 + assert writer_calls[0][1] == {"per_token_scale": True} + + +def test_dynamic_gathered_chunk_writer_call_propagates_per_token_scale(): + writer_calls = [] + + def writer(*args, **kwargs): + writer_calls.append((args, kwargs)) + + impl = _enabled_impl(writer) + impl._nvfp4_dynamic_scale = True + impl._ckv_current_chunk_kv_c = torch.zeros((2, 512), dtype=torch.bfloat16) + impl._ckv_current_chunk_kpe = torch.zeros((2, 64), dtype=torch.bfloat16) + impl.device = torch.device("cpu") + impl.cp_kv_cache_interleave_size = 1 + impl.dcp_world_size = 1 + impl.kv_cache_dtype = "nvfp4_ds_mla" + metadata = SimpleNamespace( + num_reqs=1, + global_cache_seq_lens_per_req=torch.tensor([2], dtype=torch.int32), + req_id_per_token=torch.tensor([0, 0], dtype=torch.int32), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32), + dcp_rank_req_starts=torch.tensor([[0]], dtype=torch.int32), + dcp_padded_total_tokens=64, + ) + layer = SimpleNamespace(_k_scale=torch.tensor(1.0)) + + impl._append_current_chunk_to_gathered( + torch.empty((64, 368), dtype=torch.uint8), + metadata, + layer, + num_actual_toks=2, + ) + + assert len(writer_calls) == 1 + assert writer_calls[0][1] == {"per_token_scale": True} + assert torch.equal(writer_calls[0][0][3], torch.tensor([0, 1], dtype=torch.int64)) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 2367e3c756df..1d4af44c8bc7 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, ) +from vllm.model_executor.layers.mla_cache_format import Nvfp4MlaCacheFormat from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -447,6 +448,27 @@ def test_offloading_config_preserves_data_parallel_index(): assert offloading_config.parallel.data_parallel_index == 2 +def test_offloading_config_carries_nvfp4_record_abi(): + config = _make_layout_vllm_config() + config.cache_config.cache_dtype = "nvfp4_ds_mla" + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=True, + fp8_rope=True, + scales_file="", + ) + + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.offloading.config." + "NVFP4_MLA_CACHE_FORMAT", + cache_format, + ): + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.model.kv_cache_abi == ( + "nvfp4_ds_mla:fp8-rope-368:dynamic-token-v1" + ) + + def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes(): config = _make_layout_vllm_config(cpu_bytes_to_use=65536) config.cache_config.block_size = 4 diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 6c11f2d465fb..4ff4d7f91c25 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -29,13 +29,14 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: ) for tokens_per_block, layer_name in kwargs.get("groups", ()) ), - worker_kv_bytes_per_block=0, + worker_kv_bytes_per_block=kwargs.get("worker_kv_bytes_per_block", 0), enable_kv_cache_events=False, extra_config={}, engine_id="test-engine", model=OffloadingModelConfig( name=kwargs.get("model_name", "test-model"), dtype=kwargs.get("dtype", "float16"), + kv_cache_abi=kwargs.get("kv_cache_abi", "vllm-default-v1"), ), cache=OffloadingCacheConfig( tokens_per_hash=kwargs.get("tokens_per_hash", 16), @@ -143,6 +144,32 @@ def test_hybrid_file_identity_uses_resolved_tokens_per_hash(): ] +def test_record_abi_and_geometry_separate_persistent_namespaces(): + static = make_mapper_from_offloading_spec( + kv_cache_abi="nvfp4_ds_mla:fp8-rope-368:static-calibrated-v1:abc", + worker_kv_bytes_per_block=23552, + ) + dynamic = make_mapper_from_offloading_spec( + kv_cache_abi="nvfp4_ds_mla:fp8-rope-368:dynamic-token-v1", + worker_kv_bytes_per_block=23552, + ) + different_geometry = make_mapper_from_offloading_spec( + kv_cache_abi="nvfp4_ds_mla:fp8-rope-368:dynamic-token-v1", + worker_kv_bytes_per_block=27648, + ) + + assert static.base_path != dynamic.base_path + assert dynamic.base_path != different_geometry.base_path + assert dynamic.fields["kv_cache_abi"].endswith("dynamic-token-v1") + assert dynamic.fields["worker_kv_bytes_per_block"] == 23552 + + +def test_default_record_abi_preserves_existing_namespace(): + default = make_mapper_from_offloading_spec() + assert "kv_cache_abi" not in default.fields + assert "worker_kv_bytes_per_block" not in default.fields + + # --------------------------------------------------------------------------- # parallel_agnostic: opt-in honored only when the config marks the layout # parallelism-agnostic (predicate computation is covered in test_factory.py) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index f10654c493d1..f2f17c1e8c14 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -4,6 +4,9 @@ from typing import TYPE_CHECKING +from vllm.model_executor.layers.mla_cache_format import ( + NVFP4_MLA_CACHE_FORMAT, +) from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec from vllm.v1.kv_offload.config import ( @@ -144,6 +147,9 @@ def _tokens_per_block(kv_cache_spec: "KVCacheSpec") -> int: model=OffloadingModelConfig( name=vllm_config.model_config.model, dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""), + kv_cache_abi=NVFP4_MLA_CACHE_FORMAT.record_abi( + str(vllm_config.cache_config.cache_dtype) + ), ), cache=OffloadingCacheConfig( tokens_per_hash=tokens_per_hash, diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index d084e99b7abf..b01ceb97c252 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import math -import os import re from dataclasses import dataclass from functools import cache @@ -12,15 +11,19 @@ from vllm.config import CacheConfig, get_current_vllm_config from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.mla_cache_format import ( + NVFP4_MLA_CACHE_FORMAT, + NVFP4_MLA_SCALES_ENV, +) from vllm.model_executor.layers.quantization import QuantizationConfig -_NVFP4_MLA_SCALES_ENV = "VLLM_NVFP4_MLA_SCALES_FILE" +_NVFP4_MLA_SCALES_ENV = NVFP4_MLA_SCALES_ENV _NVFP4_MLA_SCALES_FORMAT = "nvfp4_ds_mla_outer_scale_v1" _NVFP4_MLA_NUM_LAYERS = 78 _NVFP4_MLA_LATENT_DIM = 512 _NVFP4_MLA_SCALE_DENOMINATOR = 6.0 * 448.0 _NVFP4_MLA_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)") -_KV_FP8_ROPE_ENABLED = os.getenv("KV_FP8_ROPE", "0") == "1" +_KV_FP8_ROPE_ENABLED = NVFP4_MLA_CACHE_FORMAT.fp8_rope _IS_GLM_MOE_DSA_CACHE: bool | None = None @@ -45,9 +48,7 @@ def _is_glm_moe_dsa_model() -> bool: _IS_GLM_MOE_DSA_CACHE = True return True speculative_config = getattr(vllm_config, "speculative_config", None) - target_model_config = getattr( - speculative_config, "target_model_config", None - ) + target_model_config = getattr(speculative_config, "target_model_config", None) target_model_type = ( getattr(target_model_config.hf_config, "model_type", None) if target_model_config is not None @@ -67,28 +68,27 @@ def _load_nvfp4_mla_outer_scales(path: str) -> tuple[float, ...]: raise ValueError(f"{_NVFP4_MLA_SCALES_ENV} must contain a JSON object") if payload.get("format") != _NVFP4_MLA_SCALES_FORMAT: raise ValueError( - f"{_NVFP4_MLA_SCALES_ENV} has unsupported format " - f"{payload.get('format')!r}" + f"{_NVFP4_MLA_SCALES_ENV} has unsupported format {payload.get('format')!r}" ) if type(payload.get("num_layers")) is not int or ( payload["num_layers"] != _NVFP4_MLA_NUM_LAYERS ): raise ValueError( - f"{_NVFP4_MLA_SCALES_ENV} must declare " - f"num_layers={_NVFP4_MLA_NUM_LAYERS}" + f"{_NVFP4_MLA_SCALES_ENV} must declare num_layers={_NVFP4_MLA_NUM_LAYERS}" ) if type(payload.get("latent_dim")) is not int or ( payload["latent_dim"] != _NVFP4_MLA_LATENT_DIM ): raise ValueError( - f"{_NVFP4_MLA_SCALES_ENV} must declare " - f"latent_dim={_NVFP4_MLA_LATENT_DIM}" + f"{_NVFP4_MLA_SCALES_ENV} must declare latent_dim={_NVFP4_MLA_LATENT_DIM}" ) denominator = payload.get("denominator") - if isinstance(denominator, bool) or not isinstance( - denominator, (int, float) - ) or not math.isclose( - float(denominator), _NVFP4_MLA_SCALE_DENOMINATOR, rel_tol=0.0, abs_tol=0.0 + if ( + isinstance(denominator, bool) + or not isinstance(denominator, (int, float)) + or not math.isclose( + float(denominator), _NVFP4_MLA_SCALE_DENOMINATOR, rel_tol=0.0, abs_tol=0.0 + ) ): raise ValueError( f"{_NVFP4_MLA_SCALES_ENV} must declare " @@ -225,13 +225,8 @@ def __init__( # per-token second-level scale stored in the record, so the host-side # divide stays identity and a scales file must not also be supplied. self._nvfp4_mla_outer_scale = 1.0 - scale_file = os.getenv(_NVFP4_MLA_SCALES_ENV, "").strip() - if scale_file and os.getenv("VLLM_NVFP4_MLA_DYNAMIC_SCALE", "0") == "1": - raise ValueError( - f"{_NVFP4_MLA_SCALES_ENV} and VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 " - "are mutually exclusive: the dynamic mode derives per-token " - "scales in the writer and ignores static calibration" - ) + NVFP4_MLA_CACHE_FORMAT.validate() + scale_file = NVFP4_MLA_CACHE_FORMAT.scales_file if scale_file and ( self.mla_attn.kv_cache_dtype == "nvfp4_ds_mla" and self.mla_attn.attn_backend.get_name() == "B12X_MLA_SPARSE" @@ -248,9 +243,9 @@ def __init__( # That layer is deep/late (not underflowing) and its KV is transient, # so identity there is a safe no-op for KLD. if 0 <= layer_idx < _NVFP4_MLA_NUM_LAYERS: - self._nvfp4_mla_outer_scale = _load_nvfp4_mla_outer_scales( - scale_file - )[layer_idx] + self._nvfp4_mla_outer_scale = _load_nvfp4_mla_outer_scales(scale_file)[ + layer_idx + ] # forward_mqa receives this MLAAttention object as ``layer``. Keep a # host float here so no device .item() or per-call scale tensor is needed. self.mla_attn._nvfp4_mla_outer_scale = self._nvfp4_mla_outer_scale @@ -323,9 +318,7 @@ def forward( q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( positions, q[..., self.qk_nope_head_dim :], k_pe ) - if self._kv_fp8_rope and ( - k_pe.dtype != torch.bfloat16 or k_pe.shape[-1] != 64 - ): + if self._kv_fp8_rope and (k_pe.dtype != torch.bfloat16 or k_pe.shape[-1] != 64): raise RuntimeError( "KV_FP8_ROPE POST-RoPE writer requires BF16 k_pe[...,64], got " f"dtype={k_pe.dtype}, shape={tuple(k_pe.shape)}" diff --git a/vllm/model_executor/layers/mla_cache_format.py b/vllm/model_executor/layers/mla_cache_format.py new file mode 100644 index 000000000000..e3e7ac2375c6 --- /dev/null +++ b/vllm/model_executor/layers/mla_cache_format.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Server-static NVFP4 MLA cache-format configuration and ABI identity.""" + +from __future__ import annotations + +import hashlib +import os +from dataclasses import dataclass +from pathlib import Path + +NVFP4_MLA_SCALES_ENV = "VLLM_NVFP4_MLA_SCALES_FILE" +NVFP4_MLA_DYNAMIC_SCALE_ENV = "VLLM_NVFP4_MLA_DYNAMIC_SCALE" +KV_FP8_ROPE_ENV = "KV_FP8_ROPE" + + +@dataclass(frozen=True) +class Nvfp4MlaCacheFormat: + """Immutable writer/reader configuration captured at process import.""" + + dynamic_scale: bool + fp8_rope: bool + scales_file: str + + @classmethod + def from_env(cls) -> Nvfp4MlaCacheFormat: + return cls( + dynamic_scale=os.getenv(NVFP4_MLA_DYNAMIC_SCALE_ENV, "0") == "1", + fp8_rope=os.getenv(KV_FP8_ROPE_ENV, "0") == "1", + scales_file=os.getenv(NVFP4_MLA_SCALES_ENV, "").strip(), + ) + + def validate(self) -> None: + if self.dynamic_scale and self.scales_file: + raise ValueError( + f"{NVFP4_MLA_SCALES_ENV} and " + f"{NVFP4_MLA_DYNAMIC_SCALE_ENV}=1 are mutually exclusive" + ) + if self.dynamic_scale and not self.fp8_rope: + raise ValueError( + f"{NVFP4_MLA_DYNAMIC_SCALE_ENV}=1 requires {KV_FP8_ROPE_ENV}=1" + ) + + def record_abi(self, cache_dtype: str) -> str: + """Return an identity suitable for persistent external-cache keys.""" + normalized_dtype = str(cache_dtype).replace("torch.", "") + if normalized_dtype != "nvfp4_ds_mla": + return f"{normalized_dtype}:default-v1" + + self.validate() + layout = "fp8-rope-368" if self.fp8_rope else "bf16-rope-432" + if self.dynamic_scale: + scale_mode = "dynamic-token-v1" + elif self.scales_file: + try: + scale_digest = hashlib.sha256( + Path(self.scales_file).read_bytes() + ).hexdigest() + except OSError as exc: + raise ValueError( + f"Cannot fingerprint {NVFP4_MLA_SCALES_ENV}={self.scales_file!r}" + ) from exc + scale_mode = f"static-calibrated-v1:{scale_digest}" + else: + scale_mode = "implicit-unit-v1" + return f"nvfp4_ds_mla:{layout}:{scale_mode}" + + +# All consumers import this one frozen value, so a process cannot configure +# the writer, readers, and external-cache namespace from different env reads. +NVFP4_MLA_CACHE_FORMAT = Nvfp4MlaCacheFormat.from_env() diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 94a28be56296..9407c4b64574 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -40,6 +40,9 @@ from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import get_mla_dims +from vllm.model_executor.layers.mla_cache_format import ( + NVFP4_MLA_CACHE_FORMAT, +) from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( @@ -78,7 +81,7 @@ _EXTEND_PREWARM_DONE: set[ tuple[int | None, int, int, int, int, int, bool, str, bool] ] = set() -_KV_FP8_ROPE_REQUESTED = os.getenv("KV_FP8_ROPE", "0") == "1" +_KV_FP8_ROPE_REQUESTED = NVFP4_MLA_CACHE_FORMAT.fp8_rope _IS_GLM_MOE_DSA_CACHE: bool | None = None @@ -127,14 +130,29 @@ def _kv_fp8_rope_enabled() -> bool: return _KV_FP8_ROPE_REQUESTED and _is_glm_moe_dsa_model() -# Self-describing two-level NVFP4 records: the writer derives a per-token +# Inline-scale two-level NVFP4 records: the writer derives a per-token # second-level scale (fp32 at record bytes [292, 296)) and the readers consume # it from the record instead of a per-layer launch scalar. Requires the # 368-byte KV_FP8_ROPE=1 record and a SparkInfer build with the matching # writer/reader mode. Mutually exclusive with VLLM_NVFP4_MLA_SCALES_FILE. -_NVFP4_DYNAMIC_SCALE_REQUESTED = ( - os.getenv("VLLM_NVFP4_MLA_DYNAMIC_SCALE", "0") == "1" -) +_NVFP4_DYNAMIC_SCALE_REQUESTED = NVFP4_MLA_CACHE_FORMAT.dynamic_scale + + +def _require_callable_parameters( + feature: str, + callables: tuple[tuple[str, Any], ...], + required: frozenset[str], +) -> None: + unsupported = [ + name + for name, function in callables + if not required.issubset(inspect.signature(function).parameters) + ] + if unsupported: + raise RuntimeError( + f"{feature} requires SparkInfer callables accepting " + f"{sorted(required)!r}; unsupported: {', '.join(unsupported)}" + ) def _cdiv(x: int, y: int) -> int: @@ -1304,16 +1322,16 @@ def __init__( f"{self.kv_cache_dtype!r}, KV_FP8_ROPE=" f"{'1' if _kv_fp8_rope_enabled() else '0'})" ) - if self._nvfp4_dynamic_scale and ( - "per_token_scale" - not in inspect.signature( - self._concat_and_cache_nvfp4_mla_fp8_rope - ).parameters - ): - raise RuntimeError( - "VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 requires a SparkInfer build " - "whose concat_and_cache_nvfp4_mla_fp8_rope supports " - "per_token_scale" + if self._nvfp4_dynamic_scale: + _require_callable_parameters( + "VLLM_NVFP4_MLA_DYNAMIC_SCALE=1 writer", + ( + ( + "concat_and_cache_nvfp4_mla_fp8_rope", + self._concat_and_cache_nvfp4_mla_fp8_rope, + ), + ), + frozenset({"per_token_scale"}), ) # MLA dims (absorbed: Q post-projection is [T, H, kv_lora_rank + rope]). @@ -1460,20 +1478,14 @@ def __init__( required_kwargs = {"latent_scale", "scale_format"} if self._nvfp4_dynamic_scale: required_kwargs = required_kwargs | {"latent_scale_per_token"} - unsupported_forwards = [ - mode - for mode, forward in ( + _require_callable_parameters( + "B12X_MLA_SPARSE with kv_cache_dtype='nvfp4_ds_mla'", + ( ("decode", sparse_mla_decode_forward), ("extend", sparse_mla_extend_forward), - ) - if not required_kwargs.issubset(inspect.signature(forward).parameters) - ] - if unsupported_forwards: - raise RuntimeError( - "B12X_MLA_SPARSE with kv_cache_dtype='nvfp4_ds_mla' " - "requires a b12x build with NVFP4 sparse-MLA API support; " - "unsupported forwards: " + ", ".join(unsupported_forwards) - ) + ), + frozenset(required_kwargs), + ) # Eager PLAN -> BIND -> KERNEL (no b12x workspace/arena, ever). We build a # caller-owned-scratch PLAN once per mode; each forward maps a vLLM diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index cd7b3ee2075a..266a27caaa37 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -22,6 +22,9 @@ class OffloadingModelConfig: name: str # KV cache data type (e.g. "float16"). dtype: str + # Versioned identity of the bytes stored in external KV caches. Formats + # with identical tensor shapes but different reader semantics must differ. + kv_cache_abi: str = "vllm-default-v1" @dataclass(frozen=True) diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index b85d4d069790..c5d0f7e2b8ab 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -32,6 +32,8 @@ def __init__( dcp_size: int, rank: int, dtype: str, + kv_cache_abi: str = "vllm-default-v1", + worker_kv_bytes_per_block: int = 0, kv_cache_groups: list[dict] | None = None, inference_engine: str = "vllm", parallel_agnostic: bool = False, @@ -58,6 +60,9 @@ def __init__( "kv_cache_groups": kv_cache_groups or [], "inference_engine": inference_engine, } + if kv_cache_abi != "vllm-default-v1": + self.fields["kv_cache_abi"] = str(kv_cache_abi) + self.fields["worker_kv_bytes_per_block"] = int(worker_kv_bytes_per_block) self.base_path: str = self._compute_base_path(root_dir, self.fields) @classmethod @@ -89,6 +94,8 @@ def from_offloading_spec( dcp_size=parallel.dcp_size, rank=parallel.rank, dtype=config.model.dtype, + kv_cache_abi=config.model.kv_cache_abi, + worker_kv_bytes_per_block=config.worker_kv_bytes_per_block, kv_cache_groups=kv_cache_groups, parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic), ) From 1435b01512fb3e9a5dcac884877ae520a92e302a Mon Sep 17 00:00:00 2001 From: derek Date: Tue, 28 Jul 2026 10:18:53 -0400 Subject: [PATCH 3/4] glm52: register NVFP4 MLA scale environment Recognize the dynamic and static calibration variables so validated configurations do not emit unknown-vLLM-environment warnings. Assisted-by: OpenAI Codex --- tests/model_executor/layers/test_mla_cache_format.py | 6 ++++++ vllm/envs.py | 8 ++++++++ 2 files changed, 14 insertions(+) diff --git a/tests/model_executor/layers/test_mla_cache_format.py b/tests/model_executor/layers/test_mla_cache_format.py index 861ef3350e50..13a87d8a4ec0 100644 --- a/tests/model_executor/layers/test_mla_cache_format.py +++ b/tests/model_executor/layers/test_mla_cache_format.py @@ -5,6 +5,7 @@ import pytest +from vllm import envs from vllm.model_executor.layers.mla_cache_format import ( KV_FP8_ROPE_ENV, NVFP4_MLA_DYNAMIC_SCALE_ENV, @@ -13,6 +14,11 @@ ) +def test_cache_format_envs_are_registered(): + assert NVFP4_MLA_DYNAMIC_SCALE_ENV in envs.environment_variables + assert NVFP4_MLA_SCALES_ENV in envs.environment_variables + + def test_from_env_captures_one_server_static_mode(monkeypatch): monkeypatch.setenv(NVFP4_MLA_DYNAMIC_SCALE_ENV, "1") monkeypatch.setenv(KV_FP8_ROPE_ENV, "1") diff --git a/vllm/envs.py b/vllm/envs.py index 19bb0fe60c12..44d898c3be98 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -61,6 +61,8 @@ VLLM_USE_B12X_SPARSE_INDEXER: bool = False VLLM_USE_B12X_MHC: bool = False VLLM_USE_B12X_FP8_GEMM: bool = False + VLLM_NVFP4_MLA_DYNAMIC_SCALE: bool = False + VLLM_NVFP4_MLA_SCALES_FILE: str = "" VLLM_B12X_ABSORB_BMM: bool = False VLLM_DSPARK_FP8_DRAFT_HEAD: bool = False VLLM_USE_B12X_WO_PROJECTION: bool = False @@ -1116,6 +1118,12 @@ def _resolve_rust_frontend_path() -> str | None: # Use b12x for FP4 MoE experts. # This is opt-in while the b12x subsystems are brought over one at a time. "VLLM_USE_B12X_MOE": lambda: bool(int(os.getenv("VLLM_USE_B12X_MOE", "0"))), + "VLLM_NVFP4_MLA_DYNAMIC_SCALE": lambda: bool( + int(os.getenv("VLLM_NVFP4_MLA_DYNAMIC_SCALE", "0")) + ), + "VLLM_NVFP4_MLA_SCALES_FILE": lambda: os.getenv( + "VLLM_NVFP4_MLA_SCALES_FILE", "" + ).strip(), # Exact TP4 GLM-5.2 E64-NVFP4/E192-NF3 one-grid decode specialization. "VLLM_NF3_GRID188_DECODE": lambda: bool( int(os.getenv("VLLM_NF3_GRID188_DECODE", "1")) From b57062274c3f53bec69b431bfae7230977f5f10c Mon Sep 17 00:00:00 2001 From: derek Date: Tue, 28 Jul 2026 10:52:50 -0400 Subject: [PATCH 4/4] glm52: preserve default offload cache namespaces --- .../layers/test_mla_cache_format.py | 15 ++++++++-- tests/v1/kv_offload/test_factory.py | 29 +++++++++++++++++++ .../model_executor/layers/mla_cache_format.py | 12 +++++--- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/layers/test_mla_cache_format.py b/tests/model_executor/layers/test_mla_cache_format.py index 13a87d8a4ec0..9a128a5fc7e9 100644 --- a/tests/model_executor/layers/test_mla_cache_format.py +++ b/tests/model_executor/layers/test_mla_cache_format.py @@ -87,10 +87,21 @@ def test_missing_static_scale_file_cannot_form_persistent_abi(tmp_path): cache_format.record_abi("nvfp4_ds_mla") -def test_non_nvfp4_cache_abi_is_unaffected_by_nvfp4_mode(): +@pytest.mark.parametrize("cache_dtype", ["bfloat16", "float16", "auto"]) +def test_non_nvfp4_cache_abi_preserves_existing_namespace(cache_dtype): cache_format = Nvfp4MlaCacheFormat( dynamic_scale=True, fp8_rope=False, scales_file="/does/not/matter", ) - assert cache_format.record_abi("bfloat16") == "bfloat16:default-v1" + assert cache_format.record_abi(cache_dtype) == "vllm-default-v1" + + +@pytest.mark.parametrize("fp8_rope", [False, True]) +def test_implicit_nvfp4_cache_abi_preserves_existing_namespace(fp8_rope): + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=False, + fp8_rope=fp8_rope, + scales_file="", + ) + assert cache_format.record_abi("nvfp4_ds_mla") == "vllm-default-v1" diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 1d4af44c8bc7..98191db57b73 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -469,6 +469,35 @@ def test_offloading_config_carries_nvfp4_record_abi(): ) +@pytest.mark.parametrize("cache_dtype", [torch.float16, torch.bfloat16, "auto"]) +def test_offloading_config_preserves_default_record_abi(cache_dtype): + config = _make_layout_vllm_config() + config.cache_config.cache_dtype = cache_dtype + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.model.kv_cache_abi == "vllm-default-v1" + + +def test_offloading_config_preserves_implicit_nvfp4_record_abi(): + config = _make_layout_vllm_config() + config.cache_config.cache_dtype = "nvfp4_ds_mla" + cache_format = Nvfp4MlaCacheFormat( + dynamic_scale=False, + fp8_rope=True, + scales_file="", + ) + + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.offloading.config." + "NVFP4_MLA_CACHE_FORMAT", + cache_format, + ): + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.model.kv_cache_abi == "vllm-default-v1" + + def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes(): config = _make_layout_vllm_config(cpu_bytes_to_use=65536) config.cache_config.block_size = 4 diff --git a/vllm/model_executor/layers/mla_cache_format.py b/vllm/model_executor/layers/mla_cache_format.py index e3e7ac2375c6..c428b906eece 100644 --- a/vllm/model_executor/layers/mla_cache_format.py +++ b/vllm/model_executor/layers/mla_cache_format.py @@ -45,13 +45,19 @@ def record_abi(self, cache_dtype: str) -> str: """Return an identity suitable for persistent external-cache keys.""" normalized_dtype = str(cache_dtype).replace("torch.", "") if normalized_dtype != "nvfp4_ds_mla": - return f"{normalized_dtype}:default-v1" + return "vllm-default-v1" self.validate() + if not self.dynamic_scale and not self.scales_file: + # Preserve the existing namespace for every unconfigured/default + # deployment. Only modes that change the record's scale semantics + # opt into a new external-cache identity. + return "vllm-default-v1" + layout = "fp8-rope-368" if self.fp8_rope else "bf16-rope-432" if self.dynamic_scale: scale_mode = "dynamic-token-v1" - elif self.scales_file: + else: try: scale_digest = hashlib.sha256( Path(self.scales_file).read_bytes() @@ -61,8 +67,6 @@ def record_abi(self, cache_dtype: str) -> str: f"Cannot fingerprint {NVFP4_MLA_SCALES_ENV}={self.scales_file!r}" ) from exc scale_mode = f"static-calibrated-v1:{scale_digest}" - else: - scale_mode = "implicit-unit-v1" return f"nvfp4_ds_mla:{layout}:{scale_mode}"