[Model] Support DeepSeek-V4.1-Flash (#56214)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Signed-off-by: andyluo7 <andy.luo@amd.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

tests/fusion/test_quant_activation_contract.py

@@ -9,8 +9,13 @@ from vllm.model_executor.kernels.linear import (
     _POSSIBLE_FP8_BLOCK_KERNELS,
     _POSSIBLE_FP8_KERNELS,
     _POSSIBLE_INT8_KERNELS,
+    _POSSIBLE_MXFP8_KERNELS,
     _POSSIBLE_NVFP4_KERNELS,
 )
+from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+    FlashInferCutedslMxfp8LinearKernel,
+    FlashInferCutlassMxfp8LinearKernel,
+)
 from vllm.model_executor.kernels.linear.nvfp4.base import (
     NvFp4LinearKernel,
     NvFp4LinearLayerConfig,
@@ -50,6 +55,8 @@ SUPPORTING = {
     FlashInferFP8ScaledMMLinearKernel,
     FlashInferCutlassNvFp4LinearKernel,
     PerTensorTorchFP8ScaledMMLinearKernel,
+    FlashInferCutedslMxfp8LinearKernel,
+    FlashInferCutlassMxfp8LinearKernel,
 }
 
 
@@ -60,6 +67,7 @@ def _all_kernel_classes() -> list[type]:
         _POSSIBLE_FP8_BLOCK_KERNELS,
         _POSSIBLE_INT8_KERNELS,
         _POSSIBLE_NVFP4_KERNELS,
+        _POSSIBLE_MXFP8_KERNELS,
     ):
         for kernels in registry.values():
             for cls in kernels:

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -1268,6 +1268,31 @@ def test_get_cached_wo_a_bf16_plain_caches() -> None:
     torch.testing.assert_close(out2, expected, atol=0, rtol=0)
 
 
+@pytest.mark.parametrize("scale_attr", ["weight_scale", "weight_scale_inv"])
+@torch.inference_mode()
+def test_get_cached_wo_a_bf16_dequantized_ignores_retained_scale(
+    scale_attr: str,
+) -> None:
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
+
+    device = torch.device("cuda")
+    torch.manual_seed(6)
+    n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
+    weight = torch.randn(
+        n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
+    )
+    retained_scale = torch.full(
+        (n_local_groups, 1), 0.25, dtype=torch.float32, device=device
+    )
+    wo_a = _FakeWoA(weight)
+    setattr(wo_a, scale_attr, retained_scale)
+
+    out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
+
+    expected = weight.view(n_local_groups, o_lora_rank, hidden_dim)
+    torch.testing.assert_close(out, expected, atol=0, rtol=0)
+
+
 @torch.inference_mode()
 def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
     from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16

tests/kernels/core/test_fused_q_kv_rmsnorm.py

@@ -112,3 +112,192 @@ def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
             rtol=1e-2,
             atol=1e-2,
         )
+
+
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
+)
+@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 129, 1024])
+@pytest.mark.parametrize("q_size", [384, 1280, 1344])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype):
+    """Preserve native MXFP8 rounding and all swizzled scale padding bytes."""
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+        fused_q_kv_rmsnorm_quant,
+    )
+
+    torch.manual_seed(0)
+    kv_size = 512
+    x = torch.randn(num_tokens, q_size + kv_size, device="cuda", dtype=dtype)
+    x[:1] = 0
+    x[1:2] *= 1e-30
+    qr, kv = x.split([q_size, kv_size], -1)
+    qw = torch.randn(q_size, device="cuda", dtype=x.dtype)
+    kvw = torch.randn(kv_size, device="cuda", dtype=x.dtype)
+    eps = 1e-20
+    result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
+    assert result[0].orig_dtype == dtype
+    assert result[0].orig_shape == qr.shape
+    if not num_tokens:
+        assert result[0].data.shape == qr.shape
+        assert result[0].scale.numel() == 0
+        assert result[1].shape == kv.shape
+        return
+    qr_ref, kv_ref = fused_q_kv_rmsnorm(qr, kv, qw, kvw, eps)
+    q_ref, scale_ref = mxfp8_e4m3_quantize(qr_ref, is_sf_swizzled_layout=True)
+    for replay in (False, True):
+        if replay:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                result = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, eps)
+            graph.replay()
+        torch.testing.assert_close(
+            result[0].data.view(torch.uint8), q_ref.view(torch.uint8), rtol=0, atol=0
+        )
+        torch.testing.assert_close(result[0].scale, scale_ref, rtol=0, atol=0)
+        torch.testing.assert_close(result[1], kv_ref, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(100), reason="MXFP8 needs Blackwell"
+)
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass"])
+@pytest.mark.parametrize("num_tokens", [1, 17, 256])
+@pytest.mark.parametrize("projection", ["attention", "indexer"])
+@pytest.mark.parametrize("bias", [False, True])
+def test_shared_query_quant_preserves_projection(
+    num_tokens, backend, projection, bias, monkeypatch
+):
+    """The native-checkpoint projection keeps identical GEMM outputs."""
+    from dataclasses import replace
+    from types import SimpleNamespace
+    from typing import Any, cast
+
+    from vllm.config import (
+        CompilationConfig,
+        CompilationMode,
+        VllmConfig,
+        set_current_vllm_config,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8 import Mxfp8LinearLayerConfig
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutedslMxfp8LinearKernel,
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+    from vllm.model_executor.layers.fusion.quant_activation import (
+        expose_input_quant_key,
+    )
+    from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
+    from vllm.model_executor.layers.quantization.modelopt import (
+        ModelOptLinearMethod,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        mxfp8_e4m3_quantize,
+    )
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
+    from vllm.models.deepseek_v4_1.attention import (
+        DeepseekV4Attention,
+        DeepseekV4Indexer,
+    )
+    from vllm.models.deepseek_v4_1.common.ops.query_quant import (
+        can_fuse_query_quant,
+        fused_q_kv_rmsnorm_quant,
+    )
+    from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config
+
+    cls: (
+        type[FlashInferCutedslMxfp8LinearKernel]
+        | type[FlashInferCutlassMxfp8LinearKernel]
+    )
+    if backend == "cute-dsl":
+        cls = FlashInferCutedslMxfp8LinearKernel
+    else:
+        cls = FlashInferCutlassMxfp8LinearKernel
+    supported, reason = cls.is_supported()
+    if not supported:
+        pytest.skip(reason)
+    torch.manual_seed(0)
+    with set_current_vllm_config(
+        VllmConfig(compilation_config=CompilationConfig(mode=CompilationMode.NONE))
+    ):
+        kernel = cls(Mxfp8LinearLayerConfig())
+        # Match the checkpoint's TP4 attention and replicated indexer widths.
+        if projection == "attention":
+            linear = ColumnParallelLinear.__new__(ColumnParallelLinear)
+            output_size = 8192
+        else:
+            linear = ReplicatedLinear.__new__(ReplicatedLinear)
+            output_size = 4096
+        torch.nn.Module.__init__(linear)
+        linear.bias = (
+            torch.nn.Parameter(
+                torch.randn(output_size, device="cuda", dtype=torch.bfloat16)
+            )
+            if bias
+            else None
+        )
+        linear.skip_bias_add = False
+        linear.return_bias = projection == "indexer"
+        linear.gather_output = False
+        w, ws = mxfp8_e4m3_quantize(
+            torch.randn(output_size, 1280, device="cuda", dtype=torch.bfloat16)
+        )
+        linear.weight = torch.nn.Parameter(w, requires_grad=False)
+        linear.weight_scale = torch.nn.Parameter(ws, requires_grad=False)
+        quant_config = DeepseekV4FP8Config.from_config(
+            {
+                "quant_method": "fp8",
+                "activation_scheme": "dynamic",
+                "weight_block_size": [32, 32],
+                "scale_fmt": "ue8m0",
+                "expert_dtype": "fp4",
+            }
+        )
+        method = quant_config.get_quant_method(linear, "model.layers.2.attn.wq_b")
+        assert type(method) is ModelOptLinearMethod
+        method.out_dtype = torch.bfloat16
+        method.kernel = kernel
+        linear.quant_method = method
+        kernel.process_weights_after_loading(linear)
+        expose_input_quant_key(linear, kernel)
+        assert can_fuse_query_quant([linear, linear])
+        x = torch.randn(num_tokens, 1792, device="cuda", dtype=torch.bfloat16)
+        qr, kv = x.split([1280, 512], -1)
+        qw = torch.randn(1280, device="cuda", dtype=x.dtype)
+        kvw = torch.randn(512, device="cuda", dtype=x.dtype)
+        qr_ref, _ = fused_q_kv_rmsnorm(qr, kv, qw, kvw, 1e-20)
+        q, _ = fused_q_kv_rmsnorm_quant(qr, kv, qw, kvw, 1e-20)
+        owner = cast(Any, SimpleNamespace(wq_b=linear))
+        project = (
+            DeepseekV4Attention._wq_b_proj
+            if projection == "attention"
+            else DeepseekV4Indexer._wq_b_proj
+        )
+        expected = project(owner, qr_ref)
+
+        def reject_requantization(*args, **kwargs):
+            raise AssertionError("pre-quantized input must skip quantization")
+
+        monkeypatch.setattr(
+            "vllm.model_executor.kernels.linear.mxfp8.flashinfer.mxfp8_e4m3_quantize",
+            reject_requantization,
+        )
+        actual = project(owner, q)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = project(owner, q)
+        graph.replay()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        assert not can_fuse_query_quant(
+            [linear, torch.nn.Linear(1280, 256, bias=False)]
+        )
+        reshaped = project(
+            owner, replace(q, orig_shape=torch.Size((1, num_tokens, 1280)))
+        )
+        torch.testing.assert_close(reshaped, expected.unsqueeze(0), rtol=0, atol=0)
+        with pytest.raises(AssertionError, match="consumer kernel"):
+            linear(replace(q, quant_key=kNvfp4Dynamic))

tests/kernels/quantization/test_rocm_mxfp8_linear.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native CDNA4 MXFP8 linear (``RocmDotScaledMxfp8LinearKernel``).
+
+DeepSeek V4.1 routes every attention linear through this kernel. It used to
+carry two BF16-dequantize workarounds on ROCm (one for all ``.attn.`` linears,
+one for the row-parallel ``wo_b``, whose comment blamed NaNs at small token
+counts). These tests pin the two properties those workarounds stood in for --
+finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at
+TP=4 and the small token counts decode actually runs.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    _mxfp8_e4m3_quantize_torch,
+    _mxfp8_e4m3_quantize_triton,
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm() or not current_platform.supports_mx(),
+    reason="native MX linear requires CDNA4 (gfx95x)",
+)
+
+# (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4:
+# fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is
+# row-parallel (so K, not N, is sharded).
+V41_ATTN_SHAPES = [
+    ("fused_wqa_wkv", 1280 + 512, 5120),
+    ("wq_b", 64 * 512 // 4, 1280),
+    ("wo_a", 8 * 1024 // 4, 4096),
+    ("wo_b", 5120, 8192 // 4),
+]
+
+
+def _make_layer(n: int, k: int, device: str) -> torch.nn.Module:
+    weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05
+    weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16)
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+    return layer
+
+
+@pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v))
+@pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024])
+@torch.inference_mode()
+def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens):
+    """Output is finite and matches an exact MXFP8 GEMM of the same inputs.
+
+    The reference quantizes the activation with the same quantizer and then
+    contracts in FP32, so the tolerance covers only kernel/accumulation error,
+    not the MXFP8 quantization error both sides share.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        RocmDotScaledMxfp8LinearKernel,
+    )
+
+    torch.manual_seed(num_tokens)
+    device = "cuda"
+    layer = _make_layer(n, k, device)
+    kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
+    kernel.process_weights_after_loading(layer)
+
+    x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5
+    out = kernel.apply_weights(layer, x)
+
+    assert out.shape == (num_tokens, n)
+    assert torch.isfinite(out).all(), f"{name}: non-finite output"
+
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    expected = (
+        dequant_mxfp8_to_bf16(x_q, x_scale).float()
+        @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T
+    )
+    rel = (out.float() - expected).norm() / expected.norm()
+    assert rel < 5e-3, f"{name}: relative error {rel:.4f}"
+
+
+@torch.inference_mode()
+def test_rocm_mxfp8_linear_handles_extreme_magnitudes():
+    """Rows that are all-zero or near the E4M3 range must not produce NaNs.
+
+    A degenerate block makes the dynamic activation quantizer pick an extreme
+    E8M0 scale; a zero row makes it pick the tiny-clamped one. Both feed
+    ``tl.dot_scaled`` directly.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        RocmDotScaledMxfp8LinearKernel,
+    )
+
+    torch.manual_seed(0)
+    device = "cuda"
+    n, k = 5120, 2048  # wo_b at TP=4
+    layer = _make_layer(n, k, device)
+    kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
+    kernel.process_weights_after_loading(layer)
+
+    x = torch.randn(8, k, device=device, dtype=torch.bfloat16)
+    x[0] = 0.0
+    x[1] = 1e-30
+    x[2] = 400.0
+    x[3, ::2] = 0.0
+
+    out = kernel.apply_weights(layer, x)
+    assert torch.isfinite(out).all()
+    assert (out[0] == 0).all(), "an all-zero activation row must give a zero row"
+
+
+@pytest.mark.parametrize("num_rows", [1, 64, 65])
+@torch.inference_mode()
+def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
+    """The fused ROCm quantizer must agree with the torch one, zeros included.
+
+    An all-zero 32-element block clamps the E8M0 scale to 0, i.e. 2**-127 --
+    subnormal in fp32 and flushed to zero on CDNA. Dividing by it turned the
+    block into 0/0 == NaN, which then poisoned every output column of the GEMM
+    that consumed it.
+    """
+    torch.manual_seed(num_rows)
+    device = "cuda"
+    x = torch.randn(num_rows, 128, device=device, dtype=torch.bfloat16)
+    x[0] = 0.0  # all-zero row: every block clamps to scale byte 0
+    x[-1, :32] = 0.0  # single zero block in an otherwise normal row
+    x[-1, 32:64] = 1e-38  # denormal-ish block
+
+    q_triton, s_triton = _mxfp8_e4m3_quantize_triton(x)
+    q_torch, s_torch = _mxfp8_e4m3_quantize_torch(x)
+
+    assert not torch.isnan(q_triton.float()).any()
+    torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+    torch.testing.assert_close(
+        q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0
+    )

tests/kernels/test_compressor_kv_cache.py

