mla: wire dynamic per-token NVFP4 scaling and cache ABI

tests/model_executor/layers/test_mla_cache_format.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import hashlib
+
+import pytest
+
+from vllm import envs
+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_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")
+    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")
+
+
+@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(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"

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))

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,56 @@ 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"
+    )
+
+
+@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

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)

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,

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"))