models: load MXFP8 sparse indexer WK scales

tests/models/test_deepseek_v2_indexer_loader.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    _mxfp8_e4m3_quantize_torch,
+    dequant_mxfp8_to_bf16,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    scaled_dequantize,
+)
+from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk
+
+
+class _LoadedParam:
+
+    def __init__(self):
+        self.loaded_weight = None
+        self.loaded_shard_id = None
+
+    def weight_loader(self, _param, weight, shard_id):
+        self.loaded_weight = weight
+        self.loaded_shard_id = shard_id
+
+
+def test_try_load_fp8_indexer_wk_consumes_mxfp8_weight_scale():
+    prefix = "layers.0.self_attn.indexer"
+    param = _LoadedParam()
+    params_dict = {f"{prefix}.wk_weights_proj.weight": param}
+    loaded_params: set[str] = set()
+    pending = {}
+
+    weight_bf16 = torch.arange(64, dtype=torch.float32).view(1, 64).to(torch.bfloat16)
+    weight_fp8, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16)
+
+    assert _try_load_fp8_indexer_wk(
+        f"{prefix}.wk.weight",
+        weight_fp8,
+        pending,
+        params_dict,
+        loaded_params,
+        [],
+    )
+    assert pending
+    assert _try_load_fp8_indexer_wk(
+        f"{prefix}.wk.weight_scale",
+        weight_scale,
+        pending,
+        params_dict,
+        loaded_params,
+        [],
+    )
+
+    assert pending == {}
+    assert loaded_params == {f"{prefix}.wk_weights_proj.weight"}
+    assert param.loaded_shard_id == 0
+    torch.testing.assert_close(
+        param.loaded_weight,
+        dequant_mxfp8_to_bf16(weight_fp8, weight_scale),
+    )
+
+
+def test_try_load_fp8_indexer_wk_preserves_fp8_weight_scale_inv_path():
+    prefix = "layers.0.self_attn.indexer"
+    param = _LoadedParam()
+    params_dict = {f"{prefix}.wk_weights_proj.weight": param}
+    loaded_params: set[str] = set()
+    pending = {}
+
+    weight_fp8 = torch.linspace(-2.0, 2.0, 32 * 64, dtype=torch.float32).view(
+        32, 64
+    ).to(torch.float8_e4m3fn)
+    scale_inv = torch.full((1, 2), 0.25, dtype=torch.float32)
+
+    assert _try_load_fp8_indexer_wk(
+        f"{prefix}.wk.weight_scale_inv",
+        scale_inv,
+        pending,
+        params_dict,
+        loaded_params,
+        [],
+    )
+    assert pending
+    assert _try_load_fp8_indexer_wk(
+        f"{prefix}.wk.weight",
+        weight_fp8,
+        pending,
+        params_dict,
+        loaded_params,
+        [],
+    )
+
+    assert pending == {}
+    assert loaded_params == {f"{prefix}.wk_weights_proj.weight"}
+    assert param.loaded_shard_id == 0
+    torch.testing.assert_close(
+        param.loaded_weight,
+        scaled_dequantize(
+            weight_fp8,
+            scale_inv,
+            group_shape=GroupShape(32, 32),
+            out_dtype=torch.bfloat16,
+        ),
+    )
+
+
+def test_try_load_fp8_indexer_wk_ignores_unrelated_mxfp8_scale():
+    assert not _try_load_fp8_indexer_wk(
+        "layers.0.self_attn.indexer.wq_b.weight_scale",
+        torch.ones((1, 2), dtype=torch.uint8),
+        {},
+        {},
+        set(),
+        [],
+    )

vllm/model_executor/models/deepseek_v2.py

@@ -822,19 +822,37 @@ def _should_skip_index_topk(
 def _try_load_fp8_indexer_wk(
     name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names
 ):
-    """
-    We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
-    in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
-    Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
-    weights and scale, and when both are available, dequantizes to BF16 and stores into
-    the fused wk_weights_proj.weight parameter.
+    """Load isolated FP8/MXFP8 indexer WK tensors into fused WK weights.
+
+    The model fuses WK and weights_proj projections, but some checkpoints store
+    WK separately in FP8 with ``weight_scale_inv`` or MXFP8 with
+    ``weight_scale`` while weights_proj stays BF16. This loader buffers the
+    isolated WK weight and scale until both are available, dequantizes WK to
+    BF16, and stores it into the fused ``wk_weights_proj.weight`` parameter.
+
+    Args:
+        name: Checkpoint tensor name.
+        tensor: Checkpoint tensor value.
+        buf: Pending WK weight/scale tensors keyed by layer prefix.
+        params_dict: Model parameters keyed by checkpoint tensor name.
+        loaded_params: Names of parameters already loaded by a special loader.
+        pp_missing_layer_names: Pipeline-parallel layer prefixes to skip.
+
+    Returns:
+        ``True`` when the tensor was consumed or intentionally skipped by this
+        special loader, otherwise ``False``.
+
+    Raises:
+        KeyError: If a matching isolated WK pair is ready but the fused
+            ``wk_weights_proj.weight`` parameter is missing.
     """
     if "indexer.wk." not in name or "wk_weights" in name:
         return False  # Weight is not an isolated WK weight for the indexer, ignore.
     is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
-    is_scale = "weight_scale_inv" in name
-    if not is_weight and not is_scale:
-        return False  # WK is not in FP8 format, ignore.
+    is_scale_inv = "weight_scale_inv" in name
+    is_mxfp8_scale = name.endswith(".weight_scale") and tensor.dtype == torch.uint8
+    if not is_weight and not is_scale_inv and not is_mxfp8_scale:
+        return False  # WK is not in a fused FP8/MXFP8 format, ignore.
     # Buffer this tensor (weight or scale) until both have arrived.
     layer_prefix = name.rsplit(".wk.", 1)[0]  # e.g. "model.layers.0.self_attn.indexer"
     fused_name = f"{layer_prefix}.wk_weights_proj.weight"
@@ -844,20 +862,35 @@ def _try_load_fp8_indexer_wk(
     ):
         return True
     entry = buf.setdefault(layer_prefix, {})
-    entry["weight" if is_weight else "scale"] = tensor
-    if "weight" not in entry or "scale" not in entry:
+    if is_weight:
+        entry["weight"] = tensor
+    elif is_scale_inv:
+        entry["scale_inv"] = tensor
+    else:
+        entry["mxfp8_scale"] = tensor
+    if "weight" not in entry or (
+        "scale_inv" not in entry and "mxfp8_scale" not in entry
+    ):
         return True  # still waiting for the other param
 
     # We have both weight and scale: dequantize FP8 to BF16.
-    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    weight_fp8 = entry["weight"]
     del buf[layer_prefix]
-    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
-    weight_bf16 = scaled_dequantize(
-        weight_fp8,
-        scale_inv,
-        group_shape=GroupShape(block_size, block_size),
-        out_dtype=torch.bfloat16,
-    )
+    if "scale_inv" in entry:
+        scale_inv = entry["scale_inv"]
+        block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+        weight_bf16 = scaled_dequantize(
+            weight_fp8,
+            scale_inv,
+            group_shape=GroupShape(block_size, block_size),
+            out_dtype=torch.bfloat16,
+        )
+    else:
+        from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+            dequant_mxfp8_to_bf16,
+        )
+
+        weight_bf16 = dequant_mxfp8_to_bf16(weight_fp8, entry["mxfp8_scale"])
 
     # Load the dequantized weight into shard 0 of the fused buffer.
     param = params_dict[fused_name]