@@ -42,6 +42,438 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="graph capture coverage"
+)
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize(
+    "lengths",
+    [(1, 5, 13), (1, 1, 1), (1, 2, 0), (7, 8, 8), (2, 3, 3)],
+    ids=["prefill", "decode", "empty", "wrap_threshold", "packed_boundary"],
+)
+def test_v41_fused_save_compress_and_insert(
+    compress_ratio: int, use_graph: bool, lengths: tuple[int, ...]
+):
+    """Ring states: a group boundary whose previous token is outside the chunk
+    pools the ring row that chunk left behind, even when this chunk overwrites
+    that row; only the last ``capacity`` tokens of a chunk are stored.
+
+    Exercise historical group starts, adjacent requests, a chunk longer than
+    the ring, single-token decode pairs, empty requests, wrap thresholds at
+    odd and even starts, padded graph rows, and padded physical pages against
+    a Torch reference with BF16 rounding before RoPE. Replays change both
+    inputs and history so stale scratch/cache values cannot satisfy the test.
+    """
+    from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
+        fused_save_compress_norm,
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(42)
+    device = "cuda"
+    capacity = 8
+    starts = [7, 8, 13]
+    actual = sum(lengths)
+    num_tokens = actual + 5
+    positions = torch.zeros(num_tokens, dtype=torch.int64, device=device)
+    req_ids = torch.zeros(actual + 2, dtype=torch.int32, device=device)
+    state_slots = torch.full((actual + 2,), -1, dtype=torch.int64, device=device)
+    cache_slots = torch.full_like(state_slots, -1)
+    query_start_loc = torch.tensor(
+        [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32
+    ).to(device)
+    ring_blocks = torch.randperm(6, device=device)[:3]
+    raw = torch.randn(num_tokens, 512 * compress_ratio, device=device)
+    norm = torch.randn(512, dtype=torch.bfloat16, device=device)
+    latent = torch.empty(num_tokens, 512, dtype=torch.bfloat16, device=device)
+    angles = torch.randn(32, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    # capacity x 1024 floats rounded up to 576-byte alignment.
+    state_backing = torch.empty(6, 8208, device=device)
+    state = state_backing.as_strided((6, capacity, 1024), (8208, 1024, 1))
+    cache_block = 256 // compress_ratio
+    cache_stride = math.ceil(cache_block * 584 / 576) * 576
+    cache_backing = torch.empty(3, cache_stride, dtype=torch.uint8, device=device)
+    cache = cache_backing.as_strided((3, cache_block, 584), (cache_stride, 584, 1))
+    cursor = 0
+    for req, (start, length) in enumerate(zip(starts, lengths)):
+        pos = torch.arange(start, start + length, device=device)
+        rows = slice(cursor, cursor + length)
+        positions[rows] = pos
+        req_ids[rows] = req
+        state_slots[rows] = ring_blocks[req] * capacity + pos % capacity
+        cache_slots[rows] = (2 - req) * cache_block + pos // compress_ratio
+        cursor += length
+
+    has_ring = compress_ratio == 2
+
+    def run():
+        fused_save_compress_norm(
+            raw,
+            positions,
+            state if has_ring else None,
+            state_slots,
+            query_start_loc if has_ring else None,
+            req_ids if has_ring else None,
+            norm,
+            1e-20,
+            compress_ratio,
+            latent,
+        )
+        rope_quant_insert(
+            latent, positions, cos_sin, cache, cache_slots, compress_ratio
+        )
+
+    state_backing.normal_()
+    run()
+    graph = None
+    if use_graph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run()
+
+    for _ in range(3):
+        raw.normal_()
+        state_backing.normal_()
+        history = state.clone()
+        expected_state = state_backing.clone()
+        expected_rows = expected_state.as_strided((6, capacity, 1024), (8208, 1024, 1))
+        cache_backing.fill_(165)
+        expected_cache = cache_backing.clone()
+        latent.fill_(float("nan"))
+        # Ratio 1 has no ring; ratio 2 stores each chunk's last `capacity` rows.
+        for t in range(actual if has_ring else 0):
+            req = req_ids[t].item()
+            if query_start_loc[req + 1].item() - t > capacity:
+                continue
+            slot = state_slots[t].item()
+            expected_rows[slot // capacity, slot % capacity, :512] = raw[t, :512]
+            expected_rows[slot // capacity, slot % capacity, 512:] = raw[t, 512:]
+
+        if graph is None:
+            run()
+        else:
+            graph.replay()
+
+        torch.testing.assert_close(state_backing, expected_state, rtol=0, atol=0)
+        for t in range(num_tokens):
+            pos = positions[t].item()
+            if t >= actual or (pos + 1) % compress_ratio:
+                assert torch.isnan(latent[t]).all()
+                continue
+            req = req_ids[t].item()
+            group_rows = []
+            for k in range(compress_ratio - 1, -1, -1):
+                if t - k >= query_start_loc[req].item():
+                    kv_score = raw[t - k]
+                    if compress_ratio == 1:
+                        kv_score = torch.cat((kv_score, torch.zeros_like(kv_score)))
+                    group_rows.append(kv_score)
+                else:
+                    ring = ring_blocks[req].item()
+                    group_rows.append(history[ring, (pos - k) % capacity])
+            group = torch.stack(group_rows)
+            pooled = (group[:, :512] * group[:, 512:].softmax(0)).sum(0)
+            normed = pooled * torch.rsqrt(pooled.square().mean() + 1e-20) * norm
+            torch.testing.assert_close(latent[t].float(), normed, rtol=0.004, atol=1e-6)
+
+            slot = cache_slots[t].item()
+            page, row = divmod(slot, cache_block)
+            values = expected_cache[page, row * 576 : (row + 1) * 576]
+            quantized, scales = _ue8m0_reference(latent[t, :448], 64, 448.0)
+            values[:448] = quantized.view(torch.uint8)
+            c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
+            rope_input = latent[t, 448:].float()
+            rotated = (
+                torch.stack(
+                    (
+                        rope_input[0::2] * c - rope_input[1::2] * s,
+                        rope_input[1::2] * c + rope_input[0::2] * s,
+                    ),
+                    dim=-1,
+                )
+                .flatten()
+                .to(torch.bfloat16)
+            )
+            actual_rope = cache_backing[page, row * 576 + 448 : (row + 1) * 576].view(
+                torch.bfloat16
+            )
+            torch.testing.assert_close(actual_rope, rotated, rtol=0.008, atol=1e-6)
+            values[448:] = actual_rope.view(torch.uint8)
+            scale_offset = cache_block * 576 + row * 8
+            expected_cache[page, scale_offset : scale_offset + 7] = (
+                scales.log2() + 127
+            ).to(torch.uint8)
+            expected_cache[page, scale_offset + 7] = 0
+        # Includes unwritten rows, page padding, quantized values and scales.
+        torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("store_fp8", [False, True])
+def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool):
+    """The FlashInfer plain-row compressed cache gets [448 NoPE | 64 RoPE] rows.
+
+    bf16 caches store the latent verbatim with the RoPE tail rotated; per-tensor
+    fp8 caches scale the bf16-rounded row by 1/fp8_scale and clamp to e4m3.
+    Non-boundary positions, negative slots and untouched rows must keep their
+    prior contents.
+    """
+    from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import (
+        rope_quant_insert,
+    )
+
+    torch.manual_seed(3)
+    device = "cuda"
+    num_tokens = 9
+    cache_block = 4
+    positions = torch.arange(5, 5 + num_tokens, dtype=torch.int64, device=device)
+    cache_slots = torch.randperm(3 * cache_block, device=device)[:num_tokens]
+    cache_slots[2] = -1
+    latent = torch.randn(num_tokens, 512, device=device).to(torch.bfloat16)
+    angles = torch.randn(32, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
+    dtype = torch.float8_e4m3fn if store_fp8 else torch.bfloat16
+    # Pad the page stride so the kernel must honour strides, not shape.
+    cache_backing = torch.full((3, (cache_block + 1) * 512), 3.0, device=device)
+    cache_backing = cache_backing.to(dtype)
+    cache = cache_backing.as_strided(
+        (3, cache_block, 512), ((cache_block + 1) * 512, 512, 1)
+    )
+    fp8_scale = torch.tensor([0.5], dtype=torch.float32, device=device)
+    expected = cache_backing.clone().float()
+
+    rope_quant_insert(
+        latent,
+        positions,
+        cos_sin,
+        cache,
+        cache_slots,
+        compress_ratio,
+        fp8_scale=fp8_scale if store_fp8 else None,
+    )
+
+    for t in range(num_tokens):
+        pos = positions[t].item()
+        slot = cache_slots[t].item()
+        if slot < 0 or (pos + 1) % compress_ratio:
+            continue
+        c, s = cos_sin[pos // compress_ratio * compress_ratio].chunk(2)
+        rope_input = latent[t, 448:].float()
+        rotated = torch.stack(
+            (
+                rope_input[0::2] * c - rope_input[1::2] * s,
+                rope_input[1::2] * c + rope_input[0::2] * s,
+            ),
+            dim=-1,
+        ).flatten()
+        row = torch.cat((latent[t, :448], rotated.to(torch.bfloat16))).float()
+        if store_fp8:
+            row = (row * (1.0 / fp8_scale)).clamp(-448.0, 448.0)
+            row = row.to(torch.float8_e4m3fn).float()
+        page, idx = divmod(slot, cache_block)
+        expected[page, idx * 512 : (idx + 1) * 512] = row
+
+    actual = cache_backing.float()
+    if store_fp8:
+        # One e4m3 ulp of slack for FMA-vs-separate rounding in the RoPE tail.
+        torch.testing.assert_close(actual, expected, rtol=0.13, atol=1e-2)
+    else:
+        nope = torch.arange(512, device=device) < 448
+        nope = nope.repeat(cache_block + 1)
+        torch.testing.assert_close(actual[:, nope], expected[:, nope], rtol=0, atol=0)
+        torch.testing.assert_close(
+            actual[:, ~nope], expected[:, ~nope], rtol=0.008, atol=1e-6
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel")
+def test_v41_compressor_metadata_maps_tokens_to_their_ring():
+    """The ring group's generic slot mapping is disabled (all PAD), so the
+    builder must map every real token to ``ring_block * capacity + pos %
+    capacity`` and keep padding tokens at PAD."""
+    from unittest.mock import MagicMock
+
+    from vllm.models.deepseek_v4_1.compressor import CompressorMetadataBuilder
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.kv_cache_interface import CircularBufferSpec
+
+    capacity = 8
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_batched_tokens = 16
+    spec = CircularBufferSpec(
+        block_size=capacity,
+        num_kv_heads=1,
+        head_size=1024,
+        head_size_v=0,
+        dtype=torch.float32,
+    )
+    device = torch.device("cuda")
+    builder = CompressorMetadataBuilder(spec, ["state"], vllm_config, device)
+
+    # Two requests: 3 tokens at positions 13..15 on ring block 5, then 2
+    # tokens at positions 7..8 on ring block 2; three padding tokens.
+    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
+    positions = torch.tensor([13, 14, 15, 7, 8, 0, 0, 0], device=device)
+    block_table = torch.tensor([[5], [2]], dtype=torch.int32, device=device)
+    common = CommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc.cpu(),
+        seq_lens=torch.tensor([16, 9], dtype=torch.int32, device=device),
+        num_reqs=2,
+        num_actual_tokens=5,
+        max_query_len=3,
+        max_seq_len=16,
+        block_table_tensor=block_table,
+        slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device),
+        positions=positions,
+    )
+    metadata = builder.build(0, common)
+
+    expected = [5 * 8 + 5, 5 * 8 + 6, 5 * 8 + 7, 2 * 8 + 7, 2 * 8 + 0, -1, -1, -1]
+    assert metadata.slot_mapping.tolist() == expected
+    assert metadata.query_start_loc is query_start_loc
+    assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage")
+@pytest.mark.parametrize(
+    "use_aux,use_graph", [(False, False), (True, False), (True, True)]
+)
+@torch.inference_mode()
+def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph):
+    """Both reused-event joins must publish this forward's states and cache rows."""
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.deepseek_v4_1.attention import (
+        DeepseekV4Attention,
+        DeepseekV4Indexer,
+    )
+    from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor
+
+    torch.manual_seed(43)
+    raw = torch.randn(19, 1024, device="cuda")
+    q = torch.zeros(19, 512, dtype=torch.bfloat16, device="cuda")
+    positions = torch.arange(7, 26, device="cuda")
+    state_slots = positions.clone()
+    state_slots[-2:] = -1
+    cache_slots = torch.where(state_slots >= 0, positions // 2, -1)
+    state = torch.empty(4, 8, 1024, device="cuda")
+    main = torch.empty(1, 128, 584, dtype=torch.uint8, device="cuda")
+    index = torch.empty(1, 128, 132, dtype=torch.uint8, device="cuda")
+    caches = (state, main, index)
+    observed = tuple(torch.empty_like(cache) for cache in caches)
+    rotary = SimpleNamespace(cos_sin_cache=torch.randn(32, 64, device="cuda"))
+    metadata = {
+        "state": SimpleNamespace(
+            slot_mapping=state_slots,
+            query_start_loc=torch.tensor([0, 19], dtype=torch.int32, device="cuda"),
+            token_to_req_indices=torch.zeros(19, dtype=torch.int32, device="cuda"),
+        ),
+        "main": SimpleNamespace(slot_mapping=cache_slots),
+        "index": SimpleNamespace(slot_mapping=cache_slots),
+    }
+    context = ForwardContext({}, metadata, {})
+    compressor = DeepseekCompressor.__new__(DeepseekCompressor)
+    torch.nn.Module.__init__(compressor)
+    compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = (
+        512,
+        64,
+        2,
+    )
+    compressor.rms_norm_eps = 1e-20
+    compressor.norm = SimpleNamespace(
+        weight=torch.ones(512, dtype=torch.bfloat16, device="cuda")
+    )
+    compressor.state_cache = SimpleNamespace(prefix="state", kv_cache=state)
+    compressor.k_cache_prefix = "main"
+    compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=main)}
+    indexer_weight = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda")
+    indexer = SimpleNamespace(
+        owns_k=True,
+        wk=lambda latent: (torch.nn.functional.linear(latent, indexer_weight), None),
+        k_norm=SimpleNamespace(
+            weight=torch.ones(128, dtype=torch.bfloat16, device="cuda"),
+            variance_epsilon=1e-20,
+        ),
+        k_cache=SimpleNamespace(prefix="index", kv_cache=index),
+        compress_ratio=2,
+        use_fp4_kv=False,
+    )
+
+    def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale):
+        DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary)
+        return None, None, None
+
+    def observe(*args):
+        for output, cache in zip(observed, caches):
+            output.copy_(cache)
+
+    auxiliary = [torch.cuda.Stream()] if use_aux else None
+    attention = SimpleNamespace(
+        compressor=compressor,
+        indexer=prepare_indexer,
+        aux_stream_list=auxiliary,
+        ln_events=[torch.cuda.Event(), torch.cuda.Event()],
+        rotary_emb=rotary,
+        indexer_rotary_emb=rotary,
+        n_local_heads=1,
+        head_dim=512,
+        _wq_b_proj=lambda qr, scale: qr.clone(),
+        _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q,
+        _sparse_indexer_and_attn=observe,
+    )
+
+    def run():
+        with override_forward_context(context):
+            DeepseekV4Attention._prepare_and_attn(
+                attention, q, q, q, None, raw, q, positions, q
+            )
+
+    def reset():
+        state.zero_()
+        main.fill_(165)
+        index.fill_(165)
+
+    reset()
+    run()
+    graph = None
+    if use_graph:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run()
+
+    previous = None
+    for _ in range(3):
+        raw.normal_()
+        reset()
+        attention.aux_stream_list = None
+        run()
+        expected = tuple(output.clone() for output in observed)
+        reset()
+        attention.aux_stream_list = auxiliary
+        if graph is None:
+            run()
+        else:
+            graph.replay()
+        for actual, reference in zip(observed, expected):
+            torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+        if previous is not None:
+            assert all(not torch.equal(a, b) for a, b in zip(expected, previous))
+        previous = expected
+
+
 def _on_gfx950() -> bool:
     if not current_platform.is_rocm():
         return False

tests/kernels/test_engram.py

