[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/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(