@@ -0,0 +1,778 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import bisect
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4_1.common import engram as engram_ops
+from vllm.models.deepseek_v4_1.common.engram import (
+    Engram,
+    NgramHashState,
+    ParallelEngramEmbedding,
+)
+from vllm.platforms import current_platform
+
+
+def _reference_engram_post_wkv(
+    hidden_states,
+    kv,
+    q_weight,
+    k_weight,
+    token_mask,
+    eps,
+    clamp_value,
+    kv_start,
+):
+    num_tokens, hc_mult, dim = hidden_states.shape
+    source = torch.arange(num_tokens, device=kv.device) + kv_start
+    valid = source < kv.shape[0]
+    local_kv = kv.new_zeros(num_tokens, kv.shape[1])
+    local_kv[valid] = kv[source[valid]]
+
+    key = local_kv[:, : hc_mult * dim].float().view(num_tokens, hc_mult, dim)
+    value = local_kv[:, hc_mult * dim :].float()
+    h = hidden_states.float()
+    weight = q_weight.float() * k_weight.float()
+    rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(
+        key.square().mean(-1) + eps
+    )
+    dot = (h * weight * key).sum(-1) * rstd * dim**-0.5
+    gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot))
+    if token_mask is not None:
+        local_mask = torch.zeros(num_tokens, dtype=torch.bool, device=kv.device)
+        local_mask[valid] = token_mask[source[valid]]
+        gate = gate.masked_fill(~local_mask.unsqueeze(-1), 0)
+    return (h + gate.unsqueeze(-1) * value.unsqueeze(-2)).to(hidden_states.dtype)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "num_tokens,num_kv_tokens,hc_mult,dim,tp_size,tp_rank,use_mask",
+    [
+        (1, 1, 4, 5120, 1, 0, False),
+        (3, 5, 4, 96, 2, 0, True),
+        (3, 5, 2, 257, 2, 1, False),
+        (3, 5, 4, 96, 2, 1, True),
+        (1, 1, 4, 96, 4, 3, True),
+        (0, 0, 4, 96, 4, 0, False),
+    ],
+)
+def test_fused_engram_post_wkv_matches_reference(
+    num_tokens,
+    num_kv_tokens,
+    hc_mult,
+    dim,
+    tp_size,
+    tp_rank,
+    use_mask,
+    monkeypatch,
+):
+    """Preserve gate math, WKV layout, masking, and SP padding semantics."""
+    torch.manual_seed(0)
+    device = "cuda"
+    hidden_states = torch.randn(
+        num_tokens, hc_mult, dim, dtype=torch.bfloat16, device=device
+    )
+    kv = torch.randn(
+        num_kv_tokens,
+        (hc_mult + 1) * dim,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    q_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
+    k_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device)
+    token_mask = None
+    if use_mask:
+        token_mask = torch.arange(num_kv_tokens, device=device) % 2 == 0
+    eps = 1e-20
+    clamp_value = 1e-6
+    use_sequence_parallel = tp_size > 1
+    shard_size = (num_kv_tokens + tp_size - 1) // tp_size
+    kv_start = tp_rank * shard_size if use_sequence_parallel else 0
+
+    expected = _reference_engram_post_wkv(
+        hidden_states,
+        kv,
+        q_weight,
+        k_weight,
+        token_mask,
+        eps,
+        clamp_value,
+        kv_start,
+    )
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
+
+    module = Engram.__new__(Engram)
+    torch.nn.Module.__init__(module)
+    module.dim = dim
+    module.hc_mult = hc_mult
+    module.eps = eps
+    module.clamp_value = clamp_value
+    module.use_sequence_parallel = use_sequence_parallel
+    module.embed_tokens = torch.nn.Identity()
+    # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
+    module.embed_tokens.tp_size = 1
+    module.staged_rows = kv.unsqueeze(1)
+    if use_sequence_parallel:
+        padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
+        module.staged_rows = padded[
+            tp_rank * shard_size : (tp_rank + 1) * shard_size
+        ].unsqueeze(1)
+    module.wkv = torch.nn.Identity()
+
+    def check_local_projection(_, args):
+        assert args[0].shape[0] == num_tokens
+
+    module.wkv.register_forward_pre_hook(check_local_projection)
+    module.q_weight = torch.nn.Parameter(q_weight, requires_grad=False)
+    module.k_weight = torch.nn.Parameter(k_weight, requires_grad=False)
+
+    actual = module(hidden_states, kv.unsqueeze(1), token_mask)
+
+    assert actual.shape == hidden_states.shape
+    assert actual.dtype == hidden_states.dtype
+    torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
+
+
+def _hash_state(use_slot_cache: bool) -> NgramHashState:
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.block_size = 64
+    state.use_slot_cache = use_slot_cache
+    state._cache = None
+    state._kv_cache_ref = None
+    state.swa_cache_module = torch.nn.Module()
+    state.swa_cache_module.kv_cache = torch.empty(0)
+    return state
+
+
+@pytest.mark.parametrize("num_blocks", [4, 100])
+def test_engram_cache_follows_kv_cache_binding(num_blocks):
+    """Discard profiling history when the real KV cache replaces its storage."""
+    state = _hash_state(use_slot_cache=True)
+    assert not state.ensure_cache()
+
+    state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
+    assert state.ensure_cache()
+    profiling_cache = state._cache
+    profiling_cache.fill_(7)
+    assert state.ensure_cache()
+    assert state._cache is profiling_cache
+    assert torch.all(state._cache == 7)
+
+    state.swa_cache_module.kv_cache = torch.empty(num_blocks, 64, 1)
+    assert state.ensure_cache()
+    assert state._cache is not profiling_cache
+    assert state._cache.shape == (num_blocks * 64,)
+    assert torch.count_nonzero(state._cache) == 0
+
+    state.swa_cache_module.kv_cache = torch.empty(0)
+    assert not state.ensure_cache()
+    assert state._cache is None
+
+
+def test_engram_without_slot_cache_only_gates_on_kv_binding():
+    """The V2 runner supplies every lookback, so no slot cache is allocated."""
+    state = _hash_state(use_slot_cache=False)
+    assert not state.ensure_cache()
+    state.swa_cache_module.kv_cache = torch.empty(4, 64, 1)
+    assert state.ensure_cache()
+    assert state._cache is None
+
+
+def _reference_hashes(
+    ids,
+    positions,
+    slots,
+    starts,
+    tables,
+    dead,
+    cache,
+    token_map,
+    multipliers,
+    primes,
+    offsets,
+):
+    """Scalar n-gram oracle with persistent, physically addressed token history."""
+    for token, slot, is_dead in zip(ids, slots, dead):
+        if slot >= 0:
+            cache[slot] = -1 if is_dead else token_map[token]
+    result = []
+    for i, position in enumerate(positions):
+        req = min(bisect.bisect_right(starts[1:], i), len(tables) - 1)
+        history, blocked = [], False
+        for shift in range(4):
+            lookback = position - shift
+            p = max(0, min(lookback, len(tables[req]) * 16 - 1))
+            source = cache[tables[req][p // 16] * 16 + p % 16]
+            blocked |= lookback < 0 or source == -1
+            history.append(0 if blocked else source)
+        layers = []
+        for layer, row in enumerate(multipliers):
+            rolling, hashes = history[0] * row[0], []
+            for j, head_primes in enumerate(primes[layer], start=1):
+                rolling ^= history[j] * row[j]
+                for head, prime in enumerate(head_primes):
+                    col = (j - 1) * len(head_primes) + head
+                    hashes.append(rolling % prime + offsets[layer][col])
+            layers.append(hashes)
+        result.append(layers)
+    return torch.tensor(result, dtype=torch.int32)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("query_len", [1, 17, 257])
+@pytest.mark.parametrize("capture", [False, True])
+@pytest.mark.parametrize("strided_slots", [False, True])
+@pytest.mark.parametrize("large_hashes", [False, True])
+def test_engram_hash_cache_replay(query_len, capture, strided_slots, large_hashes):
+    """Preserve history, dead tokens and padding as inputs change across replay."""
+    num_tokens = 2 * query_len + 3
+    blocks = (3 * query_len + 15) // 16
+    tables = [list(range(blocks, 2 * blocks)), list(range(blocks))]
+    starts = [0, query_len, 2 * query_len]
+
+    def tensor(data, dtype=torch.int64):
+        return torch.tensor(data, dtype=dtype, device="cuda")
+
+    input_ids = tensor([0] * num_tokens)
+    positions = torch.zeros_like(input_ids)
+    slots = tensor([-1] * (num_tokens * (2 if strided_slots else 1)))
+    if strided_slots:
+        slots = slots[::2]
+    dead_mask = tensor([False] * num_tokens, torch.bool)
+    cache = tensor([-1] * (2 * blocks * 16), torch.int32)
+    token_map = [i * 3091 if large_hashes else i % 7 for i in range(32)]
+    multipliers = (
+        engram_ops.compute_hash_multipliers((1, 14), 4, 99092).tolist()
+        if large_hashes
+        else [[3, 5, 7, 9], [11, 13, 15, 17]]
+    )
+    primes = (
+        [
+            [
+                [16000057 + 2 * (layer * 24 + j * 8 + h) for h in range(8)]
+                for j in range(3)
+            ]
+            for layer in range(2)
+        ]
+        if large_hashes
+        else [[[19], [23], [29]]] * 2
+    )
+    offsets = []
+    for layer in primes:
+        row, offset = [], 0
+        for heads in layer:
+            for prime in heads:
+                row.append(offset)
+                offset += prime
+        offsets.append(row)
+    state = NgramHashState.__new__(NgramHashState)
+    torch.nn.Module.__init__(state)
+    state.token_map = tensor(token_map, torch.int32)
+    state.multipliers = tensor(multipliers)
+    state.primes = tensor(primes)
+    state.offsets = tensor(offsets)
+    state._cache = cache
+    state.use_slot_cache = True
+    state.pad_id = 0
+    state.block_size = 16
+    lookback = tensor([[-1] * 3] * 2, torch.int32)
+    args = (
+        input_ids,
+        positions,
+        tensor(starts, torch.int32),
+        dead_mask,
+        lookback,
+        torch.zeros_like(lookback, dtype=torch.bool),
+        slots,
+        tensor(tables, torch.int32),
+    )
+    op = state
+    graph = None
+    if capture:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                op(*args)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            output = op(*args)
+        torch.cuda.current_stream().wait_stream(stream)
+    cache.fill_(-1)
+    expected_cache = [-1] * cache.numel()
+
+    for step in range(3):
+        ids = [(i + 3 * step) % 32 for i in range(num_tokens)]
+        pos = list(range(step * query_len, (step + 1) * query_len)) * 2
+        physical_slots = [
+            tables[i // query_len][p // 16] * 16 + p % 16 for i, p in enumerate(pos)
+        ]
+        pos += [0] * 3
+        physical_slots += [-1] * 3
+        # A fully padded replay must leave every cache slot untouched.
+        if step == 2:
+            physical_slots = [-1] * num_tokens
+        dead = [(i + step) % 5 == 0 for i in range(num_tokens)]
+        input_ids.copy_(tensor(ids))
+        positions.copy_(tensor(pos))
+        slots.copy_(tensor(physical_slots))
+        dead_mask.copy_(tensor(dead, torch.bool))
+        if graph is None:
+            output = op(*args)
+        else:
+            graph.replay()
+        expected = _reference_hashes(
+            ids,
+            pos,
+            physical_slots,
+            starts,
+            tables,
+            dead,
+            expected_cache,
+            token_map,
+            multipliers,
+            primes,
+            offsets,
+        )
+        if step != 2:
+            real = 2 * query_len
+            torch.testing.assert_close(
+                output[:real].cpu(), expected[:real], rtol=0, atol=0
+            )
+        assert cache.cpu().tolist() == expected_cache
+
+    empty_output = state(
+        input_ids[:0], positions[:0], args[2], dead_mask[:0], *args[4:]
+    )
+    assert empty_output.shape == (0, 2, len(offsets[0]))
+    assert cache.cpu().tolist() == expected_cache
+
+
+def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False):
+    """Run the op on a batch of (token_ids, start, window) requests; `window`
+    is the runner's [depth] lookback for that request (None = all unknown).
+    Token IDs 14 and 22 stand in for image tokens that break n-grams.
+    Returns one [len(token_ids), ...] hash tensor per request."""
+    token_map, multipliers, primes, offsets, block_size = tables_meta
+    depth = multipliers.shape[1] - 1
+
+    def tensor(data, dtype=torch.int64):
+        return torch.tensor(data, dtype=dtype, device="cuda")
+
+    input_ids, positions, slots, windows, starts = [], [], [], [], [0]
+    for req, (token_ids, start, window) in enumerate(requests):
+        pos = list(range(start, start + len(token_ids)))
+        input_ids += token_ids
+        positions += pos
+        slots += [
+            int(block_table[req, p // block_size]) * block_size + p % block_size
+            for p in pos
+        ]
+        windows.append(window or [-1] * depth)
+        starts.append(starts[-1] + len(token_ids))
+    window = tensor(windows, torch.int32)
+    state = _hash_state(use_slot_cache=cache is not None)
+    state.token_map = token_map
+    state.multipliers = multipliers
+    state.primes = primes
+    state.offsets = offsets
+    state.pad_id = 0
+    state.block_size = block_size
+    state._cache = cache
+    input_tensor = tensor(input_ids, torch.int32)
+    args = (
+        input_tensor,
+        tensor(positions),
+        tensor(starts, torch.int32),
+        (input_tensor == 14) | (input_tensor == 22),
+        window,
+        (window == 14) | (window == 22),
+        tensor(slots) if cache is not None else None,
+        block_table if cache is not None else None,
+    )
+    out = state(*args)
+    if capture:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            out = state(*args)
+        graph.replay()
+    return [out[a:b] for a, b in zip(starts, starts[1:])]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("runner", ["v2_window_only", "v1_prompt_window_plus_cache"])
+@pytest.mark.parametrize("capture", [False, True])
+def test_lookback_window_reproduces_single_instance(runner, capture):
+    """Decoding on a fresh instance whose prompt KV came from elsewhere (P/D,
+    offload) must hash like the instance that processed the whole sequence,
+    with two requests batched and chunks of 1-2 tokens so in-batch, window
+    and (V1) cache lookbacks mix within one step.
+
+    V2 supplies every lookback from its device token history and needs no
+    slot cache. V1 supplies prompt positions only and reads generated
+    positions from the slot cache it fills itself. Without any window the
+    first decode token hashes stale slots (negative control)."""
+    block_size = 4
+    block_table = torch.tensor(
+        [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda"
+    )
+    tables_meta = (
+        torch.arange(50, dtype=torch.int32, device="cuda"),
+        torch.tensor([[3, 5, 7, 9]], device="cuda"),
+        torch.tensor([[[97, 89], [83, 79], [73, 71]]], device="cuda"),
+        torch.tensor([[0, 97, 186, 269, 348, 421]], device="cuda"),
+        block_size,
+    )
+    depth = 3
+    prompts = [[11, 12, 13, 14, 15, 16], [31, 32, 33, 34, 35]]
+    decodes = [[21, 22, 23, 24], [41, 42, 43, 44]]
+    chunk_sizes = [1, 2, 1]
+    v2 = runner.startswith("v2")
+
+    def fresh_cache():
+        return torch.zeros(8 * block_size, dtype=torch.int32, device="cuda")
+
+    reference = _hash_ids(
+        [(p + d, 0, None) for p, d in zip(prompts, decodes)],
+        tables_meta,
+        block_table,
+        cache=fresh_cache(),
+    )
+
+    histories = [list(p) for p in prompts]
+    cache = None if v2 else fresh_cache()
+    consumed = 0
+    for chunk in chunk_sizes:
+        batch = []
+        for history, prompt, decode in zip(histories, prompts, decodes):
+            start = len(history)
+            tokens = decode[consumed : consumed + chunk]
+            window = [
+                history[start - 1 - j] if v2 or start - 1 - j < len(prompt) else -1
+                for j in range(depth)
+            ]
+            batch.append((tokens, start, window))
+        got = _hash_ids(batch, tables_meta, block_table, cache=cache, capture=capture)
+        for req, (tokens, start, _) in enumerate(batch):
+            assert torch.equal(got[req], reference[req][start : start + chunk]), (
+                f"{runner} request {req} chunk at {start}"
+            )
+            histories[req] += tokens
+        consumed += chunk
+
+    unseeded = _hash_ids(
+        [(decodes[0][:1], len(prompts[0]), None)],
+        tables_meta,
+        block_table,
+        cache=fresh_cache(),
+    )
+    assert not torch.equal(unseeded[0][0], reference[0][len(prompts[0])])
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_v2_model_state_gathers_lookback_window():
+    """The window is gathered on device from the runner's token history."""
+    from vllm.models.deepseek_v4_1.nvidia.model_state import DeepseekV41ModelState
+
+    depth, max_num_reqs, max_model_len = 3, 4, 16
+    state = DeepseekV41ModelState.__new__(DeepseekV41ModelState)
+    state.rope_state = None
+    state.supports_mm_inputs = False
+    state.prompt_embeds_state = None
+    state.lookback_token_ids = torch.full(
+        (max_num_reqs, depth), -1, dtype=torch.int32, device="cuda"
+    )
+
+    all_token_ids = torch.zeros(max_num_reqs, max_model_len, dtype=torch.int32)
+    all_token_ids[2, :8] = torch.tensor([11, 12, 13, 14, 15, 16, 21, 22])
+    all_token_ids[0, :5] = torch.arange(1, 6)
+    req_states = SimpleNamespace(
+        num_computed_tokens=SimpleNamespace(
+            gpu=torch.tensor([0, 9, 8, 6], dtype=torch.int32, device="cuda")
+        ),
+        all_token_ids=SimpleNamespace(gpu=all_token_ids.cuda()),
+    )
+    # Batch rows -> request state rows: a request two decode steps in, a fresh
+    # request, and one at the end of a 5-token prompt.
+    input_batch = SimpleNamespace(
+        idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda")
+    )
+    window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"]
+    assert window.cpu().tolist() == [
+        [22, 21, 16],
+        [-1, -1, -1],
+        [0, 0, 0],
+        [-1, -1, -1],
+    ]
+    # Graph capture runs on the dummy inputs and replays read the same buffer.
+    dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"]
+    assert dummy is window and torch.all(dummy == -1)
+
+
+def _reference_lookup(weight, scale_inv, ids, start, end, block=32):
+    """The torch expression the fused kernel replaces."""
+    mask = (ids < start) | (ids >= end)
+    local = (ids - start).masked_fill(mask, 0).long()
+    values = torch.nn.functional.embedding(local, weight)
+    scales = torch.nn.functional.embedding(local, scale_inv)
+    scales = (scales.to(torch.int32) << 23).view(torch.float32)
+    values = values.float().unflatten(-1, (-1, block))
+    values = (values * scales.unsqueeze(-1)).flatten(-2).to(torch.bfloat16)
+    return values.masked_fill(mask.unsqueeze(-1), 0)
+
+
+def _make_embedding(cpu_offload, rows=4096, dim=256, block=32):
+    layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding)
+    torch.nn.Module.__init__(layer)
+    layer.dim, layer.block_size, layer.tp_size = dim, block, 1
+    layer.n_hash_cols = layer.part_n_hash_cols = 24
+    layer.head_start = 0
+    layer.cpu_offload = cpu_offload
+    layer.part_num_embeddings = rows
+    # A window strictly inside the table, so unowned rows are exercised too.
+    layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2
+    layer._views = layer._view_src = None
+    layer._num_sms = torch.cuda.get_device_properties(
+        torch.accelerator.current_device_index()
+    ).multi_processor_count
+    kwargs = (
+        {"device": "cpu", "pin_memory": True} if cpu_offload else {"device": "cuda"}
+    )
+    owned = layer.vocab_end_idx - layer.vocab_start_idx
+    layer.weight = torch.empty(owned, dim, dtype=torch.float8_e4m3fn, **kwargs)
+    layer.weight_scale_inv = torch.empty(
+        owned, dim // block, dtype=torch.uint8, **kwargs
+    )
+    torch.manual_seed(0)
+    layer.weight.copy_((torch.randn(owned, dim) * 4).to(torch.float8_e4m3fn))
+    layer.weight_scale_inv.copy_(
+        torch.randint(120, 134, (owned, dim // block), dtype=torch.uint8)
+    )
+    return layer
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
+    """Keep complete buckets and reconstruct head order, including TP padding."""
+    head_sizes = (17, 19, 23, 29, 31, 37)
+    num_rows, dim = sum(head_sizes), 64
+    torch.manual_seed(0)
+    weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn)
+    scales = torch.randint(120, 134, (num_rows + 7, dim // 32), dtype=torch.uint8)
+    hashes = torch.empty(7, 2, len(head_sizes), dtype=torch.int32, device="cuda")
+    start = 0
+    for head, size in enumerate(head_sizes):
+        hashes[:, :, head].random_(start, start + size)
+        hashes[0, :, head] = start
+        hashes[-1, :, head] = start + size - 1
+        start += size
+    ids = hashes[:, 1]
+    expected = _reference_lookup(weight.cuda(), scales.cuda(), ids, 0, num_rows)
+    monkeypatch.setattr(
+        engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    shards = []
+    layers = []
+    for rank in range(tp_size):
+        monkeypatch.setattr(
+            engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
+        )
+        with torch.device("cuda"):
+            layer = ParallelEngramEmbedding(
+                num_rows + 7, dim, head_sizes, cpu_offload=cpu_offload
+            )
+        layer.weight.weight_loader(layer.weight, weight)
+        layer.weight_scale_inv.weight_loader(
+            layer.weight_scale_inv, scales.view(torch.float8_e8m0fnu)
+        )
+        out = torch.empty(
+            len(ids),
+            layer.part_n_hash_cols,
+            dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        layer.lookup(ids, out)
+        shards.append(out)
+        layers.append(layer)
+    gathered = torch.cat(shards, dim=1)
+    torch.testing.assert_close(gathered[:, : len(head_sizes)], expected, rtol=0, atol=0)
+    assert torch.count_nonzero(gathered[:, len(head_sizes) :]) == 0
+    assert sum(layer.part_num_embeddings for layer in layers) == num_rows
+
+    def gather(local, dim):
+        torch.testing.assert_close(local, shards[0], rtol=0, atol=0)
+        return torch.cat(shards, dim=dim)
+
+    monkeypatch.setattr(engram_ops, "tensor_model_parallel_all_gather", gather)
+    torch.testing.assert_close(layers[0](ids), expected, rtol=0, atol=0)
+    module = Engram.__new__(Engram)
+    torch.nn.Module.__init__(module)
+    module.embed_tokens = layers[0]
+    module.use_sequence_parallel = False
+    module.staged_rows = torch.empty_like(shards[0])
+    module.prepare_embeddings(ids)
+    torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0)
+    # Slicing before head reordering must preserve padded heads and empty owners.
+    module.use_sequence_parallel = True
+    chunk = (len(ids) + tp_size - 1) // tp_size
+    padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size))
+    for rank in range(tp_size):
+        monkeypatch.setattr(
+            engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank
+        )
+        torch.testing.assert_close(
+            module.embed(ids), padded[rank * chunk : (rank + 1) * chunk], rtol=0, atol=0
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize("background", [False, True])
+@pytest.mark.parametrize("num_tokens", [1, 7, 256])
+def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
+    """The fused gather must be bit-exact with the torch dequant path it
+    replaces, from HBM and from pinned host memory alike, and must contribute
+    zeros for rows another TP rank owns."""
+    layer = _make_embedding(cpu_offload)
+    cols, rows = 24, layer.part_num_embeddings
+    ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda")
+    expected = _reference_lookup(
+        layer.weight.cuda(),
+        layer.weight_scale_inv.cuda(),
+        ids,
+        layer.vocab_start_idx,
+        layer.vocab_end_idx,
+    )
+    out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda")
+    layer.lookup(ids, out, background=background)
+    assert torch.equal(out, expected)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("cpu_offload", [False, True])
+@pytest.mark.parametrize("capture", ["eager", "full", "breakable"])
+def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture):
+    """Consume early lookup results after a break, with fresh IDs each replay."""
+    _run_engram_prepared_rows(cpu_offload, capture)
+
+
+def _run_engram_prepared_rows(
+    cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False
+):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    layer = _make_embedding(cpu_offload)
+    cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64)
+    layer.n_hash_cols = cols
+    layer.tp_size = tp_size
+    layer.part_n_hash_cols = (cols + tp_size - 1) // tp_size
+    layer.head_start = rank * layer.part_n_hash_cols
+    engram = Engram.__new__(Engram)
+    torch.nn.Module.__init__(engram)
+    engram.embed_tokens = layer
+    engram.use_sequence_parallel = use_sequence_parallel
+    engram.staged_rows = torch.empty(
+        num_tokens,
+        layer.part_n_hash_cols,
+        layer.dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    # Match the non-contiguous per-layer slice of the model hash tensor.
+    hashes = torch.randint(
+        0,
+        layer.part_num_embeddings,
+        (num_tokens, 2, cols),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    src = hashes[:, 1]
+    local_tokens = (
+        (num_tokens + tp_size - 1) // tp_size if use_sequence_parallel else num_tokens
+    )
+    out = torch.empty(
+        local_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda"
+    )
+    embed = engram.embed
+    if capture == "compiled":
+        embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True)
+
+    def step(cap=None):
+        engram.prepare_embeddings(src)
+        if cap is not None:
+            cap.add_eager(lambda: None)
+        out.copy_(embed(src))
+
+    warmup = torch.cuda.Stream()
+    warmup.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(warmup):
+        step()
+    torch.cuda.current_stream().wait_stream(warmup)
+
+    graph = None
+    if capture == "full":
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            step()
+    elif capture == "breakable":
+        graph = BreakableCUDAGraphCapture()
+        with torch.cuda.stream(warmup), graph:
+            step(graph)
+        torch.cuda.current_stream().wait_stream(warmup)
+        assert graph.num_graphs == 2
+
+    for _ in range(3):
+        hashes.random_(0, layer.part_num_embeddings)
+        if graph is None:
+            step()
+        else:
+            graph.replay()
+        expected = _reference_lookup(
+            layer.weight.cuda(),
+            layer.weight_scale_inv.cuda(),
+            src,
+            layer.vocab_start_idx,
+            layer.vocab_end_idx,
+        )
+        if use_sequence_parallel:
+            expected = torch.nn.functional.pad(
+                expected, (0, 0, 0, 0, 0, (-num_tokens) % tp_size)
+            )[rank * local_tokens : (rank + 1) * local_tokens]
+        torch.testing.assert_close(out, expected, rtol=0, atol=0)
+
+
+def _engram_tp_worker(rank, tp_size, port):
+    from tests.utils import init_test_distributed_environment
+    from vllm.distributed import cleanup_dist_env_and_memory
+
+    torch.accelerator.set_device_index(rank)
+    init_test_distributed_environment(tp_size, 1, rank, str(port), local_rank=rank)
+    try:
+        for cpu_offload in (False, True):
+            for sp in (False, True):
+                for capture in ("eager", "compiled", "full", "breakable"):
+                    _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp)
+    finally:
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.distributed(num_gpus=2)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_engram_head_collectives_survive_graph_breaks():
+    """All-gather preserves head order and local SP tokens across graph replay."""
+    from vllm.utils.network_utils import get_open_port
+
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("Requires two GPUs")
+    torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2)

tests/kernels/test_fused_indexer_q_rope_quant.py

@@ -14,6 +14,7 @@ Expects bit-exact equality on both q_fp8 and weights_out.
 """
 
 import contextlib
+import importlib
 from unittest import mock
 
 import pytest
@@ -29,9 +30,10 @@ from vllm.utils.import_utils import has_cutedsl
 
 HEAD_DIM = 128
 ROPE_DIM = 64
-N_HEAD = 64
 MAX_POS = 4096
 
+CUTEDSL_MODULE = "vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl"
+
 
 def quantize_to_mxfp4(
     x: torch.Tensor,
@@ -90,6 +92,7 @@ def _reference(
     weights: torch.Tensor,
     softmax_scale: float,
     head_scale: float,
+    n_head: int,
     use_fp4: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     q_rot = q.clone()
@@ -105,7 +108,7 @@ def _reference(
     )
 
     if use_fp4:
-        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
+        q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, n_head, HEAD_DIM))
         # Pack 4 ue8m0 bytes into 1 int32
         q_scale = ue8m0.view(torch.int32).squeeze(-1)
         # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
@@ -118,8 +121,8 @@ def _reference(
             HEAD_DIM,
             use_ue8m0=True,
         )
-        q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
-        q_scale = q_scale.view(-1, N_HEAD)
+        q_fp8 = q_fp8.view(-1, n_head, HEAD_DIM)
+        q_scale = q_scale.view(-1, n_head)
 
         weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
         return q_fp8, weights_out
@@ -144,9 +147,10 @@ def _reference(
     ],
 )
 @pytest.mark.parametrize("use_cutedsl", [False, True])
+@pytest.mark.parametrize("n_head", [32, 64])
 @torch.inference_mode()
 def test_fused_indexer_q_rope_quant_matches_unfused(
-    num_tokens, cache_dtype, use_fp4, use_cutedsl
+    num_tokens, cache_dtype, use_fp4, use_cutedsl, n_head
 ):
     if use_cutedsl and not has_cutedsl():
         pytest.skip("cutedsl (cutlass) not installed")
@@ -154,17 +158,29 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
     device = "cuda"
     torch.manual_seed(0)
 
-    q = torch.randn(num_tokens, N_HEAD, HEAD_DIM, dtype=torch.bfloat16, device=device)
+    q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device)
     positions = torch.randint(
         0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device
     )
     cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device)
-    weights = torch.randn(num_tokens, N_HEAD, dtype=torch.bfloat16, device=device)
+    weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device)
     softmax_scale = HEAD_DIM**-0.5
-    head_scale = N_HEAD**-0.5
+    # head_scale must be an exact power of two for bit-exactness to be a fair
+    # expectation: the triton kernel multiplies by softmax_scale and head_scale
+    # separately (two roundings) while the cutedsl kernels fold them into one
+    # fp32 scalar first (one rounding). Those agree only when one factor is a
+    # pow2. n_head**-0.5 is a pow2 at 64 heads but not at 32.
+    head_scale = 0.125
 
     q_quant_ref, weights_ref = _reference(
-        positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
+        positions,
+        q,
+        cos_sin_cache,
+        weights,
+        softmax_scale,
+        head_scale,
+        n_head,
+        use_fp4,
     )
     # use_cutedsl=False: force the triton path even when cutedsl is installed
     # by patching the dispatcher's has_cutedsl() binding to return False.
@@ -210,3 +226,337 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
         f"weights mismatch: max abs diff "
         f"{(weights_ref - weights_fused).abs().max().item()}"
     )
+
+
+@pytest.mark.skipif(not has_cutedsl(), reason="cutedsl (cutlass) not installed")
+@pytest.mark.parametrize(
+    "use_fp4",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(
+                not (
+                    current_platform.is_cuda()
+                    and current_platform.is_device_capability_family(100)
+                ),
+                reason="MXFP4 indexer cache requires an SM100-family GPU",
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("n_head", [32, 64])
+@torch.inference_mode()
+def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head):
+    """The cutedsl kernels must not store past ``num_tokens``.
+
+    The launcher rounds ``num_tokens * threads_per_token`` up to whole 128-thread
+    blocks, so the trailing block is only partially in range unless
+    ``threads_per_token`` is itself a multiple of the block size. With 32 heads
+    and coarsen=4 (picked for >= 512 tokens) a token needs 64 threads, so an odd
+    token count leaves half of the last block addressing token ``num_tokens``.
+    Every output gets one extra sentinel row that the kernel must leave alone.
+    """
+    mod = importlib.import_module(CUTEDSL_MODULE)
+    device = "cuda"
+    torch.manual_seed(0)
+    num_tokens = 935
+
+    sentinel_u8, sentinel_f32 = 0x55, -12345.0
+    # positions/q are padded too, so an out-of-bounds *load* stays addressable
+    # and the assertions below are about stores only.
+    positions = torch.zeros(num_tokens + 1, dtype=torch.int64, device=device)
+    q = torch.randn(
+        num_tokens + 1, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device
+    )
+    cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device)
+    weights = torch.randn(num_tokens + 1, n_head, dtype=torch.bfloat16, device=device)
+    weights_out = torch.full(
+        (num_tokens + 1, n_head), sentinel_f32, dtype=torch.float32, device=device
+    )
+    outputs = {"weights_out": weights_out}
+
+    if use_fp4:
+        outputs["q_packed"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM // 2),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        outputs["q_scale"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM // 32),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        mod.fused_indexer_q_rope_quant_mxfp4_cutedsl(
+            positions[:num_tokens],
+            q[:num_tokens],
+            cos_sin_cache,
+            weights[:num_tokens],
+            1.0,
+            1.0,
+            outputs["q_packed"][:num_tokens],
+            outputs["q_scale"][:num_tokens],
+            weights_out[:num_tokens],
+        )
+    else:
+        outputs["q_fp8"] = torch.full(
+            (num_tokens + 1, n_head, HEAD_DIM),
+            sentinel_u8,
+            dtype=torch.uint8,
+            device=device,
+        )
+        mod.fused_indexer_q_rope_quant_fp8_cutedsl(
+            positions[:num_tokens],
+            q[:num_tokens],
+            cos_sin_cache,
+            weights[:num_tokens],
+            1.0,
+            1.0,
+            outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn),
+            weights_out[:num_tokens],
+        )
+    torch.accelerator.synchronize()
+
+    clobbered = {}
+    for name, tensor in outputs.items():
+        sentinel = sentinel_f32 if tensor.dtype == torch.float32 else sentinel_u8
+        n_changed = (tensor[num_tokens] != sentinel).sum().item()
+        if n_changed:
+            clobbered[name] = n_changed
+    assert not clobbered, (
+        f"kernel wrote past token {num_tokens - 1} into the sentinel row: {clobbered}"
+    )
+
+
+def _indexer_k_reference(
+    k_pre, positions, cos_sin_cache, rms_norm_weight, compress_ratio, use_fp4
+):
+    k = k_pre.float()
+    k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-20)
+    k = (k * rms_norm_weight.float()).to(torch.bfloat16).float()
+    group_positions = positions // compress_ratio * compress_ratio
+    cos, sin = cos_sin_cache[group_positions].float().chunk(2, dim=-1)
+    even, odd = k[:, 64::2], k[:, 65::2]
+    rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), dim=-1)
+    k = torch.cat((k[:, :64], rotated.flatten(1)), dim=-1).to(torch.bfloat16)
+    if use_fp4:
+        return quantize_to_mxfp4(k)
+    exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448
+    scale = exponent.log2().ceil().exp2()
+    values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return values.view(torch.uint8), scale.view(torch.uint8)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
+@pytest.mark.parametrize("num_tokens", [1, 17, 257])
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@pytest.mark.parametrize("use_fp4", [False, True])
+@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_indexer_k_inserts_only_valid_groups_into_padded_pages(
+    num_tokens, compress_ratio, use_fp4, cache_dtype
+):
+    """Insert group ends and preserve skipped slots, graph rows, and page padding."""
+    from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+
+    if use_fp4 and not current_platform.is_device_capability_family(100):
+        pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
+    torch.manual_seed(17)
+    device = "cuda"
+    k_pre = torch.randn(num_tokens + 3, 128, device=device, dtype=torch.bfloat16)
+    norm_weight = torch.randn(128, device=device, dtype=torch.bfloat16)
+    positions = torch.arange(7, num_tokens + 10, device=device)
+    slots = torch.randperm(2 * num_tokens, device=device)[:num_tokens].clone()
+    slots[5::7] = -1
+    valid = (slots >= 0) & ((positions[:num_tokens] + 1) % compress_ratio == 0)
+    k_pre[3::7] = 0
+    k_pre[8::11] *= 1e-10
+    k_pre[:num_tokens][~valid] = float("nan")
+    k_pre[num_tokens:] = float("nan")
+    positions[num_tokens:] = MAX_POS + 10
+    angles = torch.randn(MAX_POS, 32, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(cache_dtype)
+    block_size = 16
+    value_bytes = 64 if use_fp4 else 128
+    page_bytes = block_size * (value_bytes + 4)
+    num_blocks = (2 * num_tokens + block_size - 1) // block_size + 1
+    storage = torch.full(
+        (num_blocks, page_bytes + 128), 0xA5, device=device, dtype=torch.uint8
+    )
+    cache = storage.as_strided(
+        (num_blocks, block_size, value_bytes + 4),
+        (storage.stride(0), value_bytes + 4, 1),
+    )
+    expected = storage.clone()
+    values, scales = _indexer_k_reference(
+        k_pre[:num_tokens][valid],
+        positions[:num_tokens][valid],
+        cos_sin,
+        norm_weight,
+        compress_ratio,
+        use_fp4,
+    )
+    block_ids = slots[valid] // block_size
+    block_offsets = slots[valid] % block_size
+    value_offsets = block_offsets[:, None] * value_bytes + torch.arange(
+        value_bytes, device=device
+    )
+    scale_offsets = block_size * value_bytes + block_offsets[:, None] * 4
+    scale_offsets = scale_offsets + torch.arange(4, device=device)
+    expected[block_ids[:, None], value_offsets] = values
+    expected[block_ids[:, None], scale_offsets] = scales
+
+    indexer_k_norm_rope_store(
+        k_pre,
+        positions,
+        cos_sin,
+        norm_weight,
+        1e-20,
+        cache,
+        slots,
+        compress_ratio,
+        use_fp4,
+    )
+    # Fused RoPE can change zero signs without changing their numerical values.
+    actual_values = storage[block_ids[:, None], value_offsets]
+    if use_fp4:
+        for shift in (0, 4):
+            zero_mask = (actual_values >> shift) & 7 == 0
+            actual_values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
+            zero_mask = (values >> shift) & 7 == 0
+            values &= ~(zero_mask.to(torch.uint8) << (shift + 3))
+    else:
+        actual_values[actual_values == 128] = 0
+        values[values == 128] = 0
+    storage[block_ids[:, None], value_offsets] = actual_values
+    expected[block_ids[:, None], value_offsets] = values
+    torch.testing.assert_close(storage, expected, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel")
+@pytest.mark.parametrize("use_fp4", [False, True])
+@torch.inference_mode()
+def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4):
+    """Replay must consume updated keys and slot mapping through the captured API."""
+    from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+
+    if use_fp4 and not current_platform.is_device_capability_family(100):
+        pytest.skip("MXFP4 indexer cache requires an SM100-family GPU")
+    torch.manual_seed(18)
+    k_pre = torch.randn(19, 128, device="cuda", dtype=torch.bfloat16)
+    norm_weight = torch.ones(128, device="cuda", dtype=torch.bfloat16)
+    positions = torch.arange(19, device="cuda")
+    cos_sin = torch.zeros(32, 64, device="cuda")
+    cos_sin[:, :32] = 1
+    slots = torch.arange(17, device="cuda")
+    cache = torch.full(
+        (2, 16, (64 if use_fp4 else 128) + 4),
+        0xA5,
+        device="cuda",
+        dtype=torch.uint8,
+    )
+
+    def run():
+        indexer_k_norm_rope_store(
+            k_pre,
+            positions,
+            cos_sin,
+            norm_weight,
+            1e-20,
+            cache,
+            slots,
+            2,
+            use_fp4,
+        )
+
+    run()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    k_pre.neg_()
+    slots.add_(4)
+    cache.fill_(0xA5)
+    graph.replay()
+    replayed = cache.clone()
+    cache.fill_(0xA5)
+    run()
+    torch.testing.assert_close(cache, replayed, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm cache layout")
+@pytest.mark.parametrize("block_size", [16, 64, 128])
+@pytest.mark.parametrize("compress_ratio", [1, 2])
+@torch.inference_mode()
+def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_ratio):
+    """ROCm reads the indexer K cache back with the 16x16-tiled value layout.
+
+    ``cp_gather_indexer_k_quant_cache_triton`` (prefill) and aiter's
+    ``Preshuffle=True`` paged kernel (decode) both pick that layout whenever
+    ``block_size > 1``, so a row-major store would feed the indexer permuted
+    key bytes and randomize its top-k. Assert the write/read pair is exact.
+    """
+    from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import (
+        indexer_k_norm_rope_store,
+    )
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        cp_gather_indexer_k_quant_cache_triton,
+    )
+
+    torch.manual_seed(0)
+    device = "cuda"
+    num_tokens = 3 * block_size + 5
+    num_blocks = num_tokens // block_size + 2
+    k_pre = torch.randn(num_tokens, HEAD_DIM, device=device, dtype=torch.bfloat16)
+    norm_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16)
+    positions = torch.arange(num_tokens, device=device)
+    angles = torch.randn(MAX_POS, ROPE_DIM // 2, device=device)
+    cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).float()
+    cache = torch.zeros(
+        num_blocks, block_size, HEAD_DIM + 4, device=device, dtype=torch.uint8
+    )
+    slots = torch.arange(num_tokens, device=device)
+
+    indexer_k_norm_rope_store(
+        k_pre,
+        positions,
+        cos_sin,
+        norm_weight,
+        1e-20,
+        cache,
+        slots,
+        compress_ratio,
+        False,
+    )
+
+    k_fp8 = torch.zeros(num_tokens, HEAD_DIM, device=device, dtype=torch.float8_e4m3fn)
+    k_scale = torch.zeros(num_tokens, 4, device=device, dtype=torch.uint8)
+    cp_gather_indexer_k_quant_cache_triton(
+        cache,
+        k_fp8,
+        k_scale,
+        torch.arange(num_blocks, device=device, dtype=torch.int32).view(1, -1),
+        torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
+        torch.zeros(num_tokens, device=device, dtype=torch.int32),
+    )
+
+    expected_values, expected_scale = _indexer_k_reference(
+        k_pre, positions, cos_sin, norm_weight, compress_ratio, False
+    )
+    emitted = (positions + 1) % compress_ratio == 0
+    gathered = k_fp8.view(torch.uint8)
+    gathered = torch.where(gathered == 128, gathered.new_zeros(()), gathered)
+    expected_values = torch.where(
+        expected_values == 128, expected_values.new_zeros(()), expected_values
+    )
+    torch.testing.assert_close(
+        gathered[emitted], expected_values[emitted], rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        k_scale[emitted], expected_scale[emitted], rtol=0, atol=0
+    )

tests/kernels/test_fused_inv_rope_fp8_quant.py

@@ -159,13 +159,12 @@ def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Ten
       absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8
 
     Args:
-        x_f32: [..., quant_group_size] float32 — one or more 128-element blocks.
+        x_f32: [..., quant_group_size] float32 quantization blocks.
 
     Returns:
         x_fp8: same shape, float8_e4m3fn
         scales: [...] float32, one scale per block
     """
-    assert x_f32.shape[-1] == QUANT_GROUP_SIZE
     absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS)
     scale_raw = absmax * (1.0 / FP8_MAX)
     scale = torch.exp2(torch.ceil(torch.log2(scale_raw)))
@@ -253,6 +252,56 @@ def reference_inv_rope_fp8_quant(
 # =========================================================================
 
 
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 4096])
+@pytest.mark.parametrize("quant_group_size", [32, 128])
+@torch.inference_mode()
+def test_packed_quantization_preserves_rotary_dimensions(num_tokens, quant_group_size):
+    """RoPE may span multiple scale blocks; pack every scale and padding row."""
+    from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+        fused_inv_rope_fp8_quant as quantize,
+    )
+
+    torch.manual_seed(42)
+    groups, heads_per_group = 2, 2
+    x = torch.randn(
+        num_tokens,
+        groups * heads_per_group,
+        HEAD_DIM,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    cache = make_cos_sin_cache(4096)
+    positions = torch.randint(4096, (num_tokens,), device="cuda")
+    q, packed_scales = quantize(
+        x,
+        positions,
+        cache,
+        groups,
+        heads_per_group,
+        quant_group_size=quant_group_size,
+        tma_aligned_scales=True,
+    )
+    scale_bytes = torch.stack(
+        [(packed_scales >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    scales = torch.exp2(scale_bytes.float() - 127)
+    rotated = reference_inv_rope(x.float(), positions, cache)
+    ref_q, ref_scales = _ref_ue8m0_quant_block(
+        rotated.reshape(num_tokens, groups, -1, quant_group_size)
+    )
+    torch.testing.assert_close(scales, ref_scales, rtol=0, atol=0)
+    actual = q.float().unflatten(-1, (-1, quant_group_size)) * scales.unsqueeze(-1)
+    expected = ref_q.float() * ref_scales.unsqueeze(-1)
+    assert (actual - expected).norm() / expected.norm() < 1e-3
+
+    aligned_tokens = (num_tokens + 3) // 4 * 4
+    full_scales = packed_scales.transpose(0, 1).as_strided(
+        (groups, aligned_tokens, packed_scales.shape[-1]),
+        (packed_scales.stride(1), 1, packed_scales.stride(2)),
+    )
+    assert torch.count_nonzero(full_scales[:, num_tokens:]) == 0
+
+
 @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128])
 @pytest.mark.parametrize(
     "num_heads,n_groups",
@@ -260,8 +309,9 @@ def reference_inv_rope_fp8_quant(
     ids=["H64_G8", "H32_G4", "H128_G8"],
 )
 @pytest.mark.parametrize("seed", [0, 42])
+@pytest.mark.parametrize("quantize", [True, False], ids=["fp8", "bf16"])
 @torch.inference_mode()
-def test_correctness(num_tokens, num_heads, n_groups, seed):
+def test_correctness(num_tokens, num_heads, n_groups, seed, quantize):
     """Compare fused kernel against reference for FP8 values and scales."""
     torch.manual_seed(seed)
 
@@ -280,6 +330,26 @@ def test_correctness(num_tokens, num_heads, n_groups, seed):
         max_pos, ROPE_DIM, dtype=torch.float32, device=device
     )
 
+    if not quantize:
+        ref_output = reference_inv_rope(o, positions, cos_sin_cache).view(
+            num_tokens, n_groups, -1
+        )
+        fused_output, fused_scale = fused_inv_rope_fp8_quant(
+            o.clone(),
+            positions,
+            cos_sin_cache,
+            n_groups,
+            heads_per_group,
+            quantize=False,
+        )
+        d = heads_per_group * HEAD_DIM
+        assert fused_output.shape == (num_tokens, n_groups, d)
+        assert fused_scale.numel() == 0
+        assert fused_output.dtype == torch.bfloat16
+        assert fused_output.stride() == (d, num_tokens * d, 1)
+        torch.testing.assert_close(fused_output, ref_output, rtol=1e-2, atol=1e-2)
+        return
+
     # Reference
     ref_fp8, ref_scale = reference_inv_rope_fp8_quant(
         o.clone(),
@@ -726,6 +796,7 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
     This catches stride/layout bugs that only manifest when the einsum
     kernel actually consumes the quantized activations.
     """
+    from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj
     from vllm.utils.deep_gemm import (
         fp8_einsum,
         is_deep_gemm_supported,
@@ -793,23 +864,23 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
     )
 
     # -- FUSED path --
-    fused_fp8, fused_scale = fused_inv_rope_fp8_quant(
+    wo_a = torch.nn.Module()
+    wo_a.register_buffer("weight", w_fp8)
+    wo_a.register_buffer("weight_scale", w_scale_t)
+    z_fused = deep_gemm_fp8_o_proj(
         o.clone(),
         positions,
         cos_sin_cache,
-        n_groups,
-        heads_per_group,
-    )
-    z_fused = torch.empty(
-        num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16
-    )
-    fp8_einsum(
-        "bhr,hdr->bhd",
-        (fused_fp8, fused_scale),
-        (w_fp8, w_scale_t),
-        z_fused,
-        recipe=recipe,
-    )
+        wo_a,
+        torch.nn.Identity(),
+        n_groups=n_groups,
+        heads_per_group=heads_per_group,
+        nope_dim=NOPE_DIM,
+        rope_dim=ROPE_DIM,
+        o_lora_rank=o_lora_rank,
+        einsum_recipe=recipe,
+        tma_aligned_scales=True,
+    ).view_as(z_ref)
 
     # -- Checks --
     # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs

tests/kernels/test_mhc_kernels.py

@@ -12,7 +12,10 @@ import vllm.model_executor.layers.mhc as mhc_layers
 from vllm.model_executor.kernels.mhc.tilelang import (
     _tilelang_hc_prenorm_gemm,
     _torch_hc_prenorm_gemm,
+    mhc_pre_delayed_tilelang,
 )
+from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch
+from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
 from vllm.model_executor.layers.mhc import (
     HAS_AITER_MHC,
     HAS_AITER_MHC_FUSED,
@@ -26,12 +29,293 @@ from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4DecoderLayer,
     DeepseekV4Model,
 )
+from vllm.models.deepseek_v4_1.nvidia.model import (
+    DeepseekV4DecoderLayer as DeepseekV41DecoderLayer,
+)
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "num_tokens,hc_mult,hidden_size",
+    [
+        (0, 4, 5120),
+        (1, 4, 5120),
+        (33, 4, 5137),
+        (128, 4, 5120),
+        (1024, 4, 5120),
+        (2, 1, 513),
+        (7, 8, 5137),
+        (65537, 4, 65),
+    ],
+)
+@pytest.mark.parametrize("strided", [False, True])
+def test_hc_collapse_preserves_weighted_residual_sum(
+    num_tokens, hc_mult, hidden_size, strided
+):
+    """Handle tile tails and strided streams without changing FP32 mixing."""
+    set_random_seed(0)
+    x = torch.randn(
+        num_tokens,
+        hc_mult,
+        hidden_size * (2 if strided else 1),
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    pre = torch.rand(num_tokens, hc_mult, device=DEVICE)
+    if strided:
+        x = x[..., ::2]
+        pre = pre.t().contiguous().t()
+    if num_tokens:
+        pre[0].zero_()
+        pre[0, -1] = 1
+    if num_tokens > 1 and hc_mult == 4:
+        pre[1].fill_(1)
+        x[1, 1] = -x[1, 0]
+        x[1, 3] = -x[1, 2]
+    expected = (pre.unsqueeze(-1) * x.float()).sum(dim=1).to(x.dtype)
+    actual = hc_collapse_triton(x, pre)
+    assert actual.is_contiguous()
+    # BF16 rounding only: ptxas <13.1 contracts the FP32 mix into FFMA despite
+    # enable_fp_fusion=False, so compare within a ULP and pin the exact cases below.
+    torch.testing.assert_close(actual, expected, atol=1.6e-2, rtol=1e-2)
+    if num_tokens:
+        torch.testing.assert_close(actual[0], x[0, -1], atol=0, rtol=0)
+    if num_tokens > 1 and hc_mult == 4:
+        torch.testing.assert_close(
+            actual[1], torch.zeros_like(actual[1]), atol=0, rtol=0
+        )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+def test_hc_collapse_custom_op_supports_compile():
+    x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    pre = torch.rand(2, 4, device=DEVICE)
+    torch.library.opcheck(torch.ops.vllm.hc_collapse_triton.default, (x, pre))
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
+@pytest.mark.parametrize("num_tokens", [1, 7])
+def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
+    """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm."""
+    from vllm.models.deepseek_v4_1.nvidia import dspark
+
+    set_random_seed(0)
+    hidden_size, hc_mult = 5120, 4
+    streams = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    mixes = [torch.rand(num_tokens, hc_mult, device=DEVICE) for _ in range(2)]
+    carried_mixes = []
+
+    def make_layer(mix):
+        def forward(hidden, positions, ids, pre, post, res, residual):
+            carried_mixes.append(pre)
+            return hidden, None, None, None, mix
+
+        return forward
+
+    monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)
+    draft = SimpleNamespace(
+        use_sequence_parallel=False,
+        hc_mult=hc_mult,
+        layers=[make_layer(mix) for mix in mixes],
+    )
+    actual = dspark.DSparkDeepseekV4Model.forward(
+        draft,
+        input_ids=torch.zeros(num_tokens, dtype=torch.long, device=DEVICE),
+        positions=torch.arange(num_tokens, device=DEVICE),
+        inputs_embeds=torch.zeros(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE
+        ),
+    )
+    expected = (streams.float() * mixes[-1].unsqueeze(-1)).sum(dim=1)
+    assert carried_mixes[0] is None
+    assert carried_mixes[1] is mixes[0]
+    torch.testing.assert_close(
+        actual, expected.to(streams.dtype), atol=1.6e-2, rtol=1e-2
+    )
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize(
+    "num_tokens,sinkhorn_repeat,norm_eps",
+    [
+        (0, 20, 1e-6),
+        (1, 1, 1e-6),
+        (1, 20, 1e-6),
+        (33, 20, 1e-20),
+        (128, 20, 1e-6),
+        (1024, 20, 1e-6),
+    ],
+)
+@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"])
+@pytest.mark.parametrize("use_deep_gemm", [False, True])
+@pytest.mark.parametrize("fused_norm", [False, True])
+def test_deepseek_v41_mhc_pre_delayed(
+    num_tokens,
+    sinkhorn_repeat,
+    norm_eps,
+    entry,
+    use_deep_gemm,
+    fused_norm,
+    monkeypatch,
+):
+    """Collapse with the carried pre-mix and save the new one for the next sublayer."""
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if use_deep_gemm and not is_deep_gemm_supported():
+        pytest.skip("DeepGEMM not available")
+    monkeypatch.setattr(
+        "vllm.utils.deep_gemm.is_deep_gemm_supported", lambda: use_deep_gemm
+    )
+    set_random_seed(0)
+    hc_mult, hidden_size = 4, 5120
+    residual = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
+    )
+    if num_tokens > 1:
+        residual[0].zero_()
+    fn = torch.randn(24, hc_mult, hidden_size, device=DEVICE) * 0.02
+    x = None
+    pre_mix = None
+    if entry == "broadcast":
+        x = residual[:, 0].contiguous()
+        residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()
+        fn = fn.sum(1)
+    else:
+        fn = fn.flatten(1)
+        if entry == "carried":
+            pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
+            if num_tokens > 1:
+                # A constant collapse just above 1 rounds to BF16 1 before RMSNorm.
+                residual[1].fill_(1)
+                pre_mix[1].zero_()
+                pre_mix[1, 0] = 1.003
+    scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
+    base = torch.randn(24, device=DEVICE)
+    args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, sinkhorn_repeat)
+    expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix, x=x)
+    norm_kwargs = {}
+    if fused_norm:
+        weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE)
+        weight.uniform_(0.5, 1.5)
+        norm_kwargs = dict(norm_weight=weight, norm_eps=norm_eps)
+        # Compare to the original collapse + CUDA RMSNorm, including BF16 rounding.
+        unfused = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x)
+        normalized = torch.empty_like(unfused[2])
+        if num_tokens:
+            torch.ops._C.rms_norm(normalized, unfused[2], weight, norm_eps)
+        expected = (*unfused[:2], normalized, unfused[3])
+    actual = mhc_pre_delayed_tilelang(*args, pre_mix=pre_mix, x=x, **norm_kwargs)
+    if fused_norm and num_tokens:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = mhc_pre_delayed_tilelang(
+                *args, pre_mix=pre_mix, x=x, **norm_kwargs
+            )
+        graph.replay()
+        for eager, replayed in zip(actual, captured, strict=True):
+            torch.testing.assert_close(eager, replayed, atol=0, rtol=0)
+    if fused_norm and entry == "carried" and num_tokens > 1:
+        torch.testing.assert_close(actual[2][1], expected[2][1], atol=0, rtol=0)
+    # DeepGEMM uses TF32 for the projection; the TileLang fallback uses FP32.
+    atol, rtol = (1e-3, 1e-3) if use_deep_gemm else (1e-5, 1e-4)
+    for i in (0, 1, 3):
+        torch.testing.assert_close(actual[i], expected[i], atol=atol, rtol=rtol)
+    if pre_mix is None and not fused_norm:
+        torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)
+    else:
+        torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])
+def test_deepseek_v41_decoder_mixes_match_torch(
+    entry, monkeypatch, default_vllm_config
+):
+    """Preserve carried pre-mixes and Engram ordering at every decoder entry."""
+    set_random_seed(0)
+    decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)
+    nn.Module.__init__(decoder)
+    decoder.hc_mult = 4
+    decoder.hc_sinkhorn_iters = 20
+    decoder.hc_eps = 1e-6
+    decoder.rms_norm_eps = 1e-20
+    decoder.hc_post_alpha = 2.0
+    decoder.use_sequence_parallel = False
+    decoder.engram = None
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to(
+        device=DEVICE, dtype=torch.bfloat16
+    )
+    decoder.attn = lambda positions, x, _: x * 0.5
+    decoder.ffn = lambda x, input_ids: x * 0.25
+    with torch.device(DEVICE):
+        decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02
+        decoder.hc_attn_fn_broadcast = decoder.hc_attn_fn.view(24, 4, 5120).sum(1)
+        decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)
+        decoder.hc_attn_base = torch.randn(24)
+        decoder.hc_ffn_base = torch.randn(24)
+        x = torch.randn(3, 5120, dtype=torch.bfloat16)
+        positions = torch.arange(3)
+        kwargs = {}
+        if entry != "broadcast":
+            kwargs["pre_mix"] = torch.rand(3, 4)
+        if entry == "pipeline":
+            x = torch.randn(3, 4, 5120, dtype=torch.bfloat16)
+        if entry in ("residual", "engram"):
+            kwargs.update(
+                residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16),
+                post_mix=torch.rand(3, 4, 1),
+                res_mix=torch.rand(3, 4, 4),
+            )
+        if entry == "engram":
+
+            class FakeEngram(nn.Module):
+                layer_hash_index = 0
+
+                def forward(self, residual, hashes, mask):
+                    return residual + 0.125
+
+            decoder.engram = FakeEngram()
+            kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32)
+
+    actual = decoder(x, positions, None, **kwargs)
+
+    def reference(*args, norm_weight, norm_eps, **kwargs):
+        post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)
+        return post, res, decoder.attn_norm(collapsed), pre
+
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v4_1.nvidia.model.mhc_pre_delayed_tilelang",
+        reference,
+    )
+    expected = decoder(x, positions, None, **kwargs)
+    for result, ref in zip(actual, expected, strict=True):
+        torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
+@pytest.mark.parametrize("carried", [False, True])
+def test_mhc_pre_delayed_custom_op_supports_compile(carried):
+    set_random_seed(0)
+    x = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
+    fn = torch.randn(24, 20480, device=DEVICE) * 0.02
+    pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None
+    scale = torch.ones(3, device=DEVICE)
+    base = torch.zeros(24, device=DEVICE)
+    torch.library.opcheck(
+        torch.ops.vllm.mhc_pre_delayed_tilelang.default,
+        (x, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20, pre_mix),
+    )
+
+
 def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor:
     x = x.softmax(-1) + eps
     x = x / (x.sum(-2, keepdim=True) + eps)
@@ -203,7 +487,7 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
     _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
 
     torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
-    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
+    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=1e-2, rtol=1e-6)
 
 
 @pytest.mark.skipif(

tests/model_executor/layers/test_mla_short_prefill_indexer.py

@@ -290,3 +290,113 @@ def test_deepseek_v32_dispatches_selected_mha(
             expected_args[1:],
         )
     )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_select_candidate_blocks_tolerates_empty_rows():
+    """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest
+    block pin must not index -1 for them (device-side assert); they select
+    no candidate blocks and real rows still pin their newest block."""
+    block_size, topk_blocks = 8, 3
+    logits = torch.zeros(3, 64, device="cuda")
+    logits[0, 3] = 5.0  # block 0 scores highest for row 0
+    logits[2, 9] = 5.0  # block 1 scores highest for row 2
+    row_ks = torch.zeros(3, dtype=torch.int64, device="cuda")
+    row_ke = torch.tensor([40, 0, 17], device="cuda")
+    out = torch.empty(3, topk_blocks, dtype=torch.int32, device="cuda")
+
+    sparse_indexer._select_candidate_blocks(
+        logits, row_ks, row_ke, topk_blocks, block_size, out
+    )
+
+    assert out[1].tolist() == [-1, -1, -1]
+    assert out[0, 0].item() == 4 and 0 in out[0].tolist()  # newest block pinned
+    assert out[2, 0].item() == 2 and 1 in out[2].tolist()
+    assert (out[0] >= 0).all() and (out[2, :2] >= 0).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "width,block_size,k,decode",
+    [
+        (73, 8, 16, False),
+        (97, 3, 7, False),
+        (32768, 8, 2048, False),
+        (32768, 8, 2048, True),
+    ],
+)
+def test_candidate_kernels_preserve_packed_bounds_and_padding(
+    width, block_size, k, decode
+):
+    """Preserve top-k ties, newest blocks, empty rows and candidate clamping."""
+    from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+        apply_candidate_mask,
+        select_candidate_blocks,
+    )
+
+    torch.manual_seed(42)
+    rows = 6
+    logits = torch.randn(rows, width * 2, device="cuda")[:, ::2]
+    logits[0, :16] = 0
+    logits[2, 5] = float("nan")
+    starts = torch.tensor([0, 5, 3, 0, 7, 1], device="cuda", dtype=torch.int32)
+    ends = torch.tensor([0, width, width - 1, 1, 11, width], device="cuda")
+    repeat = 1
+    if decode:
+        starts = None
+        ends = torch.tensor([0, 1, width], device="cuda", dtype=torch.int32)
+        repeat = 2
+    ks = torch.zeros(rows, device="cuda", dtype=torch.int64) if decode else starts
+    ke = ends.repeat_interleave(repeat)
+    cols = torch.arange(width, device="cuda")
+    valid = (cols >= ks[:, None]) & (cols < ke[:, None])
+    scores = logits.masked_fill(~valid, -torch.inf)
+    blocks = ((cols - ks[:, None]) // block_size).clamp(min=0).long()
+    nblocks = (width + block_size - 1) // block_size
+    reduced = logits.new_full((rows, nblocks), -torch.inf)
+    reduced.scatter_reduce_(1, blocks, scores, reduce="amax", include_self=True)
+    lengths = ke - ks
+    last = ((lengths - 1) // block_size).clamp(min=0).long()
+    reduced.scatter_(
+        1, last[:, None], torch.where(lengths > 0, torch.inf, -torch.inf)[:, None]
+    )
+    top = reduced.topk(min(k, nblocks), dim=-1)
+    expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32)
+    expected[:, : top.indices.shape[1]] = torch.where(
+        top.values > -torch.inf, top.indices, -1
+    ).int()
+    actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2]
+    select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    candidates = actual.clone()
+    candidates[:, 0] = nblocks + 10
+    candidates[:, 1] = 0
+    candidates[:, 2] = 0
+    positions = ks[:, None, None] + candidates.long()[:, :, None] * block_size
+    positions = positions + torch.arange(block_size, device="cuda")
+    keep = torch.zeros(rows, width, device="cuda", dtype=torch.int8)
+    keep.scatter_reduce_(
+        1,
+        positions.clamp(0, width - 1).reshape(rows, -1),
+        (candidates >= 0)[:, :, None]
+        .expand(-1, -1, block_size)
+        .reshape(rows, -1)
+        .to(torch.int8),
+        reduce="amax",
+        include_self=True,
+    )
+    reference = logits.masked_fill((keep == 0) | ~valid, -torch.inf)
+    apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
+    torch.testing.assert_close(logits, reference, rtol=0, atol=0, equal_nan=True)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat)
+    for block in (1, -1):
+        candidates.fill_(-1)
+        candidates[:, 0] = block
+        logits.fill_(3.0)
+        graph.replay()
+        keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block)
+        reference = torch.where(keep, 3.0, -torch.inf)
+        torch.testing.assert_close(logits, reference, rtol=0, atol=0)

tests/models/multimodal/generation/test_maverick.py

@@ -616,6 +616,8 @@ def test_dummy_maverick(
 ) -> None:
     # Disable multiprocessing allows us to access model executor from LLM engine
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    # The forked test process may inherit CUDA state from test collection.
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
     model_path = create_reduced_maverick_model(
         original_model_name=original_model_name,

tests/models/registry.py

@@ -858,6 +858,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "DeepseekV4ForConditionalGeneration": _HfExamplesInfo(
         "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp",
     ),
+    "DeepseekV41ForCausalLM": _HfExamplesInfo(
+        "deepseek-ai/DeepSeek-V4.1-Flash",
+    ),
     "Dots3NoteForCausalLM": _HfExamplesInfo(
         "dots-studio/dots3-note-prev",
         is_available_online=False,
@@ -1537,6 +1540,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         is_available_online=False,
         use_original_num_layers=True,  # DSpark has >1 draft block
     ),
+    "DSparkV41DraftModel": _HfExamplesInfo(
+        "deepseek-ai/DeepSeek-V4.1-Flash",
+        speculative_model="deepseek-ai/DeepSeek-V4.1-Flash",  # draft in mtp.*
+        is_available_online=False,
+        use_original_num_layers=True,  # DSpark has >1 draft block
+    ),
     "K3DSparkModel": _HfExamplesInfo(
         "moonshotai/Kimi-K3",
         speculative_model="Inferact/Kimi-K3-DSpark",

tests/models/test_deepseek_v4_mega_moe.py

@@ -18,7 +18,13 @@ from vllm.models.deepseek_v4.nvidia.model import (
 )
 from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
+from vllm.models.deepseek_v4_1.common.mm_preprocess import (
+    IMAGE_PAD_ID,
+    IMAGE_SENTINEL_BASE_ID,
+)
+from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE
 from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
 
 pytestmark = pytest.mark.skipif(
     not current_platform.is_cuda(),
@@ -26,6 +32,133 @@ pytestmark = pytest.mark.skipif(
 )
 
 
+@pytest.fixture
+def v41_moe_config(dist_init):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=DeepseekV41Config(
+                text_config=dict(
+                    hidden_size=128,
+                    num_hidden_layers=2,
+                    n_routed_experts=8,
+                    num_experts_per_tok=2,
+                    dspark_n_routed_experts=4,
+                    dspark_num_experts_per_tok=3,
+                    n_shared_experts=1,
+                    moe_intermediate_size=128,
+                    hidden_act="silu",
+                    swiglu_limit=10.0,
+                    norm_topk_prob=True,
+                    topk_method="noaux_tc",
+                    routed_scaling_factor=1.5,
+                ),
+            ),
+        ),
+        quant_config=None,
+        kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"),
+        parallel_config=SimpleNamespace(
+            enable_expert_parallel=True,
+            enable_eplb=False,
+            eplb_config=SimpleNamespace(num_redundant_experts=0),
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+    )
+
+
+@pytest.mark.parametrize("vision", [False, True])
+@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)])
+def test_deepseek_v41_moe_routes_without_hash_table(
+    v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k
+):
+    """Main and draft layers select experts by score and preserve image routing."""
+    config = v41_moe_config.model_config.hf_config
+    config.vision_n_layers = int(vision)
+    with torch.device("cuda"):
+        moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn")
+        hidden_states = torch.randn(4, config.hidden_size)
+        input_ids = (
+            torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257])
+            if vision
+            else None
+        )
+    assert moe.gate.tid2eid is None
+    assert moe.gate.weight.shape == (num_experts, config.hidden_size)
+    assert isinstance(moe.experts, DeepseekV4MegaMoEExperts)
+    assert moe.experts.w13_weight.shape == (num_experts, 256, 64)
+    assert moe.experts.top_k == top_k
+    assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2)
+
+    with torch.no_grad():
+        moe.gate.weight.normal_(std=0.01)
+        moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda"))
+        if vision:
+            moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias)
+
+    scores = torch.nn.functional.softplus(
+        torch.nn.functional.linear(hidden_states, moe.gate.weight)
+    ).sqrt()
+    bias = moe.gate.e_score_correction_bias
+    if vision:
+        image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID)
+        bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias)
+    expected_ids = (scores + bias).topk(top_k, dim=-1).indices
+    expected_weights = scores.gather(1, expected_ids)
+    expected_weights *= config.routed_scaling_factor / expected_weights.sum(
+        dim=-1, keepdim=True
+    )
+
+    def check_routing(x, weights, ids, *, activation_clamp):
+        torch.testing.assert_close(ids, expected_ids)
+        torch.testing.assert_close(weights, expected_weights)
+        assert activation_clamp == config.swiglu_limit
+        return x.clone()
+
+    monkeypatch.setattr(moe.experts, "forward", check_routing)
+    monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x)
+    torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states)
+
+
+@pytest.mark.parametrize(
+    "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))]
+)
+def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults(
+    v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected
+):
+    config = v41_moe_config.model_config.hf_config
+    config.dspark_n_routed_experts = draft_experts
+    config.dspark_num_experts_per_tok = draft_top_k
+    v41_moe_config.kernel_config.moe_backend = "auto"
+    captured = {}
+
+    def make_experts(**kwargs):
+        captured.update(kwargs)
+        return torch.nn.Identity()
+
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts
+    )
+    moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn")
+    assert moe.gate.weight.shape == (expected[0], 128)
+    assert (captured["num_experts"], captured["top_k"]) == expected
+    assert captured["hash_indices_table"] is None
+
+
+def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config):
+    config = v41_moe_config.model_config.hf_config
+    config.num_hash_layers = 1
+    config.vocab_size = 32
+    moe = DeepseekV4MoE(
+        v41_moe_config,
+        prefix="model.layers.0.ffn",
+        num_hash_layers=config.num_hash_layers,
+    )
+    assert moe.gate.tid2eid.shape == (32, 2)
+    assert moe.gate.e_score_correction_bias is None
+    with pytest.raises(ValueError, match="hash MoE routing requires input_ids"):
+        moe(torch.zeros(1, 128))
+
+
 def test_deepseek_v4_mega_moe_expert_mapping():
     mapping = make_deepseek_v4_expert_params_mapping(2)
 
@@ -169,15 +302,79 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
     assert torch.count_nonzero(experts.w13_weight[1]) == 0
 
 
+@pytest.mark.parametrize("intermediate_size", [512, 2304])
+def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size):
+    """Pad gate/up separately and zero every added down-projection column."""
+    experts = DeepseekV4MegaMoEExperts(
+        SimpleNamespace(
+            scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+            compilation_config=SimpleNamespace(static_forward_context={}),
+        ),
+        num_experts=2,
+        num_local_experts=2,
+        experts_start_idx=0,
+        top_k=1,
+        hidden_size=128,
+        intermediate_size=intermediate_size,
+    )
+    originals = []
+    for param in (
+        experts.w13_weight,
+        experts.w13_weight_scale,
+        experts.w2_weight,
+        experts.w2_weight_scale,
+    ):
+        param.data.random_(1, 128)
+        originals.append(param.data.clone())
+    monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None)
+    monkeypatch.setattr(
+        "vllm.utils.deep_gemm._import_deep_gemm",
+        lambda: SimpleNamespace(
+            transform_sf_into_required_layout=lambda sf, *args: sf,
+            transform_weights_for_mega_moe=lambda l1, l2: (l1, l2),
+        ),
+    )
+
+    experts.finalize_weights()
+    padded_size = experts.intermediate_size
+    assert padded_size == (intermediate_size + 511) // 512 * 512
+    for actual, original in zip(experts._transformed_l1_weights, originals[:2]):
+        expected = (
+            original
+            if actual.dtype == torch.int8
+            else (experts._ue8m0_uint8_to_float(original))
+        )
+        actual = actual.view(expected.dtype).unflatten(1, (2, padded_size))
+        assert torch.equal(
+            actual[:, :, :intermediate_size],
+            expected.unflatten(1, (2, intermediate_size)),
+        )
+        assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0
+    for actual, original in zip(experts._transformed_l2_weights, originals[2:]):
+        expected = (
+            original
+            if actual.dtype == torch.int8
+            else (experts._ue8m0_uint8_to_float(original))
+        )
+        actual = actual.view(expected.dtype)
+        assert torch.equal(actual[..., : original.shape[-1]], expected)
+        assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0
+
+    transformed_l1 = experts._transformed_l1_weights
+    experts.finalize_weights()
+    assert experts._transformed_l1_weights is transformed_l1
+
+
 @pytest.mark.parametrize(
-    "hidden_size,intermediate_size,block_size,fused",
+    "hidden_size,intermediate_size,block_size,fused,mxfp8",
     [
-        pytest.param(128, 512, 128, True, id="aligned-block128"),
-        pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"),
+        pytest.param(128, 512, 128, True, False, id="aligned-block128"),
+        pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"),
+        pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"),
     ],
 )
 def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
-    monkeypatch, hidden_size, intermediate_size, block_size, fused
+    monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8
 ):
     """V4.1-Flash's routed padding must preserve the serial shared-expert fallback."""
 
@@ -261,6 +458,16 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(
             ),
         ),
     )
+    if mxfp8:
+        for linear in (shared_experts.gate_up_proj, shared_experts.down_proj):
+            linear.weight_block_size = (1, 32)
+            linear.weight_scale = torch.nn.Parameter(
+                linear.weight_scale_inv.view(torch.uint8)
+                .repeat_interleave(128, dim=0)
+                .repeat_interleave(4, dim=1),
+                requires_grad=False,
+            )
+            del linear.weight_scale_inv
     monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
 
     original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()

tests/models/test_dspark_mla.py

@@ -186,3 +186,79 @@ def test_context_kv_weights_are_loaded_as_merged_linear_shards():
     assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1]
     assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr()
     assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr()
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32]
+)
+@pytest.mark.parametrize(
+    ("checkpoint_name", "runtime_module", "shard_id"),
+    [
+        ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0),
+        ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1),
+        ("mtp.0.main_proj.scale", "model.main_proj", None),
+        (
+            "mtp.0.ffn.shared_experts.w1.scale",
+            "model.layers.0.ffn.shared_experts.gate_up_proj",
+            0,
+        ),
+        (
+            "mtp.0.ffn.shared_experts.w2.scale",
+            "model.layers.0.ffn.shared_experts.down_proj",
+            None,
+        ),
+    ],
+)
+def test_v41_dspark_loads_linear_scales(
+    monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id
+):
+    """Checkpoint ``.scale`` maps to the quant method's scale parameter and
+    loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static
+    loader (see tests/quantization/test_modelopt.py), not in load_weights."""
+    from vllm.models.deepseek_v4_1.nvidia import dspark
+
+    mxfp8 = scale_dtype != torch.float32
+    scale_name = "weight_scale" if mxfp8 else "weight_scale_inv"
+    runtime_name = f"{runtime_module}.{scale_name}"
+    raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8)
+    checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float()
+    param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False)
+    shards = []
+
+    def load_scale(param, weight, *args):
+        shards.append(args)
+        assert weight.dtype == checkpoint_scale.dtype
+        param.copy_(weight)
+
+    param.weight_loader = load_scale
+    draft = SimpleNamespace(
+        config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1),
+        quant_config=SimpleNamespace(
+            weight_block_size=[32, 32] if mxfp8 else [128, 128]
+        ),
+        linear_scale_name=scale_name,
+        pad_shared_expert=False,
+        model=SimpleNamespace(
+            layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))],
+            confidence_head=None,
+        ),
+        named_parameters=lambda: [(runtime_name, param)],
+        process_weights_after_loading=lambda: None,
+    )
+    draft._remap_dspark_name = lambda name: (
+        dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name)
+    )
+    monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: []
+    )
+
+    loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights(
+        draft, [(checkpoint_name, checkpoint_scale)]
+    )
+
+    assert loaded == {runtime_name}
+    assert shards == [() if shard_id is None else (shard_id,)]
+    torch.testing.assert_close(param, checkpoint_scale)

tests/models/test_registry.py

@@ -93,6 +93,7 @@ def test_registry_imports(model_arch):
     [
         ("LlamaForCausalLM", False, False, "bi-encoder"),
         ("LlavaForConditionalGeneration", True, True, "bi-encoder"),
+        ("DeepseekV41ForCausalLM", True, False, "bi-encoder"),
         ("BertForSequenceClassification", False, False, "cross-encoder"),
         ("RobertaForSequenceClassification", False, False, "cross-encoder"),
         ("XLMRobertaForSequenceClassification", False, False, "cross-encoder"),

tests/quantization/test_fp8.py

@@ -37,6 +37,7 @@ from vllm.model_executor.layers.quantization.fp8 import (
     Fp8MoEMethod,
 )
 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
+from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod
 from vllm.model_executor.layers.quantization.online.fp8 import (
     Fp8PerTensorOnlineLinearMethod,
 )
@@ -63,6 +64,383 @@ MODELS = [
 ]
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
+@pytest.mark.parametrize("tp_rank", [0, 1])
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
+def test_deepseek_v41_mxfp8_scale_loading(
+    dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype
+):
+    """Expand linear scales before TP slicing, including padded shared experts."""
+    from vllm.model_executor.layers import linear as linear_module
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.models.deepseek_v4_1 import quant_config as quant_module
+    from vllm.models.deepseek_v4_1.nvidia import model as model_module
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    quant_config = quant_module.DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
+    )
+    quant_config._resolved_expert_dtype = "fp4"
+    for module in (model_module, linear_module):
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank)
+    tp_args = dict(quant_config=quant_config, bias=False)
+    block = torch.nn.Module()
+    block.attn = torch.nn.Module()
+    block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args)
+    block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args)
+    block.ffn = torch.nn.Module()
+    block.ffn.shared_experts = torch.nn.Module()
+    block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear(
+        128, [128, 128], **tp_args
+    )
+    block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args)
+    block.ffn.experts = torch.nn.Module()
+    block.ffn.experts.register_parameter(
+        "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False)
+    )
+    block.engram = torch.nn.Module()
+    block.engram.register_parameter(
+        "weight_scale_inv",
+        torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False),
+    )
+
+    def load_expert_scale(param, weight, *args, **kwargs):
+        param.data.copy_(weight)
+        return True
+
+    block.ffn.experts.weight_scale.weight_loader = load_expert_scale
+    model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model)
+    torch.nn.Module.__init__(model)
+    model.layers = torch.nn.ModuleList([block])
+    model.config = SimpleNamespace(num_attention_heads=2)
+    model.quant_config = quant_config
+    model.use_sequence_parallel = False
+    model.get_expert_mapping = lambda: [
+        ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")
+    ]
+
+    checkpoints = [
+        ("attn.wq_b", "attn.wq_b", 128, 128, 0, None),
+        ("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)),
+        ("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)),
+        (
+            "ffn.shared_experts.w1",
+            "ffn.shared_experts.gate_up_proj",
+            96,
+            128,
+            0,
+            slice(0, 64),
+        ),
+        (
+            "ffn.shared_experts.w3",
+            "ffn.shared_experts.gate_up_proj",
+            96,
+            128,
+            0,
+            slice(64, 128),
+        ),
+        (
+            "ffn.shared_experts.down_proj",
+            "ffn.shared_experts.down_proj",
+            128,
+            96,
+            1,
+            None,
+        ),
+    ]
+    weights = []
+    expected = []
+    for source, target, n, k, axis, shard in checkpoints:
+        weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn)
+        scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8)
+        weights.extend(
+            [
+                (f"layers.0.{source}.weight", weight),
+                (f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)),
+            ]
+        )
+        dequant = (
+            weight.float().reshape(n // 32, 32, k // 32, 32)
+            * torch.exp2(scale_bytes.float() - 127)[:, None, :, None]
+        ).reshape(n, k)
+        if "shared_experts" in source:
+            padded = torch.zeros(128, 128)
+            padded[:n, :k] = dequant
+            dequant = padded
+        expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank]))
+
+    expert_scale = torch.full((3, 4), 125, dtype=torch.uint8)
+    weights.append(
+        ("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype))
+    )
+    weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype)))
+    mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
+    loaded = model.load_weights(
+        (name.removeprefix("model."), weight) for name, weight in mapper.apply(weights)
+    )
+    assert "layers.0.ffn.experts.weight_scale" in loaded
+    assert torch.equal(block.ffn.experts.weight_scale, expert_scale)
+    assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale)
+    for target, shard, reference in expected:
+        linear = block.get_submodule(target)
+        assert isinstance(linear.quant_method, ModelOptLinearMethod)
+        assert f"layers.0.{target}.weight_scale" in loaded
+        weight = linear.weight if shard is None else linear.weight[shard]
+        scale = linear.weight_scale if shard is None else linear.weight_scale[shard]
+        actual = (
+            weight.float().unflatten(-1, (-1, 32))
+            * torch.exp2(scale.float() - 127).unsqueeze(-1)
+        ).flatten(-2)
+        torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA")
+@pytest.mark.parametrize(
+    "weight_block_size,expert_dtype,scale_name",
+    [
+        ([32, 32], "fp4", "weight_scale"),
+        ([128, 128], "fp4", "weight_scale_inv"),
+        ([32, 32], "fp8", "weight_scale_inv"),
+    ],
+)
+def test_deepseek_v41_vl_mapper_routes_linear_scales(
+    weight_block_size, expert_dtype, scale_name
+):
+    """The VL wrapper must map ``.scale`` keys to the parameter the linear
+    quant method registers, as the text model does. A hardcoded
+    ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints."""
+    from vllm.models.deepseek_v4_1.nvidia import model as model_module
+    from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module
+
+    vllm_config = SimpleNamespace(
+        quant_config=SimpleNamespace(weight_block_size=weight_block_size)
+    )
+    resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype)
+    assert resolved == scale_name
+
+    mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved)
+    weight = torch.empty(0)
+    mapped = [
+        name
+        for name, _ in mapper.apply(
+            [("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)]
+        )
+    ]
+    assert mapped == [
+        f"language_model.model.layers.0.attn.wq_a.{scale_name}",
+        "language_model.model.layers.0.attn.wkv.weight",
+    ]
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
+@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
+@pytest.mark.parametrize(
+    "weight_shape,is_bmm",
+    [
+        ((129, 160), False),
+        ((387, 160), True),
+        ((3, 129, 160), False),
+        ((3, 128, 512), False),
+    ],
+)
+def test_deepgemm_mxfp8_preserves_weight_and_scale_values(
+    scale_dtype, weight_shape, is_bmm
+):
+    """DeepGEMM layout conversion preserves native weights and E8M0 scales."""
+    from vllm.model_executor.layers.quantization.utils import fp8_utils
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
+        100
+    ):
+        pytest.skip("DeepGEMM MXFP8 requires Blackwell")
+
+    weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn)
+    scales = (
+        torch.randint(
+            1,
+            255,
+            (weight.numel() // 32 + 1,),
+            device="cuda",
+            dtype=torch.uint8,
+        )[1:]
+        .view(*weight_shape[:-1], weight_shape[-1] // 32)
+        .view(scale_dtype)
+    )
+    original_weight = weight.view(torch.uint8).clone()
+    original_scales = scales.view(torch.uint8).clone()
+    processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block(
+        weight,
+        scales,
+        quant_block_shape=(1, 32),
+        use_e8m0=True,
+        is_bmm=is_bmm,
+        bmm_batch_size=3 if is_bmm else 0,
+    )
+    assert packed.dtype == torch.int32
+    unpacked = torch.stack(
+        [(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8)
+    torch.testing.assert_close(
+        unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        processed_weight.view(torch.uint8).reshape_as(original_weight),
+        original_weight,
+        rtol=0,
+        atol=0,
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
+@pytest.mark.parametrize("prequantized", [False, True])
+@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"])
+@pytest.mark.parametrize("num_tokens", [1, 7, 128])
+def test_mxfp8_bmm_loads_and_projects_grouped_weights(
+    dist_init, default_vllm_config, prequantized, config_source, num_tokens
+):
+    """BMM metadata set after construction selects grouped weight processing."""
+    from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (
+        DeepGemmMxfp8BmmLinearKernel,
+    )
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.model_executor.layers.quantization import get_quantization_config
+    from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
+        fused_inv_rope_fp8_quant,
+    )
+    from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
+        compute_fp8_einsum_recipe,
+        deep_gemm_fp8_o_proj,
+    )
+    from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model
+    from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    if not is_deep_gemm_supported() or not current_platform.is_device_capability_family(
+        100
+    ):
+        pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell")
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    quant_config = DeepseekV4FP8Config(
+        is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]
+    )
+    quant_config._resolved_expert_dtype = "fp4"
+    if config_source == "mxfp8":
+        quant_config = get_quantization_config("mxfp8").from_config(
+            {"quant_method": "mxfp8"}
+        )
+    with torch.device("cuda"):
+        linear = ColumnParallelLinear(
+            512,
+            256,
+            bias=False,
+            quant_config=quant_config,
+            return_bias=False,
+        )
+    linear.is_bmm = True
+    linear.bmm_batch_size = 2
+    weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8)
+    checkpoint_scales = (
+        scales.view(torch.float8_e8m0fnu)
+        if config_source == "deepseek"
+        else scales.repeat_interleave(32, dim=0)
+    )
+    model = DeepseekV4Model.__new__(DeepseekV4Model)
+    torch.nn.Module.__init__(model)
+    block = torch.nn.Module()
+    block.attn = torch.nn.Module()
+    block.attn.wo_a = linear
+    model.layers = torch.nn.ModuleList([block])
+    model.config = SimpleNamespace(num_attention_heads=2)
+    model.quant_config = quant_config
+    model.use_sequence_parallel = False
+    model.get_expert_mapping = lambda: []
+    model.load_weights(
+        [
+            ("layers.0.attn.wo_a.weight", weight),
+            ("layers.0.attn.wo_a.weight_scale", checkpoint_scales),
+        ]
+    )
+    reference_weight = (
+        weight.float().reshape(8, 32, 16, 32)
+        * torch.exp2(scales.float() - 127)[:, None, :, None]
+    ).reshape(2, 128, 512)
+    linear.quant_method.process_weights_after_loading(linear)
+    linear.quant_method.process_weights_after_loading(linear)
+    assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel)
+    assert linear.weight.shape == (2, 128, 512)
+    torch.testing.assert_close(
+        linear.weight.view(torch.uint8).flatten(),
+        weight.view(torch.uint8).flatten(),
+        rtol=0,
+        atol=0,
+    )
+    unpacked_scales = torch.stack(
+        [(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1
+    ).flatten(-2)
+    torch.testing.assert_close(
+        unpacked_scales.to(torch.uint8).reshape(256, 16),
+        scales.repeat_interleave(32, dim=0),
+        rtol=0,
+        atol=0,
+    )
+
+    x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16)
+    reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight)
+    inputs = x.transpose(0, 1)
+    if prequantized:
+        cache = torch.cat(
+            (
+                torch.ones(num_tokens, 32, device="cuda"),
+                torch.zeros(num_tokens, 32, device="cuda"),
+            ),
+            dim=1,
+        )
+        inputs = fused_inv_rope_fp8_quant(
+            x.permute(1, 0, 2),
+            torch.arange(num_tokens, device="cuda"),
+            cache,
+            n_groups=2,
+            heads_per_group=1,
+            nope_dim=448,
+            rope_dim=64,
+            quant_group_size=32,
+            tma_aligned_scales=current_platform.has_device_capability(100),
+        )
+    with torch.no_grad():
+        output = linear(inputs)
+    assert output.shape == reference.shape
+    assert (output.float() - reference).norm() / reference.norm() < 0.06
+    if prequantized:
+        recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32)
+        projected = deep_gemm_fp8_o_proj(
+            x.permute(1, 0, 2),
+            torch.arange(num_tokens, device="cuda"),
+            cache,
+            linear,
+            torch.nn.Identity(),
+            n_groups=2,
+            heads_per_group=1,
+            nope_dim=448,
+            rope_dim=64,
+            o_lora_rank=128,
+            einsum_recipe=recipe,
+            tma_aligned_scales=tma_aligned_scales,
+        )
+        torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0)
+    compiled = torch.compile(linear, backend="eager", fullgraph=True)
+    with torch.no_grad():
+        torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0)
+
+
 def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch):
     monkeypatch.setattr(
         flashinfer_utils,

tests/quantization/test_modelopt.py

@@ -127,6 +127,23 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
     assert method.spec.activation is kNvfp4Dynamic
 
 
+def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch):
+    """Standard MXFP8 checkpoints already have one scale row per weight row."""
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    kernel = Mock()
+    kernel.input_quant_key.return_value = None
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel",
+        lambda **kwargs: kernel,
+    )
+    config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})
+    linear = ReplicatedLinear(64, 64, bias=False, quant_config=config)
+    scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2)
+    linear.weight_scale.weight_loader(linear.weight_scale, scales)
+    assert torch.equal(linear.weight_scale, scales)
+
+
 def test_modelopt_fp8_updates_weight_dims_after_transpose():
     """Humming reads weight.input_dim/output_dim. Swapping the
     ModelWeightParameter for a plain Parameter drops them, so the per-tensor

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -7,19 +7,61 @@ import pytest
 import torch
 
 from tests.v1.attention.utils import create_vllm_config
+from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend
+from vllm.models.deepseek_v4_1.sparse_mla import (
+    DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
+)
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.mla.compressor_utils import (
     CompressedSlotMappingKernel,
 )
 from vllm.v1.attention.backends.mla.indexer import (
     BuildPrefillChunkMetadataKernel,
+    DeepseekV4IndexerBackend,
     DeepseekV32IndexerMetadataBuilder,
+    DeepseekV41IndexerBackend,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
     ConvertReqIndexToGlobalIndexKernel,
 )
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    MLAAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+)
 from vllm.v1.worker.block_table import get_block_table_width
+from vllm.v1.worker.utils import select_common_block_size
+
+
+def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla():
+    """Packed MLA/indexer groups must retain 64 compressed rows per page."""
+    kernel_block_size = select_common_block_size(
+        256, [DeepseekV4SparseMLABackend, DeepseekV4IndexerBackend]
+    )
+    spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=132,
+        dtype=torch.uint8,
+        tokens_per_state=4,
+    )
+    assert compute_layer_kv_cache_shape_bytes(spec, 2, kernel_block_size) == (
+        2,
+        1,
+        64,
+        132,
+    )
+
+
+def test_indexer_preserves_deepseek_v41_mla_block_size():
+    """V4.1 retains its smaller pages independently of V4's C4 indexer."""
+    block_size = DeepseekV41SparseMLABackend.get_supported_kernel_block_sizes()[0]
+    assert isinstance(block_size, int)
+    assert (
+        select_common_block_size(
+            block_size, [DeepseekV41SparseMLABackend, DeepseekV41IndexerBackend]
+        )
+        == block_size
+    )
 
 
 def test_indexer_warmup_normalizes_zero_compress_ratios():

tests/v1/core/test_contiguous_kv_packing.py

@@ -398,6 +398,56 @@ class TestCSALinearGrouping:
         assert len(gdn) == 10
 
 
+class TestSlidingWindowBucketCap:
+    def test_sliding_window_bucket_is_capped_at_the_main_page(self):
+        """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an
+        unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would
+        set the block width and the MLA group would fill under a third of
+        every block; capped, it splits into groups no wider than the MLA
+        page."""
+        config = _shared_layout_config()
+        specs: dict[str, KVCacheSpec] = {}
+        for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)):
+            specs[f"layers.{layer}.attn"] = MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=ratio,
+                alignment=576,
+            )
+            specs[f"layers.{layer}.idx"] = MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=132,
+                dtype=torch.uint8,
+                tokens_per_state=ratio,
+                alignment=576,
+            )
+        for layer in range(43):
+            specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec(
+                block_size=64,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                sliding_window=128,
+                alignment=576,
+            )
+        groups = get_kv_cache_groups(config, specs)
+        pages = _pages(groups)
+        main = next(g for g in groups if "layers.2.attn" in g.layer_names)
+        main_bytes = sum(pages[n] for n in main.layer_names)
+        assert _get_kv_cache_bytes_per_block(groups) == main_bytes
+        swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")]
+        per_group = main_bytes // pages["layers.0.swa"]
+        assert len(swa_groups) == -(-43 // per_group)
+        for g in swa_groups:
+            assert sum(pages[n] for n in g.layer_names) <= main_bytes
+        assert sorted(n for g in swa_groups for n in g.layer_names) == sorted(
+            f"layers.{i}.swa" for i in range(43)
+        )
+
+
 class TestDensePacking:
     def test_bytes_per_block_is_largest_group(self):
         groups, g1, g2 = _mixed_page_groups()
@@ -548,3 +598,55 @@ class TestDensePacking:
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])
+
+
+class TestCompressorRingGroup:
+    def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self):
+        """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity
+        8) beside 128-token paged MLA groups and 64-token SWA groups. The ring
+        claims one block per request, never hashes, and does not disturb the
+        scheduler/hash block sizes the paged groups imply."""
+        config = _shared_layout_config()
+        config.cache_config.enable_prefix_caching = True
+        specs: dict[str, KVCacheSpec] = {
+            "layers.2.attn": MLAAttentionSpec(
+                block_size=128,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                tokens_per_state=2,
+                alignment=576,
+            ),
+            "layers.2.attn.compressor.state_cache": CircularBufferSpec(
+                block_size=8,
+                num_kv_heads=1,
+                head_size=1024,
+                head_size_v=0,
+                dtype=torch.float32,
+            ),
+        }
+        for layer in range(4):
+            specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec(
+                block_size=64,
+                num_kv_heads=1,
+                head_size=584,
+                dtype=torch.uint8,
+                sliding_window=128,
+                alignment=576,
+            )
+        groups = get_kv_cache_groups(config, specs)
+        ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]]
+        assert len(ring_groups) == 1
+        assert not ring_groups[0].kv_cache_spec.prefix_cacheable
+        kv_cache_config = get_kv_cache_config_from_groups(
+            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+        )
+        assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64)
+        manager = KVCacheManager(
+            generate_scheduler_kv_cache_config([kv_cache_config]),
+            max_model_len=8192,
+            enable_caching=True,
+            hash_block_size=64,
+            scheduler_block_size=128,
+        )
+        assert len(manager.coordinator.single_type_managers) == len(groups)

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -333,6 +333,57 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
     assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
 
 
+@pytest.mark.parametrize("keyword", [False, True])
+def test_quantized_activation_replay_does_not_retain_tensors(
+    cuda_capture_stream, keyword
+):
+    """Release the producer tensors while preserving both buffers on replay."""
+    import gc
+    import weakref
+
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+    from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
+
+    @eager_break_during_capture
+    def consume(x, data_out, scale_out):
+        data_out.copy_(x.data)
+        scale_out.copy_(x.scale)
+
+    source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16)
+    data_out = source.to(torch.float8_e4m3fn)
+    scale_out = source[:, ::32].to(torch.uint8)
+    output = torch.empty_like(source)
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        data = source.to(torch.float8_e4m3fn)
+        scale = source[:, ::32].to(torch.uint8)
+        activation = QuantizedActivation(
+            data, scale, source.dtype, source.shape, kMxfp8Dynamic
+        )
+        data_ref, scale_ref = weakref.ref(data), weakref.ref(scale)
+        if keyword:
+            consume(x=activation, data_out=data_out, scale_out=scale_out)
+        else:
+            consume(activation, data_out, scale_out)
+        output.copy_(data_out)
+
+    del data, scale, activation
+    gc.collect()
+    assert data_ref() is None
+    assert scale_ref() is None
+
+    for value in (1, 2, 3):
+        source.fill_(value)
+        cap.replay()
+        cuda_capture_stream.synchronize()
+        torch.testing.assert_close(output, source, rtol=0, atol=0)
+        torch.testing.assert_close(scale_out, source[:, ::32].to(torch.uint8))
+
+
 def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
     """Handle an eager attention break inside a multi-stream overlap region."""
     from vllm.compilation.breakable_cudagraph import (

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (
 )
 from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KpoolTailSpec,
     KVCacheConfig,
@@ -798,3 +799,101 @@ def test_worker_setup_tolerates_finer_scratch_group():
     # The scratch group's namespace never enters the store.
     scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
     assert not any(key.startswith(scratch_prefix) for key in store._data)
+
+
+def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
+    """A per-request ring group (CircularBufferSpec, capacity as block_size)
+    sits beside the paged group: worker setup tolerates its block size (the
+    DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged
+    blocks only, and lookup hits on the paged group alone."""
+    full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
+    ring = CircularBufferSpec(
+        block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8
+    )
+    cfg = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=4 * full.page_size_bytes,
+                layers=["L0"],
+                layer_stride=4 * full.page_size_bytes,
+                block_stride=full.page_size_bytes,
+            ),
+            KVCacheTensor(
+                size=4 * ring.page_size_bytes,
+                layers=["L1"],
+                layer_stride=4 * ring.page_size_bytes,
+                block_stride=ring.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["L0"], full),
+            KVCacheGroupSpec(["L1"], ring),
+        ],
+    )
+    vllm_config = _minimal_vllm_config(cache_block_size=16)
+    store = _DictStore()
+
+    worker = _build_worker_with_dict_store(vllm_config, cfg, store)
+    worker.tp_size = 1
+    worker.pp_size = 1
+    worker.num_kv_head = 8
+    assert worker.token_dbs[1].hash_block_size == 8
+
+    raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
+    raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
+    kv_caches = {
+        "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0],
+        "L1": dense_kv_cache_views(raw_ring, ring, 4, 1, KVCacheLayout.LBNHC)[0],
+    }
+
+    def _fake_thread_init(*args, **kwargs):
+        for v in list(args) + list(kwargs.values()):
+            if isinstance(v, threading.Event):
+                v.set()
+        m = MagicMock()
+        m.start = lambda: None
+        return m
+
+    with (
+        patch.object(
+            mooncake_store_worker,
+            "KVCacheStoreSendingThread",
+            side_effect=_fake_thread_init,
+        ),
+        patch.object(
+            mooncake_store_worker,
+            "KVCacheStoreRecvingThread",
+            side_effect=_fake_thread_init,
+        ),
+    ):
+        worker.register_kv_caches(kv_caches)
+
+    send_thread = KVCacheStoreSendingThread(
+        store=store,
+        token_databases=worker.token_dbs,
+        block_size=worker.block_size,
+        coord=worker.coord,
+        tp_rank=worker.tp_rank,
+        group_put_steps=worker._group_tp_replication_factors,
+        kv_role=worker.kv_role,
+        ready_event=threading.Event(),
+        enable_kv_event=False,
+        group_participates=[True, False],
+    )
+    hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
+    save_req = ReqMeta(
+        req_id="r0",
+        token_len_chunk=64,
+        # The ring group holds one block for the request's lifetime.
+        block_ids=([1, 2, 3, 4], [1]),
+        block_hashes=hs,
+        can_save=True,
+        store_job_id=1,
+    )
+    send_thread.add_request(save_req)
+    send_thread._handle_request(send_thread.request_queue.get())
+    worker.store = store
+
+    assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64
+    assert not any("@group:1" in key for key in store._data)

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -1251,3 +1251,112 @@ def test_csa_linear_remote_ple_is_copied_whole():
             tp_ratio=-2,
             transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
         )
+
+
+def _make_ring_worker():
+    """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a
+    model whose sliding-window KV lives in per-request rings."""
+    from unittest.mock import MagicMock
+
+    from vllm.config import set_current_vllm_config
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import (
+        CircularBufferSpec,
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MLAAttentionSpec,
+        UniformTypeKVCacheSpecs,
+    )
+
+    paged_specs = {
+        f"paged.{index}": MLAAttentionSpec(
+            block_size=4, num_kv_heads=1, head_size=64, dtype=torch.float16
+        )
+        for index in range(2)
+    }
+    ring_specs = {
+        f"ring.{index}": CircularBufferSpec(
+            block_size=6,
+            num_kv_heads=1,
+            head_size=32,
+            head_size_v=0,
+            dtype=torch.float16,
+        )
+        for index in range(2)
+    }
+
+    def group(specs):
+        uniform = UniformTypeKVCacheSpecs.from_specs(specs)
+        assert uniform is not None
+        return KVCacheGroupSpec(list(specs), uniform)
+
+    # Both rings overlay the second paged region, so the scratch group must
+    # address only that region while the paged group spans both.
+    tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1"))
+    region_size, page_size = 512, 256
+    kv_cache_config = KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=len(tensor_regions) * region_size,
+                layers=[layer_name],
+                layer_stride=region_size,
+                block_stride=page_size,
+                offset=region_index * region_size,
+            )
+            for region_index, layer_names in enumerate(tensor_regions)
+            for layer_name in layer_names
+        ],
+        kv_cache_groups=[group(paged_specs), group(ring_specs)],
+    )
+
+    vllm_config = create_vllm_config(block_size=4)
+    vllm_config.cache_config.enable_prefix_caching = False
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [4]
+    fake_backend.get_name.return_value = "TEST_ATTN"
+    fake_backend.full_cls_name.return_value = "test.AttentionBackend"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    with (
+        patch.object(bw, "NixlWrapper", _RecordingNixl),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
+        tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(2)]
+        worker.register_kv_caches(
+            {
+                layer_name: tensors[region_index]
+                for region_index, layer_names in enumerate(tensor_regions)
+                for layer_name in layer_names
+            }
+        )
+    return worker
+
+
+@pytest.mark.cpu_test
+def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions():
+    worker = _make_ring_worker()
+
+    assert worker._ple_group_index is None
+    assert worker._scratch_region_indices == [1]
+
+    desc_ids = worker._compute_desc_ids(
+        block_ids=([1], [0]),
+        dst_num_blocks=2,
+        block_size_ratio=None,
+        physical_blocks_per_logical=1,
+    )
+    # Paged block 1 in both regions; ring block 0 only in its scratch region.
+    assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0])

tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py

@@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector i
     SimpleCPUOffloadConnector,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -90,3 +91,59 @@ def test_disk_mode_rejects_non_positive_capacity() -> None:
                 "disk_capacity_bytes": 0,
             }
         )
+
+
+def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig:
+    """A paged full-attention group beside a per-request ring group."""
+    full = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    ring = CircularBufferSpec(
+        block_size=6, num_kv_heads=1, head_size=32, head_size_v=0, dtype=DTYPE
+    )
+    return KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=_BYTES_PER_BLOCK * num_blocks,
+                layers=["attn"],
+                layer_stride=_BYTES_PER_BLOCK * num_blocks,
+                block_stride=_BYTES_PER_BLOCK,
+            ),
+            KVCacheTensor(
+                size=ring.page_size_bytes * num_blocks,
+                layers=["ring"],
+                layer_stride=ring.page_size_bytes * num_blocks,
+                block_stride=ring.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["attn"], full),
+            KVCacheGroupSpec(["ring"], ring),
+        ],
+    )
+
+
+def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None:
+    """The ring group's capacity is not a token granularity, it holds no
+    hashed blocks, and it never participates in CPU hits."""
+    vllm_config = create_vllm_config(
+        kv_connector="SimpleCPUOffloadConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"cpu_bytes_to_use": _BYTES_PER_BLOCK * 8},
+    )
+    connector = SimpleCPUOffloadConnector(
+        vllm_config=vllm_config,
+        role=KVConnectorRole.SCHEDULER,
+        kv_cache_config=_make_ring_kv_cache_config(),
+    )
+    scheduler = connector.scheduler_manager
+    assert scheduler.fa_gidx == 0
+    cpu_groups = scheduler.cpu_kv_cache_config.kv_cache_groups
+    assert [g.kv_cache_spec.prefix_cacheable for g in cpu_groups] == [True, False]
+    # Only the paged group takes part in CPU-side hit lookup.
+    attention_groups = scheduler.cpu_coordinator.attention_groups
+    assert [g.group_ids for g in attention_groups] == [[0]]

tests/v1/simple_kv_offload/test_scheduler.py

@@ -1319,12 +1319,10 @@ def test_inflight_finish_deferred_cleanup() -> None:
 # Test 8: Null GPU blocks are skipped in store and load transfer pairs
 # ---------------------------------------------------------------------------
 def test_multi_group_null_blocks_skipped() -> None:
-    """Null GPU blocks (no block_hash) must not appear in store or load pairs.
+    """The null block id never appears in store or load pairs.
 
-    In eager store mode, _prepare_eager_store_specs skips blocks whose
-    block_hash is None (null blocks have no hash). We verify this by mixing
-    real hashed blocks with unhashed (null-like) blocks in a single group and
-    checking that only real blocks appear in the store list.
+    A null block-table slot whose hash is still cached on the GPU is stored
+    from the cached block instead; the null block itself is never copied.
     """
     fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False)
     sched = fix.scheduler
@@ -1356,9 +1354,11 @@ def test_multi_group_null_blocks_skipped() -> None:
         f"Null block id {null_block_id} should not appear in store transfer pairs"
     )
 
-    # Only real block should be scheduled for store
-    assert len(meta.store_gpu_blocks) == 1
+    # The nulled slot's hash is still cached on the GPU, so that block is
+    # recovered from the pool and stored alongside the real one.
+    assert len(meta.store_gpu_blocks) == 2
     assert gpu_blocks[0].block_id in meta.store_gpu_blocks
+    assert gpu_blocks[1].block_id in meta.store_gpu_blocks
 
     # Complete the store
     assert meta.store_event >= 0
@@ -1374,13 +1374,12 @@ def test_multi_group_null_blocks_skipped() -> None:
         block_hasher=req._block_hasher,
     )
     hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
-    # Only 1 block was stored (the real one)
-    assert hit_tokens == BLOCK_SIZE
+    assert hit_tokens == 2 * BLOCK_SIZE
     assert is_async is True
 
     # Allocate new GPU blocks for the load
-    gpu_blocks2 = gpu_pool.get_new_blocks(1)
-    kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],))
+    gpu_blocks2 = gpu_pool.get_new_blocks(2)
+    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
     sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
 
     sched_out2 = make_scheduler_output({req2.request_id: 1})
@@ -2800,3 +2799,47 @@ def test_boundary_handoff_keeps_block_meta_index_parallel() -> None:
     assert req_ids == [req.request_id]
     assert block_meta is not None
     assert len(block_meta) == len(gpu_ids) == len(cpu_ids)
+
+
+def test_finished_eager_store_recovers_nulled_window_tail() -> None:
+    """A sliding-window group nulls pages that left the window before the
+    connector sees the finished request's block table, while the GPU keeps
+    them hashed in its free queue. The eager store must copy those pages, or
+    the CPU lookup (which needs the whole window) misses every time."""
+    fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False)
+    sched = fix.scheduler
+    gpu_pool = fix.gpu_block_pool
+    null_block = gpu_pool.null_block
+
+    num_blocks = 4  # == sliding window of the test SWA group
+    req = make_request(num_blocks=num_blocks)
+    fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
+    swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1)
+    # remove_skipped_blocks() ran: the two oldest window pages are freed (still
+    # hashed) and their slots hold the null block.
+    gpu_pool.free_blocks(swa_blocks[:2])
+    swa_table = [null_block, null_block, *swa_blocks[2:]]
+    kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table))
+    req.num_computed_tokens = num_blocks * BLOCK_SIZE
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    block_ids = kv_blocks.get_block_ids()
+    sched.request_finished_all_groups(req, block_ids)
+    meta = sched.build_connector_meta(make_scheduler_output({}))
+
+    assert null_block.block_id not in meta.store_gpu_blocks
+    assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)}
+    assert meta.store_event >= 0
+    simulate_store_completion(sched, meta.store_event)
+
+    req2 = Request(
+        request_id="req-window-tail-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * BLOCK_SIZE
+    assert is_async is True