diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 2b43938060..ed4ea67f93 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -102,6 +102,8 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/kernels/ commands: + # Isolate mHC from FlashInfer's import-time ptxas override. + - pytest -v -s kernels/test_mhc_kernels.py - pytest -v -s kernels/ --ignore=kernels/attention --ignore=kernels/core @@ -117,6 +119,7 @@ steps: --ignore=kernels/test_kda.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_mhc_kernels.py --ignore=kernels/test_shuffle_rows.py # BROKEN on main, pending kernel fixes (B200): # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index a794ff1e3a..f698a272f5 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -9,8 +9,13 @@ from vllm.model_executor.kernels.linear import ( _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_FP8_KERNELS, _POSSIBLE_INT8_KERNELS, + _POSSIBLE_MXFP8_KERNELS, _POSSIBLE_NVFP4_KERNELS, ) +from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.base import ( NvFp4LinearKernel, NvFp4LinearLayerConfig, @@ -50,6 +55,8 @@ SUPPORTING = { FlashInferFP8ScaledMMLinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, } @@ -60,6 +67,7 @@ def _all_kernel_classes() -> list[type]: _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_INT8_KERNELS, _POSSIBLE_NVFP4_KERNELS, + _POSSIBLE_MXFP8_KERNELS, ): for kernels in registry.values(): for cls in kernels: diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index fb696edd84..1c389a1117 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/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 diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 8debe7c3a3..7e5e410487 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/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)) diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py new file mode 100644 index 0000000000..1d5e62c6c4 --- /dev/null +++ b/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 + ) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1049eb658..1f6601dc6c 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/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 diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py new file mode 100644 index 0000000000..f7aa6dd03e --- /dev/null +++ b/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) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 5b568c938e..8fd19ffb0b 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/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 + ) diff --git a/tests/kernels/test_fused_inv_rope_fp8_quant.py b/tests/kernels/test_fused_inv_rope_fp8_quant.py index 5b71e59baa..d27144a38e 100644 --- a/tests/kernels/test_fused_inv_rope_fp8_quant.py +++ b/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 diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index a146c1a5c9..61e0449cae 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/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( diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 8bd97a3ab4..de6518ebec 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -290,3 +290,113 @@ def test_deepseek_v32_dispatches_selected_mha( expected_args[1:], ) ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_select_candidate_blocks_tolerates_empty_rows(): + """Full-cudagraph decode pads the batch with seq_len-0 rows. The newest + block pin must not index -1 for them (device-side assert); they select + no candidate blocks and real rows still pin their newest block.""" + block_size, topk_blocks = 8, 3 + logits = torch.zeros(3, 64, device="cuda") + logits[0, 3] = 5.0 # block 0 scores highest for row 0 + logits[2, 9] = 5.0 # block 1 scores highest for row 2 + row_ks = torch.zeros(3, dtype=torch.int64, device="cuda") + row_ke = torch.tensor([40, 0, 17], device="cuda") + out = torch.empty(3, topk_blocks, dtype=torch.int32, device="cuda") + + sparse_indexer._select_candidate_blocks( + logits, row_ks, row_ke, topk_blocks, block_size, out + ) + + assert out[1].tolist() == [-1, -1, -1] + assert out[0, 0].item() == 4 and 0 in out[0].tolist() # newest block pinned + assert out[2, 0].item() == 2 and 1 in out[2].tolist() + assert (out[0] >= 0).all() and (out[2, :2] >= 0).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "width,block_size,k,decode", + [ + (73, 8, 16, False), + (97, 3, 7, False), + (32768, 8, 2048, False), + (32768, 8, 2048, True), + ], +) +def test_candidate_kernels_preserve_packed_bounds_and_padding( + width, block_size, k, decode +): + """Preserve top-k ties, newest blocks, empty rows and candidate clamping.""" + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + select_candidate_blocks, + ) + + torch.manual_seed(42) + rows = 6 + logits = torch.randn(rows, width * 2, device="cuda")[:, ::2] + logits[0, :16] = 0 + logits[2, 5] = float("nan") + starts = torch.tensor([0, 5, 3, 0, 7, 1], device="cuda", dtype=torch.int32) + ends = torch.tensor([0, width, width - 1, 1, 11, width], device="cuda") + repeat = 1 + if decode: + starts = None + ends = torch.tensor([0, 1, width], device="cuda", dtype=torch.int32) + repeat = 2 + ks = torch.zeros(rows, device="cuda", dtype=torch.int64) if decode else starts + ke = ends.repeat_interleave(repeat) + cols = torch.arange(width, device="cuda") + valid = (cols >= ks[:, None]) & (cols < ke[:, None]) + scores = logits.masked_fill(~valid, -torch.inf) + blocks = ((cols - ks[:, None]) // block_size).clamp(min=0).long() + nblocks = (width + block_size - 1) // block_size + reduced = logits.new_full((rows, nblocks), -torch.inf) + reduced.scatter_reduce_(1, blocks, scores, reduce="amax", include_self=True) + lengths = ke - ks + last = ((lengths - 1) // block_size).clamp(min=0).long() + reduced.scatter_( + 1, last[:, None], torch.where(lengths > 0, torch.inf, -torch.inf)[:, None] + ) + top = reduced.topk(min(k, nblocks), dim=-1) + expected = torch.full((rows, k), -1, device="cuda", dtype=torch.int32) + expected[:, : top.indices.shape[1]] = torch.where( + top.values > -torch.inf, top.indices, -1 + ).int() + actual = torch.empty(rows, k * 2, device="cuda", dtype=torch.int32)[:, ::2] + select_candidate_blocks(logits, starts, ends, k, block_size, actual, repeat) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + candidates = actual.clone() + candidates[:, 0] = nblocks + 10 + candidates[:, 1] = 0 + candidates[:, 2] = 0 + positions = ks[:, None, None] + candidates.long()[:, :, None] * block_size + positions = positions + torch.arange(block_size, device="cuda") + keep = torch.zeros(rows, width, device="cuda", dtype=torch.int8) + keep.scatter_reduce_( + 1, + positions.clamp(0, width - 1).reshape(rows, -1), + (candidates >= 0)[:, :, None] + .expand(-1, -1, block_size) + .reshape(rows, -1) + .to(torch.int8), + reduce="amax", + include_self=True, + ) + reference = logits.masked_fill((keep == 0) | ~valid, -torch.inf) + apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat) + torch.testing.assert_close(logits, reference, rtol=0, atol=0, equal_nan=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + apply_candidate_mask(logits, starts, ends, candidates, block_size, repeat) + for block in (1, -1): + candidates.fill_(-1) + candidates[:, 0] = block + logits.fill_(3.0) + graph.replay() + keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) + reference = torch.where(keep, 3.0, -torch.inf) + torch.testing.assert_close(logits, reference, rtol=0, atol=0) diff --git a/tests/models/multimodal/generation/test_maverick.py b/tests/models/multimodal/generation/test_maverick.py index ff6e523e5b..d9afd93faa 100644 --- a/tests/models/multimodal/generation/test_maverick.py +++ b/tests/models/multimodal/generation/test_maverick.py @@ -616,6 +616,8 @@ def test_dummy_maverick( ) -> None: # Disable multiprocessing allows us to access model executor from LLM engine monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + # The forked test process may inherit CUDA state from test collection. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") model_path = create_reduced_maverick_model( original_model_name=original_model_name, diff --git a/tests/models/registry.py b/tests/models/registry.py index a474a5a0fb..874fa9c11c 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -858,6 +858,9 @@ _MULTIMODAL_EXAMPLE_MODELS = { "DeepseekV4ForConditionalGeneration": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp", ), + "DeepseekV41ForCausalLM": _HfExamplesInfo( + "deepseek-ai/DeepSeek-V4.1-Flash", + ), "Dots3NoteForCausalLM": _HfExamplesInfo( "dots-studio/dots3-note-prev", is_available_online=False, @@ -1537,6 +1540,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { is_available_online=False, use_original_num_layers=True, # DSpark has >1 draft block ), + "DSparkV41DraftModel": _HfExamplesInfo( + "deepseek-ai/DeepSeek-V4.1-Flash", + speculative_model="deepseek-ai/DeepSeek-V4.1-Flash", # draft in mtp.* + is_available_online=False, + use_original_num_layers=True, # DSpark has >1 draft block + ), "K3DSparkModel": _HfExamplesInfo( "moonshotai/Kimi-K3", speculative_model="Inferact/Kimi-K3-DSpark", diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 4e95904f6b..58d5239285 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,7 +18,13 @@ from vllm.models.deepseek_v4.nvidia.model import ( ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.models.deepseek_v4_1.common.mm_preprocess import ( + IMAGE_PAD_ID, + IMAGE_SENTINEL_BASE_ID, +) +from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config pytestmark = pytest.mark.skipif( not current_platform.is_cuda(), @@ -26,6 +32,133 @@ pytestmark = pytest.mark.skipif( ) +@pytest.fixture +def v41_moe_config(dist_init): + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=DeepseekV41Config( + text_config=dict( + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + dspark_n_routed_experts=4, + dspark_num_experts_per_tok=3, + n_shared_experts=1, + moe_intermediate_size=128, + hidden_act="silu", + swiglu_limit=10.0, + norm_topk_prob=True, + topk_method="noaux_tc", + routed_scaling_factor=1.5, + ), + ), + ), + quant_config=None, + kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"), + parallel_config=SimpleNamespace( + enable_expert_parallel=True, + enable_eplb=False, + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + + +@pytest.mark.parametrize("vision", [False, True]) +@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)]) +def test_deepseek_v41_moe_routes_without_hash_table( + v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k +): + """Main and draft layers select experts by score and preserve image routing.""" + config = v41_moe_config.model_config.hf_config + config.vision_n_layers = int(vision) + with torch.device("cuda"): + moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") + hidden_states = torch.randn(4, config.hidden_size) + input_ids = ( + torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_PAD_ID, 129257]) + if vision + else None + ) + assert moe.gate.tid2eid is None + assert moe.gate.weight.shape == (num_experts, config.hidden_size) + assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) + assert moe.experts.w13_weight.shape == (num_experts, 256, 64) + assert moe.experts.top_k == top_k + assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2) + + with torch.no_grad(): + moe.gate.weight.normal_(std=0.01) + moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda")) + if vision: + moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias) + + scores = torch.nn.functional.softplus( + torch.nn.functional.linear(hidden_states, moe.gate.weight) + ).sqrt() + bias = moe.gate.e_score_correction_bias + if vision: + image_mask = (input_ids == IMAGE_SENTINEL_BASE_ID) | (input_ids == IMAGE_PAD_ID) + bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias) + expected_ids = (scores + bias).topk(top_k, dim=-1).indices + expected_weights = scores.gather(1, expected_ids) + expected_weights *= config.routed_scaling_factor / expected_weights.sum( + dim=-1, keepdim=True + ) + + def check_routing(x, weights, ids, *, activation_clamp): + torch.testing.assert_close(ids, expected_ids) + torch.testing.assert_close(weights, expected_weights) + assert activation_clamp == config.swiglu_limit + return x.clone() + + monkeypatch.setattr(moe.experts, "forward", check_routing) + monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) + torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states) + + +@pytest.mark.parametrize( + "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))] +) +def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults( + v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected +): + config = v41_moe_config.model_config.hf_config + config.dspark_n_routed_experts = draft_experts + config.dspark_num_experts_per_tok = draft_top_k + v41_moe_config.kernel_config.moe_backend = "auto" + captured = {} + + def make_experts(**kwargs): + captured.update(kwargs) + return torch.nn.Identity() + + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts + ) + moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn") + assert moe.gate.weight.shape == (expected[0], 128) + assert (captured["num_experts"], captured["top_k"]) == expected + assert captured["hash_indices_table"] is None + + +def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config): + config = v41_moe_config.model_config.hf_config + config.num_hash_layers = 1 + config.vocab_size = 32 + moe = DeepseekV4MoE( + v41_moe_config, + prefix="model.layers.0.ffn", + num_hash_layers=config.num_hash_layers, + ) + assert moe.gate.tid2eid.shape == (32, 2) + assert moe.gate.e_score_correction_bias is None + with pytest.raises(ValueError, match="hash MoE routing requires input_ids"): + moe(torch.zeros(1, 128)) + + def test_deepseek_v4_mega_moe_expert_mapping(): mapping = make_deepseek_v4_expert_params_mapping(2) @@ -169,15 +302,79 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): assert torch.count_nonzero(experts.w13_weight[1]) == 0 +@pytest.mark.parametrize("intermediate_size", [512, 2304]) +def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediate_size): + """Pad gate/up separately and zero every added down-projection column.""" + experts = DeepseekV4MegaMoEExperts( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + ), + num_experts=2, + num_local_experts=2, + experts_start_idx=0, + top_k=1, + hidden_size=128, + intermediate_size=intermediate_size, + ) + originals = [] + for param in ( + experts.w13_weight, + experts.w13_weight_scale, + experts.w2_weight, + experts.w2_weight_scale, + ): + param.data.random_(1, 128) + originals.append(param.data.clone()) + monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) + monkeypatch.setattr( + "vllm.utils.deep_gemm._import_deep_gemm", + lambda: SimpleNamespace( + transform_sf_into_required_layout=lambda sf, *args: sf, + transform_weights_for_mega_moe=lambda l1, l2: (l1, l2), + ), + ) + + experts.finalize_weights() + padded_size = experts.intermediate_size + assert padded_size == (intermediate_size + 511) // 512 * 512 + for actual, original in zip(experts._transformed_l1_weights, originals[:2]): + expected = ( + original + if actual.dtype == torch.int8 + else (experts._ue8m0_uint8_to_float(original)) + ) + actual = actual.view(expected.dtype).unflatten(1, (2, padded_size)) + assert torch.equal( + actual[:, :, :intermediate_size], + expected.unflatten(1, (2, intermediate_size)), + ) + assert torch.count_nonzero(actual[:, :, intermediate_size:]) == 0 + for actual, original in zip(experts._transformed_l2_weights, originals[2:]): + expected = ( + original + if actual.dtype == torch.int8 + else (experts._ue8m0_uint8_to_float(original)) + ) + actual = actual.view(expected.dtype) + assert torch.equal(actual[..., : original.shape[-1]], expected) + assert torch.count_nonzero(actual[..., original.shape[-1] :]) == 0 + + transformed_l1 = experts._transformed_l1_weights + experts.finalize_weights() + assert experts._transformed_l1_weights is transformed_l1 + + @pytest.mark.parametrize( - "hidden_size,intermediate_size,block_size,fused", + "hidden_size,intermediate_size,block_size,fused,mxfp8", [ - pytest.param(128, 512, 128, True, id="aligned-block128"), - pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"), + pytest.param(128, 512, 128, True, False, id="aligned-block128"), + pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"), + pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"), ], ) def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( - monkeypatch, hidden_size, intermediate_size, block_size, fused + monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8 ): """V4.1-Flash's routed padding must preserve the serial shared-expert fallback.""" @@ -261,6 +458,16 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( ), ), ) + if mxfp8: + for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): + linear.weight_block_size = (1, 32) + linear.weight_scale = torch.nn.Parameter( + linear.weight_scale_inv.view(torch.uint8) + .repeat_interleave(128, dim=0) + .repeat_interleave(4, dim=1), + requires_grad=False, + ) + del linear.weight_scale_inv monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 01ee0b0533..0d1507a033 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -186,3 +186,79 @@ def test_context_kv_weights_are_loaded_as_merged_linear_shards(): assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1] assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr() + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32] +) +@pytest.mark.parametrize( + ("checkpoint_name", "runtime_module", "shard_id"), + [ + ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0), + ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1), + ("mtp.0.main_proj.scale", "model.main_proj", None), + ( + "mtp.0.ffn.shared_experts.w1.scale", + "model.layers.0.ffn.shared_experts.gate_up_proj", + 0, + ), + ( + "mtp.0.ffn.shared_experts.w2.scale", + "model.layers.0.ffn.shared_experts.down_proj", + None, + ), + ], +) +def test_v41_dspark_loads_linear_scales( + monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id +): + """Checkpoint ``.scale`` maps to the quant method's scale parameter and + loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static + loader (see tests/quantization/test_modelopt.py), not in load_weights.""" + from vllm.models.deepseek_v4_1.nvidia import dspark + + mxfp8 = scale_dtype != torch.float32 + scale_name = "weight_scale" if mxfp8 else "weight_scale_inv" + runtime_name = f"{runtime_module}.{scale_name}" + raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8) + checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float() + param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False) + shards = [] + + def load_scale(param, weight, *args): + shards.append(args) + assert weight.dtype == checkpoint_scale.dtype + param.copy_(weight) + + param.weight_loader = load_scale + draft = SimpleNamespace( + config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1), + quant_config=SimpleNamespace( + weight_block_size=[32, 32] if mxfp8 else [128, 128] + ), + linear_scale_name=scale_name, + pad_shared_expert=False, + model=SimpleNamespace( + layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + confidence_head=None, + ), + named_parameters=lambda: [(runtime_name, param)], + process_weights_after_loading=lambda: None, + ) + draft._remap_dspark_name = lambda name: ( + dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name) + ) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: [] + ) + + loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights( + draft, [(checkpoint_name, checkpoint_scale)] + ) + + assert loaded == {runtime_name} + assert shards == [() if shard_id is None else (shard_id,)] + torch.testing.assert_close(param, checkpoint_scale) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 30eb80a3f8..14d85fc09d 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -93,6 +93,7 @@ def test_registry_imports(model_arch): [ ("LlamaForCausalLM", False, False, "bi-encoder"), ("LlavaForConditionalGeneration", True, True, "bi-encoder"), + ("DeepseekV41ForCausalLM", True, False, "bi-encoder"), ("BertForSequenceClassification", False, False, "cross-encoder"), ("RobertaForSequenceClassification", False, False, "cross-encoder"), ("XLMRobertaForSequenceClassification", False, False, "cross-encoder"), diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 1c73573d8d..a91c99bf53 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -37,6 +37,7 @@ from vllm.model_executor.layers.quantization.fp8 import ( Fp8MoEMethod, ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, ) @@ -63,6 +64,383 @@ MODELS = [ ] +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA") +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) +def test_deepseek_v41_mxfp8_scale_loading( + dist_init, default_vllm_config, monkeypatch, tp_rank, scale_dtype +): + """Expand linear scales before TP slicing, including padded shared experts.""" + from vllm.model_executor.layers import linear as linear_module + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, + ) + from vllm.models.deepseek_v4_1 import quant_config as quant_module + from vllm.models.deepseek_v4_1.nvidia import model as model_module + + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + quant_config = quant_module.DeepseekV4FP8Config( + is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32] + ) + quant_config._resolved_expert_dtype = "fp4" + for module in (model_module, linear_module): + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank) + tp_args = dict(quant_config=quant_config, bias=False) + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args) + block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args) + block.ffn = torch.nn.Module() + block.ffn.shared_experts = torch.nn.Module() + block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear( + 128, [128, 128], **tp_args + ) + block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args) + block.ffn.experts = torch.nn.Module() + block.ffn.experts.register_parameter( + "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False) + ) + block.engram = torch.nn.Module() + block.engram.register_parameter( + "weight_scale_inv", + torch.nn.Parameter(torch.empty(3, 4, dtype=scale_dtype), False), + ) + + def load_expert_scale(param, weight, *args, **kwargs): + param.data.copy_(weight) + return True + + block.ffn.experts.weight_scale.weight_loader = load_expert_scale + model = model_module.DeepseekV4Model.__new__(model_module.DeepseekV4Model) + torch.nn.Module.__init__(model) + model.layers = torch.nn.ModuleList([block]) + model.config = SimpleNamespace(num_attention_heads=2) + model.quant_config = quant_config + model.use_sequence_parallel = False + model.get_expert_mapping = lambda: [ + ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1") + ] + + checkpoints = [ + ("attn.wq_b", "attn.wq_b", 128, 128, 0, None), + ("attn.wq_a", "attn.fused_wqa_wkv", 128, 128, 0, slice(0, 64)), + ("attn.wkv", "attn.fused_wqa_wkv", 64, 128, 0, slice(64, 96)), + ( + "ffn.shared_experts.w1", + "ffn.shared_experts.gate_up_proj", + 96, + 128, + 0, + slice(0, 64), + ), + ( + "ffn.shared_experts.w3", + "ffn.shared_experts.gate_up_proj", + 96, + 128, + 0, + slice(64, 128), + ), + ( + "ffn.shared_experts.down_proj", + "ffn.shared_experts.down_proj", + 128, + 96, + 1, + None, + ), + ] + weights = [] + expected = [] + for source, target, n, k, axis, shard in checkpoints: + weight = torch.randint(-4, 5, (n, k)).to(torch.float8_e4m3fn) + scale_bytes = torch.randint(124, 131, (n // 32, k // 32), dtype=torch.uint8) + weights.extend( + [ + (f"layers.0.{source}.weight", weight), + (f"layers.0.{source}.scale", scale_bytes.view(scale_dtype)), + ] + ) + dequant = ( + weight.float().reshape(n // 32, 32, k // 32, 32) + * torch.exp2(scale_bytes.float() - 127)[:, None, :, None] + ).reshape(n, k) + if "shared_experts" in source: + padded = torch.zeros(128, 128) + padded[:n, :k] = dequant + dequant = padded + expected.append((target, shard, dequant.chunk(2, dim=axis)[tp_rank])) + + expert_scale = torch.full((3, 4), 125, dtype=torch.uint8) + weights.append( + ("layers.0.ffn.experts.0.w1.weight_scale", expert_scale.view(scale_dtype)) + ) + weights.append(("layers.0.engram.weight_scale_inv", expert_scale.view(scale_dtype))) + mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale") + loaded = model.load_weights( + (name.removeprefix("model."), weight) for name, weight in mapper.apply(weights) + ) + assert "layers.0.ffn.experts.weight_scale" in loaded + assert torch.equal(block.ffn.experts.weight_scale, expert_scale) + assert torch.equal(block.engram.weight_scale_inv.view(torch.uint8), expert_scale) + for target, shard, reference in expected: + linear = block.get_submodule(target) + assert isinstance(linear.quant_method, ModelOptLinearMethod) + assert f"layers.0.{target}.weight_scale" in loaded + weight = linear.weight if shard is None else linear.weight[shard] + scale = linear.weight_scale if shard is None else linear.weight_scale[shard] + actual = ( + weight.float().unflatten(-1, (-1, 32)) + * torch.exp2(scale.float() - 127).unsqueeze(-1) + ).flatten(-2) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DSV4.1 requires CUDA") +@pytest.mark.parametrize( + "weight_block_size,expert_dtype,scale_name", + [ + ([32, 32], "fp4", "weight_scale"), + ([128, 128], "fp4", "weight_scale_inv"), + ([32, 32], "fp8", "weight_scale_inv"), + ], +) +def test_deepseek_v41_vl_mapper_routes_linear_scales( + weight_block_size, expert_dtype, scale_name +): + """The VL wrapper must map ``.scale`` keys to the parameter the linear + quant method registers, as the text model does. A hardcoded + ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" + from vllm.models.deepseek_v4_1.nvidia import model as model_module + from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module + + vllm_config = SimpleNamespace( + quant_config=SimpleNamespace(weight_block_size=weight_block_size) + ) + resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype) + assert resolved == scale_name + + mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved) + weight = torch.empty(0) + mapped = [ + name + for name, _ in mapper.apply( + [("layers.0.attn.wq_a.scale", weight), ("layers.0.attn.wkv.weight", weight)] + ) + ] + assert mapped == [ + f"language_model.model.layers.0.attn.wq_a.{scale_name}", + "language_model.model.layers.0.attn.wkv.weight", + ] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") +@pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) +@pytest.mark.parametrize( + "weight_shape,is_bmm", + [ + ((129, 160), False), + ((387, 160), True), + ((3, 129, 160), False), + ((3, 128, 512), False), + ], +) +def test_deepgemm_mxfp8_preserves_weight_and_scale_values( + scale_dtype, weight_shape, is_bmm +): + """DeepGEMM layout conversion preserves native weights and E8M0 scales.""" + from vllm.model_executor.layers.quantization.utils import fp8_utils + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( + 100 + ): + pytest.skip("DeepGEMM MXFP8 requires Blackwell") + + weight = torch.randn(weight_shape, device="cuda").to(torch.float8_e4m3fn) + scales = ( + torch.randint( + 1, + 255, + (weight.numel() // 32 + 1,), + device="cuda", + dtype=torch.uint8, + )[1:] + .view(*weight_shape[:-1], weight_shape[-1] // 32) + .view(scale_dtype) + ) + original_weight = weight.view(torch.uint8).clone() + original_scales = scales.view(torch.uint8).clone() + processed_weight, packed = fp8_utils.deepgemm_post_process_fp8_weight_block( + weight, + scales, + quant_block_shape=(1, 32), + use_e8m0=True, + is_bmm=is_bmm, + bmm_batch_size=3 if is_bmm else 0, + ) + assert packed.dtype == torch.int32 + unpacked = torch.stack( + [(packed >> (8 * i)) & 0xFF for i in range(4)], dim=-1 + ).flatten(-2) + unpacked = unpacked[..., : scales.shape[-1]].to(torch.uint8) + torch.testing.assert_close( + unpacked.reshape_as(original_scales), original_scales, rtol=0, atol=0 + ) + torch.testing.assert_close( + processed_weight.view(torch.uint8).reshape_as(original_weight), + original_weight, + rtol=0, + atol=0, + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") +@pytest.mark.parametrize("prequantized", [False, True]) +@pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"]) +@pytest.mark.parametrize("num_tokens", [1, 7, 128]) +def test_mxfp8_bmm_loads_and_projects_grouped_weights( + dist_init, default_vllm_config, prequantized, config_source, num_tokens +): + """BMM metadata set after construction selects grouped weight processing.""" + from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import ( + DeepGemmMxfp8BmmLinearKernel, + ) + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant, + ) + from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( + compute_fp8_einsum_recipe, + deep_gemm_fp8_o_proj, + ) + from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model + from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.utils.deep_gemm import is_deep_gemm_supported + + if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( + 100 + ): + pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell") + + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) + quant_config = DeepseekV4FP8Config( + is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32] + ) + quant_config._resolved_expert_dtype = "fp4" + if config_source == "mxfp8": + quant_config = get_quantization_config("mxfp8").from_config( + {"quant_method": "mxfp8"} + ) + with torch.device("cuda"): + linear = ColumnParallelLinear( + 512, + 256, + bias=False, + quant_config=quant_config, + return_bias=False, + ) + linear.is_bmm = True + linear.bmm_batch_size = 2 + weight = torch.randn(256, 512, device="cuda").to(torch.float8_e4m3fn) + scales = torch.randint(124, 131, (8, 16), device="cuda", dtype=torch.uint8) + checkpoint_scales = ( + scales.view(torch.float8_e8m0fnu) + if config_source == "deepseek" + else scales.repeat_interleave(32, dim=0) + ) + model = DeepseekV4Model.__new__(DeepseekV4Model) + torch.nn.Module.__init__(model) + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.wo_a = linear + model.layers = torch.nn.ModuleList([block]) + model.config = SimpleNamespace(num_attention_heads=2) + model.quant_config = quant_config + model.use_sequence_parallel = False + model.get_expert_mapping = lambda: [] + model.load_weights( + [ + ("layers.0.attn.wo_a.weight", weight), + ("layers.0.attn.wo_a.weight_scale", checkpoint_scales), + ] + ) + reference_weight = ( + weight.float().reshape(8, 32, 16, 32) + * torch.exp2(scales.float() - 127)[:, None, :, None] + ).reshape(2, 128, 512) + linear.quant_method.process_weights_after_loading(linear) + linear.quant_method.process_weights_after_loading(linear) + assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel) + assert linear.weight.shape == (2, 128, 512) + torch.testing.assert_close( + linear.weight.view(torch.uint8).flatten(), + weight.view(torch.uint8).flatten(), + rtol=0, + atol=0, + ) + unpacked_scales = torch.stack( + [(linear.weight_scale >> (8 * i)) & 0xFF for i in range(4)], dim=-1 + ).flatten(-2) + torch.testing.assert_close( + unpacked_scales.to(torch.uint8).reshape(256, 16), + scales.repeat_interleave(32, dim=0), + rtol=0, + atol=0, + ) + + x = torch.randn(2, num_tokens, 512, device="cuda", dtype=torch.bfloat16) + reference = torch.einsum("gmk,gnk->mgn", x.float(), reference_weight) + inputs = x.transpose(0, 1) + if prequantized: + cache = torch.cat( + ( + torch.ones(num_tokens, 32, device="cuda"), + torch.zeros(num_tokens, 32, device="cuda"), + ), + dim=1, + ) + inputs = fused_inv_rope_fp8_quant( + x.permute(1, 0, 2), + torch.arange(num_tokens, device="cuda"), + cache, + n_groups=2, + heads_per_group=1, + nope_dim=448, + rope_dim=64, + quant_group_size=32, + tma_aligned_scales=current_platform.has_device_capability(100), + ) + with torch.no_grad(): + output = linear(inputs) + assert output.shape == reference.shape + assert (output.float() - reference).norm() / reference.norm() < 0.06 + if prequantized: + recipe, tma_aligned_scales = compute_fp8_einsum_recipe(block_size=32) + projected = deep_gemm_fp8_o_proj( + x.permute(1, 0, 2), + torch.arange(num_tokens, device="cuda"), + cache, + linear, + torch.nn.Identity(), + n_groups=2, + heads_per_group=1, + nope_dim=448, + rope_dim=64, + o_lora_rank=128, + einsum_recipe=recipe, + tma_aligned_scales=tma_aligned_scales, + ) + torch.testing.assert_close(projected, output.flatten(1), rtol=0, atol=0) + compiled = torch.compile(linear, backend="eager", fullgraph=True) + with torch.no_grad(): + torch.testing.assert_close(compiled(inputs), output, rtol=0, atol=0) + + def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch): monkeypatch.setattr( flashinfer_utils, diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 59a333bfe4..60b207c184 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -127,6 +127,23 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head(): assert method.spec.activation is kNvfp4Dynamic +def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypatch): + """Standard MXFP8 checkpoints already have one scale row per weight row.""" + from vllm.model_executor.layers.linear import ReplicatedLinear + + kernel = Mock() + kernel.input_quant_key.return_value = None + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel", + lambda **kwargs: kernel, + ) + config = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"}) + linear = ReplicatedLinear(64, 64, bias=False, quant_config=config) + scales = torch.arange(128, dtype=torch.uint8).reshape(64, 2) + linear.weight_scale.weight_loader(linear.weight_scale, scales) + assert torch.equal(linear.weight_scale, scales) + + def test_modelopt_fp8_updates_weight_dims_after_transpose(): """Humming reads weight.input_dim/output_dim. Swapping the ModelWeightParameter for a plain Parameter drops them, so the per-tensor diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index c0ce30d214..51f7e70f5f 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -7,19 +7,61 @@ import pytest import torch from tests.v1.attention.utils import create_vllm_config +from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend +from vllm.models.deepseek_v4_1.sparse_mla import ( + DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, +) from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla.compressor_utils import ( CompressedSlotMappingKernel, ) from vllm.v1.attention.backends.mla.indexer import ( BuildPrefillChunkMetadataKernel, + DeepseekV4IndexerBackend, DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, ) from vllm.v1.attention.backends.mla.sparse_utils import ( ConvertReqIndexToGlobalIndexKernel, ) -from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + MLAAttentionSpec, + compute_layer_kv_cache_shape_bytes, +) from vllm.v1.worker.block_table import get_block_table_width +from vllm.v1.worker.utils import select_common_block_size + + +def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla(): + """Packed MLA/indexer groups must retain 64 compressed rows per page.""" + kernel_block_size = select_common_block_size( + 256, [DeepseekV4SparseMLABackend, DeepseekV4IndexerBackend] + ) + spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=4, + ) + assert compute_layer_kv_cache_shape_bytes(spec, 2, kernel_block_size) == ( + 2, + 1, + 64, + 132, + ) + + +def test_indexer_preserves_deepseek_v41_mla_block_size(): + """V4.1 retains its smaller pages independently of V4's C4 indexer.""" + block_size = DeepseekV41SparseMLABackend.get_supported_kernel_block_sizes()[0] + assert isinstance(block_size, int) + assert ( + select_common_block_size( + block_size, [DeepseekV41SparseMLABackend, DeepseekV41IndexerBackend] + ) + == block_size + ) def test_indexer_warmup_normalizes_zero_compress_ratios(): diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 5249de7f8c..2bdb5e74a7 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -398,6 +398,56 @@ class TestCSALinearGrouping: assert len(gdn) == 10 +class TestSlidingWindowBucketCap: + def test_sliding_window_bucket_is_capped_at_the_main_page(self): + """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an + unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would + set the block width and the MLA group would fill under a third of + every block; capped, it splits into groups no wider than the MLA + page.""" + config = _shared_layout_config() + specs: dict[str, KVCacheSpec] = {} + for layer, ratio in ((2, 2), (8, 2), (14, 2), (20, 1)): + specs[f"layers.{layer}.attn"] = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=ratio, + alignment=576, + ) + specs[f"layers.{layer}.idx"] = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=ratio, + alignment=576, + ) + for layer in range(43): + specs[f"layers.{layer}.swa"] = SlidingWindowMLASpec( + block_size=64, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=128, + alignment=576, + ) + groups = get_kv_cache_groups(config, specs) + pages = _pages(groups) + main = next(g for g in groups if "layers.2.attn" in g.layer_names) + main_bytes = sum(pages[n] for n in main.layer_names) + assert _get_kv_cache_bytes_per_block(groups) == main_bytes + swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")] + per_group = main_bytes // pages["layers.0.swa"] + assert len(swa_groups) == -(-43 // per_group) + for g in swa_groups: + assert sum(pages[n] for n in g.layer_names) <= main_bytes + assert sorted(n for g in swa_groups for n in g.layer_names) == sorted( + f"layers.{i}.swa" for i in range(43) + ) + + class TestDensePacking: def test_bytes_per_block_is_largest_group(self): groups, g1, g2 = _mixed_page_groups() @@ -548,3 +598,55 @@ class TestDensePacking: if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +class TestCompressorRingGroup: + def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): + """DeepSeek-V4.1 compressor state: a CircularBufferSpec ring (capacity + 8) beside 128-token paged MLA groups and 64-token SWA groups. The ring + claims one block per request, never hashes, and does not disturb the + scheduler/hash block sizes the paged groups imply.""" + config = _shared_layout_config() + config.cache_config.enable_prefix_caching = True + specs: dict[str, KVCacheSpec] = { + "layers.2.attn": MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=2, + alignment=576, + ), + "layers.2.attn.compressor.state_cache": CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=1024, + head_size_v=0, + dtype=torch.float32, + ), + } + for layer in range(4): + specs[f"layers.{layer}.attn.swa_cache"] = SlidingWindowMLASpec( + block_size=64, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=128, + alignment=576, + ) + groups = get_kv_cache_groups(config, specs) + ring_groups = [g for g in groups if "state_cache" in g.layer_names[0]] + assert len(ring_groups) == 1 + assert not ring_groups[0].kv_cache_spec.prefix_cacheable + kv_cache_config = get_kv_cache_config_from_groups( + config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + ) + assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([kv_cache_config]), + max_model_len=8192, + enable_caching=True, + hash_block_size=64, + scheduler_block_size=128, + ) + assert len(manager.coordinator.single_type_managers) == len(groups) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index dc9c1ffb30..ded83d800d 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -333,6 +333,57 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) +@pytest.mark.parametrize("keyword", [False, True]) +def test_quantized_activation_replay_does_not_retain_tensors( + cuda_capture_stream, keyword +): + """Release the producer tensors while preserving both buffers on replay.""" + import gc + import weakref + + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic + + @eager_break_during_capture + def consume(x, data_out, scale_out): + data_out.copy_(x.data) + scale_out.copy_(x.scale) + + source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16) + data_out = source.to(torch.float8_e4m3fn) + scale_out = source[:, ::32].to(torch.uint8) + output = torch.empty_like(source) + cap = BreakableCUDAGraphCapture() + with cap: + data = source.to(torch.float8_e4m3fn) + scale = source[:, ::32].to(torch.uint8) + activation = QuantizedActivation( + data, scale, source.dtype, source.shape, kMxfp8Dynamic + ) + data_ref, scale_ref = weakref.ref(data), weakref.ref(scale) + if keyword: + consume(x=activation, data_out=data_out, scale_out=scale_out) + else: + consume(activation, data_out, scale_out) + output.copy_(data_out) + + del data, scale, activation + gc.collect() + assert data_ref() is None + assert scale_ref() is None + + for value in (1, 2, 3): + source.fill_(value) + cap.replay() + cuda_capture_stream.synchronize() + torch.testing.assert_close(output, source, rtol=0, atol=0) + torch.testing.assert_close(scale_out, source[:, ::32].to(torch.uint8)) + + def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): """Handle an eager attention break inside a multi-stream overlap region.""" from vllm.compilation.breakable_cudagraph import ( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index d6b0becc1d..e6b7302429 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( ) from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KpoolTailSpec, KVCacheConfig, @@ -798,3 +799,101 @@ def test_worker_setup_tolerates_finer_scratch_group(): # The scratch group's namespace never enters the store. scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0] assert not any(key.startswith(scratch_prefix) for key in store._data) + + +def test_ring_scratch_group_is_never_stored_and_does_not_block_hits(): + """A per-request ring group (CircularBufferSpec, capacity as block_size) + sits beside the paged group: worker setup tolerates its block size (the + DeepSeek-V4.1 compressor ring is 8 rows), saving a request stores paged + blocks only, and lookup hits on the paged group alone.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + ring = CircularBufferSpec( + block_size=8, num_kv_heads=1, head_size=64, head_size_v=0, dtype=torch.uint8 + ) + cfg = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=4 * full.page_size_bytes, + layers=["L0"], + layer_stride=4 * full.page_size_bytes, + block_stride=full.page_size_bytes, + ), + KVCacheTensor( + size=4 * ring.page_size_bytes, + layers=["L1"], + layer_stride=4 * ring.page_size_bytes, + block_stride=ring.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], ring), + ], + ) + vllm_config = _minimal_vllm_config(cache_block_size=16) + store = _DictStore() + + worker = _build_worker_with_dict_store(vllm_config, cfg, store) + worker.tp_size = 1 + worker.pp_size = 1 + worker.num_kv_head = 8 + assert worker.token_dbs[1].hash_block_size == 8 + + raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8) + raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8) + kv_caches = { + "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0], + "L1": dense_kv_cache_views(raw_ring, ring, 4, 1, KVCacheLayout.LBNHC)[0], + } + + def _fake_thread_init(*args, **kwargs): + for v in list(args) + list(kwargs.values()): + if isinstance(v, threading.Event): + v.set() + m = MagicMock() + m.start = lambda: None + return m + + with ( + patch.object( + mooncake_store_worker, + "KVCacheStoreSendingThread", + side_effect=_fake_thread_init, + ), + patch.object( + mooncake_store_worker, + "KVCacheStoreRecvingThread", + side_effect=_fake_thread_init, + ), + ): + worker.register_kv_caches(kv_caches) + + send_thread = KVCacheStoreSendingThread( + store=store, + token_databases=worker.token_dbs, + block_size=worker.block_size, + coord=worker.coord, + tp_rank=worker.tp_rank, + group_put_steps=worker._group_tp_replication_factors, + kv_role=worker.kv_role, + ready_event=threading.Event(), + enable_kv_event=False, + group_participates=[True, False], + ) + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)] + save_req = ReqMeta( + req_id="r0", + token_len_chunk=64, + # The ring group holds one block for the request's lifetime. + block_ids=([1, 2, 3, 4], [1]), + block_hashes=hs, + can_save=True, + store_job_id=1, + ) + send_thread.add_request(save_req) + send_thread._handle_request(send_thread.request_queue.get()) + worker.store = store + + assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64 + assert not any("@group:1" in key for key in store._data) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index a9f8007398..32908f78ad 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -1251,3 +1251,112 @@ def test_csa_linear_remote_ple_is_copied_whole(): tp_ratio=-2, transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1), ) + + +def _make_ring_worker(): + """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a + model whose sliding-window KV lives in per-request rings.""" + from unittest.mock import MagicMock + + from vllm.config import set_current_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + paged_specs = { + f"paged.{index}": MLAAttentionSpec( + block_size=4, num_kv_heads=1, head_size=64, dtype=torch.float16 + ) + for index in range(2) + } + ring_specs = { + f"ring.{index}": CircularBufferSpec( + block_size=6, + num_kv_heads=1, + head_size=32, + head_size_v=0, + dtype=torch.float16, + ) + for index in range(2) + } + + def group(specs): + uniform = UniformTypeKVCacheSpecs.from_specs(specs) + assert uniform is not None + return KVCacheGroupSpec(list(specs), uniform) + + # Both rings overlay the second paged region, so the scratch group must + # address only that region while the paged group spans both. + tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1")) + region_size, page_size = 512, 256 + kv_cache_config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[ + KVCacheTensor( + size=len(tensor_regions) * region_size, + layers=[layer_name], + layer_stride=region_size, + block_stride=page_size, + offset=region_index * region_size, + ) + for region_index, layer_names in enumerate(tensor_regions) + for layer_name in layer_names + ], + kv_cache_groups=[group(paged_specs), group(ring_specs)], + ) + + vllm_config = create_vllm_config(block_size=4) + vllm_config.cache_config.enable_prefix_caching = False + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [4] + fake_backend.get_name.return_value = "TEST_ATTN" + fake_backend.full_cls_name.return_value = "test.AttentionBackend" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(2)] + worker.register_kv_caches( + { + layer_name: tensors[region_index] + for region_index, layer_names in enumerate(tensor_regions) + for layer_name in layer_names + } + ) + return worker + + +@pytest.mark.cpu_test +def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions(): + worker = _make_ring_worker() + + assert worker._ple_group_index is None + assert worker._scratch_region_indices == [1] + + desc_ids = worker._compute_desc_ids( + block_ids=([1], [0]), + dst_num_blocks=2, + block_size_ratio=None, + physical_blocks_per_logical=1, + ) + # Paged block 1 in both regions; ring block 0 only in its scratch region. + assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0]) diff --git a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py index 69f9c2da41..af4c7858fc 100644 --- a/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py +++ b/tests/v1/kv_connector/unit/test_simple_cpu_offload_connector.py @@ -19,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector i SimpleCPUOffloadConnector, ) from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, @@ -90,3 +91,59 @@ def test_disk_mode_rejects_non_positive_capacity() -> None: "disk_capacity_bytes": 0, } ) + + +def _make_ring_kv_cache_config(num_blocks: int = 16) -> KVCacheConfig: + """A paged full-attention group beside a per-request ring group.""" + full = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ) + ring = CircularBufferSpec( + block_size=6, num_kv_heads=1, head_size=32, head_size_v=0, dtype=DTYPE + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=_BYTES_PER_BLOCK * num_blocks, + layers=["attn"], + layer_stride=_BYTES_PER_BLOCK * num_blocks, + block_stride=_BYTES_PER_BLOCK, + ), + KVCacheTensor( + size=ring.page_size_bytes * num_blocks, + layers=["ring"], + layer_stride=ring.page_size_bytes * num_blocks, + block_stride=ring.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["attn"], full), + KVCacheGroupSpec(["ring"], ring), + ], + ) + + +def test_ring_scratch_group_is_accepted_and_never_offloaded() -> None: + """The ring group's capacity is not a token granularity, it holds no + hashed blocks, and it never participates in CPU hits.""" + vllm_config = create_vllm_config( + kv_connector="SimpleCPUOffloadConnector", + kv_role="kv_both", + kv_connector_extra_config={"cpu_bytes_to_use": _BYTES_PER_BLOCK * 8}, + ) + connector = SimpleCPUOffloadConnector( + vllm_config=vllm_config, + role=KVConnectorRole.SCHEDULER, + kv_cache_config=_make_ring_kv_cache_config(), + ) + scheduler = connector.scheduler_manager + assert scheduler.fa_gidx == 0 + cpu_groups = scheduler.cpu_kv_cache_config.kv_cache_groups + assert [g.kv_cache_spec.prefix_cacheable for g in cpu_groups] == [True, False] + # Only the paged group takes part in CPU-side hit lookup. + attention_groups = scheduler.cpu_coordinator.attention_groups + assert [g.group_ids for g in attention_groups] == [[0]] diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 1952e3a4b0..901e076e09 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1319,12 +1319,10 @@ def test_inflight_finish_deferred_cleanup() -> None: # Test 8: Null GPU blocks are skipped in store and load transfer pairs # --------------------------------------------------------------------------- def test_multi_group_null_blocks_skipped() -> None: - """Null GPU blocks (no block_hash) must not appear in store or load pairs. + """The null block id never appears in store or load pairs. - In eager store mode, _prepare_eager_store_specs skips blocks whose - block_hash is None (null blocks have no hash). We verify this by mixing - real hashed blocks with unhashed (null-like) blocks in a single group and - checking that only real blocks appear in the store list. + A null block-table slot whose hash is still cached on the GPU is stored + from the cached block instead; the null block itself is never copied. """ fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False) sched = fix.scheduler @@ -1356,9 +1354,11 @@ def test_multi_group_null_blocks_skipped() -> None: f"Null block id {null_block_id} should not appear in store transfer pairs" ) - # Only real block should be scheduled for store - assert len(meta.store_gpu_blocks) == 1 + # The nulled slot's hash is still cached on the GPU, so that block is + # recovered from the pool and stored alongside the real one. + assert len(meta.store_gpu_blocks) == 2 assert gpu_blocks[0].block_id in meta.store_gpu_blocks + assert gpu_blocks[1].block_id in meta.store_gpu_blocks # Complete the store assert meta.store_event >= 0 @@ -1374,13 +1374,12 @@ def test_multi_group_null_blocks_skipped() -> None: block_hasher=req._block_hasher, ) hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) - # Only 1 block was stored (the real one) - assert hit_tokens == BLOCK_SIZE + assert hit_tokens == 2 * BLOCK_SIZE assert is_async is True # Allocate new GPU blocks for the load - gpu_blocks2 = gpu_pool.get_new_blocks(1) - kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],)) + gpu_blocks2 = gpu_pool.get_new_blocks(2) + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,)) sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens) sched_out2 = make_scheduler_output({req2.request_id: 1}) @@ -2800,3 +2799,47 @@ def test_boundary_handoff_keeps_block_meta_index_parallel() -> None: assert req_ids == [req.request_id] assert block_meta is not None assert len(block_meta) == len(gpu_ids) == len(cpu_ids) + + +def test_finished_eager_store_recovers_nulled_window_tail() -> None: + """A sliding-window group nulls pages that left the window before the + connector sees the finished request's block table, while the GPU keeps + them hashed in its free queue. The eager store must copy those pages, or + the CPU lookup (which needs the whole window) misses every time.""" + fix = make_scheduler(num_cpu_blocks=16, num_gpu_blocks=32, num_groups=2, lazy=False) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + null_block = gpu_pool.null_block + + num_blocks = 4 # == sliding window of the test SWA group + req = make_request(num_blocks=num_blocks) + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0) + swa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=1) + # remove_skipped_blocks() ran: the two oldest window pages are freed (still + # hashed) and their slots hold the null block. + gpu_pool.free_blocks(swa_blocks[:2]) + swa_table = [null_block, null_block, *swa_blocks[2:]] + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, swa_table)) + req.num_computed_tokens = num_blocks * BLOCK_SIZE + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) + + block_ids = kv_blocks.get_block_ids() + sched.request_finished_all_groups(req, block_ids) + meta = sched.build_connector_meta(make_scheduler_output({})) + + assert null_block.block_id not in meta.store_gpu_blocks + assert set(meta.store_gpu_blocks) == {b.block_id for b in (*fa_blocks, *swa_blocks)} + assert meta.store_event >= 0 + simulate_store_completion(sched, meta.store_event) + + req2 = Request( + request_id="req-window-tail-load", + prompt_token_ids=req.prompt_token_ids, + sampling_params=req.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req._block_hasher, + ) + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens == num_blocks * BLOCK_SIZE + assert is_async is True diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 6b63588150..2e1e8fc619 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2143,13 +2143,15 @@ class rocm_aiter_ops: if not current_platform.is_rocm(): return - from vllm.platforms.rocm import on_gfx11 + from vllm.platforms.rocm import on_gfx11, on_gfx950 - if on_gfx11() and not _OPS_REGISTERED: + # This op has self-contained Triton/C++ implementations on gfx11 + # and gfx950. Only its optional top-k fast path comes from aiter. + if (on_gfx11() or on_gfx950()) and not _OPS_REGISTERED: direct_register_custom_op( op_name="rocm_aiter_sparse_attn_indexer", op_func=rocm_aiter_sparse_attn_indexer, - mutates_args=["topk_indices_buffer"], + mutates_args=["topk_indices_buffer", "candidate_blocks"], fake_impl=rocm_aiter_sparse_attn_indexer_fake, dispatch_key=current_platform.dispatch_key, ) @@ -2327,7 +2329,7 @@ class rocm_aiter_ops: direct_register_custom_op( op_name="rocm_aiter_sparse_attn_indexer", op_func=rocm_aiter_sparse_attn_indexer, - mutates_args=["topk_indices_buffer"], + mutates_args=["topk_indices_buffer", "candidate_blocks"], fake_impl=rocm_aiter_sparse_attn_indexer_fake, dispatch_key=current_platform.dispatch_key, ) diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 343b9ef398..8e25d05c28 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -42,6 +42,7 @@ from vllm.forward_context import ( is_forward_context_available, ) from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors @@ -56,6 +57,12 @@ def is_breakable_cudagraph_enabled() -> bool: F = TypeVar("F", bound=Callable[..., Any]) +def _weak_ref_capture_arg(arg: Any) -> Any: + if isinstance(arg, QuantizedActivation): + return arg.weak_ref() + return weak_ref_tensor(arg) + + def eager_break_during_capture(fn: F) -> F: """Decorator that turns a custom-op Python kernel into a "break point" for the breakable cudagraph capture. @@ -105,13 +112,8 @@ def eager_break_during_capture(fn: F) -> F: # Weak-ref args: strong refs in the replay lambda pin cudagraph-pool # slots across batch descriptors. cudagraph owns the slot, so the # weak_ref is safe to deref on replay. - weak_args = tuple( - weak_ref_tensor(a) if isinstance(a, torch.Tensor) else a for a in args - ) - weak_kwargs = { - k: weak_ref_tensor(v) if isinstance(v, torch.Tensor) else v - for k, v in kwargs.items() - } + weak_args = tuple(_weak_ref_capture_arg(a) for a in args) + weak_kwargs = {k: _weak_ref_capture_arg(v) for k, v in kwargs.items()} return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs)) return wrapper # type: ignore[return-value] diff --git a/vllm/config/engram.py b/vllm/config/engram.py index b46311315f..2aede1868d 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -15,6 +15,14 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Architecture -> the hf_text_config field naming its n-gram layers. A model is +# only configurable here if it actually has such layers to store. +_NGRAM_LAYER_FIELDS = { + "DeepseekV41ForCausalLM": "engram_layer_ids", + "Qwen4ExpForCausalLM": "ple_layer_ids", + "Qwen4ExpForConditionalGeneration": "ple_layer_ids", +} + def _default_cpu_offload() -> bool: """Honor the legacy environment variable only when the field is omitted.""" @@ -40,23 +48,25 @@ class EngramConfig: Otherwise, each DP rank has a separate TP-sharded embedding replica.""" def verify_model_config(self, model_config: "ModelConfig | None") -> None: - """Reject Engram configuration for models without supported embeddings.""" + """Reject Engram configuration for models without n-gram embeddings.""" from vllm.platforms import current_platform - supported_architectures = { - "Qwen4ExpForCausalLM", - "Qwen4ExpForConditionalGeneration", - } + field = ( + _NGRAM_LAYER_FIELDS.get(model_config.architecture) + if model_config is not None + else None + ) if ( model_config is None - or model_config.architecture not in supported_architectures + or field is None or not current_platform.is_cuda() - or not getattr(model_config.hf_text_config, "ple_layer_ids", None) + or not getattr(model_config.hf_text_config, field, None) ): raise ValueError( "EngramConfig requires a model with supported Engram " - "embeddings. Currently only the CUDA Qwen4Exp implementation " - "with non-empty ple_layer_ids is supported." + "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek " + "V4.1 implementations with non-empty n-gram layer ids are " + "supported." ) def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 1312cbb241..a6dddc3178 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -165,6 +165,7 @@ FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( { "DeepseekV4ForCausalLM", "DeepSeekV4MTPModel", + "DeepseekV41ForCausalLM", } ) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index d4085d7d52..ae56e58803 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -679,12 +679,18 @@ class SpeculativeConfig: ], } ) - if hf_config.model_type == "deepseek_v4": + if hf_config.model_type in ("deepseek_v4", "deepseek_v41"): + # V4.1 has no classic-MTP draft: its checkpoints ship DSpark stages + # under ``mtp.*``, so only V4 gets an MTP architecture here. The + # DSpark path rewrites ``architectures`` itself and needs only + # ``n_predict``; ``method="mtp"`` on V4.1 is rejected below. + is_v41 = hf_config.model_type == "deepseek_v41" hf_config.model_type = "deepseek_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", None) - hf_config.update( - {"n_predict": n_predict, "architectures": ["DeepSeekV4MTPModel"]} - ) + overrides = {"n_predict": n_predict} + if not is_v41: + overrides["architectures"] = ["DeepSeekV4MTPModel"] + hf_config.update(overrides) if hf_config.model_type in ("pangu_ultra_moe"): hf_config.model_type = "pangu_ultra_moe_mtp" if hf_config.model_type == "pangu_ultra_moe_mtp": @@ -1335,6 +1341,18 @@ class SpeculativeConfig: MTPModelTypes ): self.method = "mtp" + if ( + self.target_model_config is not None + and self.target_model_config.hf_config.model_type + == "deepseek_v41" + ): + raise ValueError( + "DeepSeek V4.1 has no classic-MTP draft: its " + "checkpoints ship DSpark stages under mtp.* " + "(main_proj/markov_head/confidence_head) and carry " + "no e_proj/h_proj/enorm/hnorm/hc_head weights. Use " + "speculative method 'dspark' instead of 'mtp'." + ) if ( self.num_speculative_tokens > 1 and self.draft_model_config.hf_config.model_type @@ -1398,12 +1416,26 @@ class SpeculativeConfig: and "Gemma4DSparkModel" not in self.draft_model_config.architectures and "K3DSparkModel" not in self.draft_model_config.architectures ): - # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config + # DeepSeek-V4(.1) DSpark reuses the full target config # and its weights ship in the target checkpoint. - self.draft_model_config.hf_config.model_type = "deepseek_v4" - self.draft_model_config.hf_config.architectures = [ - "DSparkDraftModel" + is_v41 = ( + self.target_model_config.hf_config.model_type == "deepseek_v41" + ) + draft_hf_config = self.draft_model_config.hf_config + draft_hf_config.model_type = ( + "deepseek_v41" if is_v41 else "deepseek_v4" + ) + draft_hf_config.architectures = [ + "DSparkV41DraftModel" if is_v41 else "DSparkDraftModel" ] + if is_v41: + # hf_config_override set n_predict to the number of + # MTP stages (3), but one DSpark round drafts + # dspark_block_size tokens; num_speculative_tokens + # divisibility is checked against n_predict below. + draft_hf_config.n_predict = getattr( + draft_hf_config, "dspark_block_size", None + ) or getattr(draft_hf_config, "n_predict", None) self.draft_model_config.quantization = ( self.target_model_config.quantization ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5bd6af20de..39c2cb4d52 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -81,6 +81,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( "DeepseekV4ForCausalLM", "DeepseekV4ForConditionalGeneration", "DeepSeekV4MTPModel", + "DeepseekV41ForCausalLM", "Dots3NoteForCausalLM", "Dots3NoteMTPModel", "Glm5NextForCausalLM", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 3f70082013..67a935d7ee 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -91,6 +91,7 @@ class MooncakeStoreCoordinator: assert all( scheduler_block_size % g.kv_cache_spec.block_size == 0 for g in kv_cache_groups + if g.kv_cache_spec.prefix_cacheable ), "scheduler_block_size must be a multiple of each group's block_size" self.kv_cache_groups = kv_cache_groups self.mamba_group_ids = { diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index ceef10e9fd..dd2e97c965 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -460,11 +460,13 @@ class NixlBaseConnectorWorker: for spec in iter_layer_specs(group.kv_cache_spec) ) ] - if len(ple_groups) != 1 or len(ple_groups[0][1].layer_names) != 1: + if len(ple_groups) > 1 or any( + len(group.layer_names) != 1 for _, group in ple_groups + ): raise ValueError( - "CSA-linear NIXL requires exactly one PLE cache owner." + "CSA-linear NIXL requires at most one PLE cache owner." ) - self._ple_group_index = ple_groups[0][0] + self._ple_group_index = ple_groups[0][0] if ple_groups else None if self._has_mamba: assert self._is_hma_required diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index 4be2f4f9ff..bd6cab668e 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -8,7 +8,7 @@ expose_input_quant_key; the kernel validates and reads the activation via as_quantized_activation. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace import torch @@ -34,6 +34,16 @@ class QuantizedActivation: orig_shape: torch.Size quant_key: QuantKey + def weak_ref(self) -> "QuantizedActivation": + """Return a copy with non-owning tensor references for CUDA graph replay.""" + from vllm.utils.torch_utils import weak_ref_tensor + + return replace( + self, + data=weak_ref_tensor(self.data), + scale=weak_ref_tensor(self.scale), + ) + def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: """Advertise the kernel's pre-quantized input key on the layer, if any. diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index a2c65738e2..fd1d047074 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -115,7 +115,20 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # lazy import to avoid triggering `torch.compile` too early from vllm.config.quantization import _ONLINE_SHORTHANDS from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig - from vllm.models.deepseek_v4 import DeepseekV4FP8Config + + if current_platform.is_cuda() or current_platform.is_rocm(): + # The v4.1 class is the v4 one extended to accept model_type + # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has + # platform-specific implementations for both CUDA and ROCm. + from vllm.models.deepseek_v4_1 import ( + DeepseekV4FP8Config as DeepseekV41FP8Config, + ) + + deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config + else: + from vllm.models.deepseek_v4 import DeepseekV4FP8Config + + deepseek_config = DeepseekV4FP8Config from .auto_awq import AutoAWQConfig from .auto_gptq import AutoGPTQConfig @@ -161,7 +174,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "inc": INCConfig, "mxfp4": Mxfp4Config, "gpt_oss_mxfp4": GptOssMxfp4Config, - "deepseek_v4_fp8": DeepseekV4FP8Config, + "deepseek_v4_fp8": deepseek_config, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 199d856e3b..e42f43b39c 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -937,14 +937,13 @@ def w8a8_triton_block_scaled_mm( assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] - # Triton cannot currently bind E8M0 scale tensors directly. On ROCm, - # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format, - # so decode them to fp32 before launching the kernel. - if current_platform.is_rocm() or current_platform.is_xpu(): - if As.dtype == torch.float8_e8m0fnu: - As = _upcast_e8m0_to_fp32(As).contiguous() - if Bs.dtype == torch.float8_e8m0fnu: - Bs = _upcast_e8m0_to_fp32(Bs).contiguous() + # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints + # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32 + # before launching the kernel. + if As.dtype == torch.float8_e8m0fnu: + As = _upcast_e8m0_to_fp32(As).contiguous() + if Bs.dtype == torch.float8_e8m0fnu: + Bs = _upcast_e8m0_to_fp32(Bs).contiguous() assert A.shape[-1] == B.shape[-1] assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() @@ -1089,8 +1088,6 @@ def deepgemm_post_process_weight_scale_block( is_sfa: bool = False, ) -> torch.Tensor: if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): - # Scales already in E8M0 from checkpoint; upcast to fp32 and let - # DeepGEMM pack the layout expected by the target architecture. ws = _upcast_e8m0_to_fp32(ws) else: assert ws.dtype == torch.float32, ( @@ -1121,12 +1118,7 @@ def deepgemm_post_process_fp8_weight_block( f"to be torch.float8_e4m3fn, got {wq.dtype} instead." ) - if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): - # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0 - # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization - # (weights already have power-of-two scales). - ws = _upcast_e8m0_to_fp32(ws) - else: + if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8): assert ws.dtype == torch.float32, ( f"Expected tensor scales dtype to be torch.float32 or " f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead" @@ -1137,7 +1129,7 @@ def deepgemm_post_process_fp8_weight_block( if is_bmm: # Reshape 2D weight/scale to 3D for grouped BMM (einsum): # wq: (g*r, d) -> (g, r, d) - # ws: (g*r/128, d/128) -> (g, r/128, d/128) + # Scale dimensions follow quant_block_shape. g = bmm_batch_size assert wq.ndim == 2 and ws.ndim == 2 d = wq.size(1) diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index 0c9324a549..bfcc69e4b1 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -126,8 +126,14 @@ def _mxfp8_quant_triton_kernel(): amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY) # [BLOCK_M] sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) - descale = tl.exp2(sb - 127.0) - xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty) + # Scale by the reciprocal rather than dividing by ``exp2(sb - 127)``: + # sb == 0 (an all-zero / denormal block) makes that divisor 2**-127, + # which is subnormal in fp32 and flushes to zero on CDNA, turning the + # block's zeros into 0/0 == NaN. The reciprocal is normal for every + # reachable sb, and both forms are exact powers of two, so nothing + # else changes. + rescale = tl.exp2(127.0 - sb) + xq = (x * rescale[:, None]).to(xq_ptr.dtype.element_ty) tl.store( xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, xq, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 485d15354d..bbbe98bc8d 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -330,7 +330,7 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): quant_config = getattr(model_config.hf_config, "quantization_config", None) if quant_config is not None and quant_config.get("quant_method") == "fp8": model_type = getattr(model_config.hf_config, "model_type", None) - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v41"): model_config.hf_config.quantization_config["quant_method"] = ( "deepseek_v4_fp8" ) @@ -343,7 +343,7 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): and hf_text_quant_config.get("quant_method") == "fp8" ): model_type = getattr(model_config.hf_text_config, "model_type", None) - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v41"): model_config.hf_text_config.quantization_config["quant_method"] = ( "deepseek_v4_fp8" ) @@ -1003,6 +1003,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "ColQwen3_5": ColQwen3_5Config, "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig, + "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index a9fe7c2ec2..69dd88f349 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -376,6 +376,10 @@ _MULTIMODAL_MODELS = { "vllm.models.deepseek_v4", "DeepseekV4ForConditionalGeneration", ), + "DeepseekV41ForCausalLM": ( + "vllm.models.deepseek_v4_1", + "DeepseekV41ForCausalLM", + ), "Dots3NoteForCausalLM": ( "vllm.models.dots3_note", "Dots3NoteForCausalLM", @@ -640,6 +644,10 @@ _SPECULATIVE_DECODING_MODELS = { "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), + "DSparkV41DraftModel": ( + "vllm.models.deepseek_v4_1", + "DSparkDeepseekV4ForCausalLM", + ), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "K3DSparkModel": ( diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index 6ca8c94fff..fc7731509d 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -173,7 +173,10 @@ def deepseek_v4_mhc_warmup( # large checkpoints; bail out for any model that is not DeepSeek V4. config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if config is not None else None - if model_type is not None and model_type != "deepseek_v4": + if model_type is not None and model_type not in ( + "deepseek_v4", + "deepseek_v41", + ): return layer = _find_first_mhc_layer(model) diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 80dfd51996..5209b94434 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -28,15 +28,20 @@ _DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( "FLASHMLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV4", "ROCM_FLASHMLA_SPARSE_DSV4", + "FLASHMLA_SPARSE_DSV41", + "FLASHINFER_MLA_SPARSE_DSV41", "DEEPSEEK_SPARSE_SWA", } ) _FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"}) -_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"}) +_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset( + {"FLASHINFER_MLA_SPARSE_DSV4", "FLASHINFER_MLA_SPARSE_DSV41"} +) _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = { "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2", "FLASHINFER_MLA_SPARSE_DSV4": "DSv4", + "FLASHINFER_MLA_SPARSE_DSV41": "DSv4.1", } _SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v4_1/attention.py index 77e98a376e..9bb53526c9 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v4_1/attention.py @@ -58,7 +58,7 @@ from vllm.utils.multi_stream_utils import ( ) from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV4IndexerBackend, + DeepseekV41IndexerBackend, dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) @@ -1005,7 +1005,7 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: - return DeepseekV4IndexerBackend + return DeepseekV41IndexerBackend class DeepseekV4Indexer(nn.Module): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 5ec07ad211..c20eaab0e5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -245,9 +245,20 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows. return [256] +class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_INDEXER" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [64 if current_platform.is_device_capability_family(90) else 128] + + @dataclass class DeepseekV32IndexerPrefillChunkMetadata: block_table: torch.Tensor @@ -1239,9 +1250,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) block_table = self.indexer_decode_block_table_buffer[:rows, :cols] - seq_lens_is_buffer_view = (use_native and next_n > 1) or ( - not use_native and max_decode_len > 1 - ) + # Flattening always returns a buffer view, including single-token + # batches. Keep its address stable across varlen graph replays. + seq_lens_is_buffer_view = not use_native or next_n > 1 # DCP: localize the now-expanded per-token global bounds to this # rank's owned KV. Done here (after expansion) so each token's global diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 45a381d4a3..dd2392d6c2 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -11,7 +11,7 @@ import torch.nn.functional as F import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CUDAGraphMode +from vllm.config import CUDAGraphMode, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.platforms import current_platform @@ -819,6 +819,32 @@ def rocm_fp8_mqa_logits( return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) +def _max_decode_logits_rows(num_batched_tokens: int) -> int: + """Upper bound on decode rows the paged-MQA logits buffer can ever hold. + + ``rocm_fp8_paged_mqa_logits`` sizes its workspace as + ``(batch_size * next_n, max_model_len)``. ``batch_size`` is bounded by + ``max_num_seqs`` and ``next_n`` by ``1 + num_speculative_tokens``, which is + far tighter than ``max_num_batched_tokens`` -- 192 vs 16384 for a typical + 32-seq DSpark-5 deployment. The loose bound is harmless at short contexts + but scales with ``max_model_len``, so at the model's full context it asks + for tens of TiB and the engine cannot start. Take whichever valid bound is + smaller; the workspace is locked after profiling, so it must not be under- + estimated. + """ + try: + vllm_config = get_current_vllm_config() + except Exception: + return num_batched_tokens + scheduler_config = getattr(vllm_config, "scheduler_config", None) + max_num_seqs = getattr(scheduler_config, "max_num_seqs", None) + if not max_num_seqs: + return num_batched_tokens + speculative_config = getattr(vllm_config, "speculative_config", None) + num_spec = getattr(speculative_config, "num_speculative_tokens", 0) or 0 + return min(num_batched_tokens, max_num_seqs * (1 + num_spec)) + + def rocm_aiter_sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: LayerNameType, @@ -835,6 +861,9 @@ def rocm_aiter_sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool = False, compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, ) -> torch.Tensor: return topk_indices_buffer @@ -856,6 +885,9 @@ def rocm_aiter_sparse_attn_indexer( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool = False, compress_ratio: int = 1, + candidate_blocks: torch.Tensor | None = None, + candidate_block_size: int = 0, + candidate_write: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -881,15 +913,15 @@ def rocm_aiter_sparse_attn_indexer( ) # Decode logits buffer, used by rocm_fp8_paged_mqa_logits. - # batch_size * next_n <= hidden_states.shape[0] == max_num_batched_tokens + decode_rows = _max_decode_logits_rows(hidden_states.shape[0]) if _ON_GFX942 or _ON_GFX950: workspace_manager.get_simultaneous( - ((hidden_states.shape[0], max_model_len), torch.float32), + ((decode_rows, max_model_len), torch.float32), ) else: workspace_manager.get_simultaneous( ( - (q_fp8.shape[1], hidden_states.shape[0], max_model_len), + (q_fp8.shape[1], decode_rows, max_model_len), torch.float32, ), ) @@ -919,6 +951,9 @@ def rocm_aiter_sparse_attn_indexer( topk_indices_buffer, skip_k_cache_insert, compress_ratio, + candidate_blocks, + candidate_block_size, + candidate_write, ) layer_attn_metadata = attn_metadata[k_cache_prefix] assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata) @@ -974,6 +1009,30 @@ def rocm_aiter_sparse_attn_indexer( chunk.cu_seqlen_ks, chunk.cu_seqlen_ke, ) + if candidate_blocks is not None: + from vllm.model_executor.layers.sparse_attn_indexer import ( + _apply_candidate_mask, + _select_candidate_blocks, + ) + + chunk_candidates = candidate_blocks[chunk.token_start : chunk.token_end] + if candidate_write: + _select_candidate_blocks( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + chunk_candidates.shape[1], + candidate_block_size, + chunk_candidates, + ) + else: + _apply_candidate_mask( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + chunk_candidates, + candidate_block_size, + ) topk_indices = topk_indices_buffer[ chunk.token_start : chunk.token_end, :topk_tokens ] @@ -1041,6 +1100,37 @@ def rocm_aiter_sparse_attn_indexer( max_model_len=max_model_len, ) + if candidate_blocks is not None: + from vllm.model_executor.layers.sparse_attn_indexer import ( + _apply_candidate_mask, + _select_candidate_blocks, + ) + + num_rows = logits.shape[0] + visible = decode_metadata.seq_lens.reshape(-1) + if visible.numel() != num_rows: + visible = visible.repeat_interleave(next_n) + visible = visible[:num_rows].to(torch.int64) + row_starts = torch.zeros_like(visible) + decode_candidates = candidate_blocks[:num_rows] + if candidate_write: + _select_candidate_blocks( + logits, + row_starts, + visible, + decode_candidates.shape[1], + candidate_block_size, + decode_candidates, + ) + else: + _apply_candidate_mask( + logits, + row_starts, + visible, + decode_candidates, + candidate_block_size, + ) + topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] num_rows = logits.shape[0] @@ -1097,6 +1187,10 @@ def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor: ) return _upcast_e8m0_to_fp32(scale).contiguous() + if scale.dtype == torch.uint8: + # MXFP8 parameters preserve E8M0 scales as their raw exponent bytes. + # They are biased exponents, not numeric uint8 scale values. + return torch.exp2(scale.to(torch.int16).to(torch.float32) - 127.0) return scale.to(torch.float32) @@ -1226,7 +1320,15 @@ def _get_cached_wo_a_bf16( ) wo_a_scale_param = get_fp8_block_weight_scale(wo_a) - if wo_a_scale_param is not None: + if wo_a_scale_param is None: + # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the + # historical ``_inv`` suffix. + wo_a_scale_param = getattr(wo_a, "weight_scale", None) + # Emulated MXFP8 kernels can replace the original one-byte weight with an + # already-dequantized BF16 tensor while retaining the scale attribute for + # metadata. Applying that retained scale again would double-dequantize the + # weight. Block scaling is only valid while the one-byte FP8 storage remains. + if wo_a_scale_param is not None and wo_a.weight.element_size() == 1: wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to( torch.float32 ) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 38bae03909..e2e38ec5d5 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -22,6 +22,7 @@ from vllm.utils.torch_utils import get_dtype_size from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KpoolTailSpec, @@ -2013,15 +2014,22 @@ def _get_packed_kv_cache_groups( cdiv(len(names), n) * page for page, names in page_size_layers.items() ) - # Bytes a block must hold however the mamba buckets end up split: a mamba - # bucket can go down to one state per group, every other bucket's split is - # already fixed by the repeat pattern. + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + anchor_bytes = max( ( widest_group_bytes( page_size_layers, len(spec.kv_cache_specs) - if isinstance(spec.first_spec, MambaSpec) + if is_state_bucket(spec) else num_groups_for(spec, balanced), ) for spec, page_size_layers, balanced in bucketed @@ -2032,10 +2040,9 @@ def _get_packed_kv_cache_groups( groups = [] for spec, page_size_layers, balanced in bucketed: num_groups = num_groups_for(spec, balanced) - # `_align_hybrid_block_size` pads a mamba state up to one attention - # page, so cap a mamba group at the states a block already fits rather - # than let it widen the block. - if anchor_bytes and isinstance(spec.first_spec, MambaSpec): + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) num_groups = max( num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) @@ -2071,8 +2078,9 @@ def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: if spec_config is None or not spec_config.use_eagle(): return False model_config = vllm_config.model_config - return ( - model_config is not None and model_config.hf_config.model_type == "deepseek_v4" + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", ) @@ -2093,9 +2101,9 @@ def _annotate_eagle_groups( spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses - the target's own decoder layer and so carries no spec marker. Its draft - attention layer is always the last registered layer, so flag whichever + 2. Model-scoped positional fallback for DeepseekV4/V4.1, whose MTP block + reuses the target's own decoder layer and so carries no spec marker. Its + draft attention layer is always the last registered layer, so flag whichever group holds it. This rule is only valid where the groups partition exactly the layers of ``kv_cache_spec``, which is true on the packed grouping path and not in general; other callers must leave @@ -2108,7 +2116,8 @@ def _annotate_eagle_groups( kv_cache_spec: The kv cache spec of each attention layer, in layer registration order. Only read by rule 2. kv_cache_groups: Groups to annotate in place. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group. + use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4/V4.1 packed + group. """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle_block_drop(): diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 9cae2f714c..bc7e6c3a13 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -51,7 +51,7 @@ from vllm.v1.simple_kv_offload.metadata import ( if TYPE_CHECKING: from vllm.v1.core.kv_cache_manager import KVCacheBlocks - from vllm.v1.core.kv_cache_utils import KVCacheBlock + from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -783,6 +783,18 @@ class SimpleCPUOffloadScheduler: req_ids = list(dict.fromkeys(req_ids)) return merged_gpu_block_ids, merged_cpu_block_ids, req_ids, merged_block_meta + def _cached_gpu_block( + self, resolved_hashes: "BlockHashList", block_idx: int, group_id: int + ) -> "KVCacheBlock | None": + """Return the GPU block still cached under this group's block hash.""" + if block_idx >= len(resolved_hashes): + return None + assert self._gpu_block_pool is not None + blocks = self._gpu_block_pool.get_cached_block( + resolved_hashes[block_idx], [group_id] + ) + return blocks[0] if blocks else None + def _select_eager_blocks_to_store( self, state: StoreRequestState, @@ -827,18 +839,25 @@ class SimpleCPUOffloadScheduler: # the middle of the request. group_size = self.group_block_sizes[g] ready = min(len(group_gpu_ids), aligned_tokens // group_size) - resolved_hashes = ( - resolve_block_hashes( - request.block_hashes, self.hash_block_size, group_size - ) - if block_meta is not None - else None + resolved_hashes = resolve_block_hashes( + request.block_hashes, self.hash_block_size, group_size ) curr_mm_idx = 0 secondary_mm_idx = 0 start = state.num_stored_blocks[g] for i, gpu_block_id in enumerate(group_gpu_ids[start:ready], start=start): gpu_block = self._gpu_block_pool.blocks[gpu_block_id] + if gpu_block.is_null: + # Sliding-window groups null pages that left the window + # before the connector sees the block table, but the + # retained prefix-cache tail stays hashed in the GPU free + # queue; store it from there. + recovered = self._cached_gpu_block(resolved_hashes, i, g) + if recovered is None: + advanced_per_group[g] += 1 + continue + gpu_block = recovered + gpu_block_id = gpu_block.block_id if ( self._classify_store_candidate(gpu_block) is not _StoreAdmission.READY @@ -854,7 +873,6 @@ class SimpleCPUOffloadScheduler: if block_meta is not None: token_start = i * group_size token_end = token_start + group_size - assert resolved_hashes is not None parent_hash = ( None if i == 0 diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index f252a35a51..5b13530ef7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -17,9 +17,9 @@ def _resolve_dspark_attention_backend( ) -> AttentionBackendEnum | None: if draft_backend is not None: return draft_backend - # DeepSeek-V4 draft layers share the target's KV-cache layout. Other + # DeepSeek-V4(.1) draft layers share the target's KV-cache layout. Other # DSpark architectures may use a different attention kind. - if draft_model_config.hf_config.model_type == "deepseek_v4": + if draft_model_config.hf_config.model_type in ("deepseek_v4", "deepseek_v41"): if target_backend is not None: logger.info_once( "Using the target model's %s attention backend for the " diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index 0dca4588c3..6b0c5fe51d 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -112,7 +112,14 @@ def get_eagle3_aux_layers_from_config( if not layer_ids: dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None) if dspark_layer_ids: - layer_ids = [i + 1 for i in dspark_layer_ids] + if getattr(hf_config, "model_type", None) == "deepseek_v41": + # v4.1 reads the attention *inputs* of its target layers, and + # the target model captures the entry stream of layer L when + # idx+1 == L, so the ids are used as-is. (v4's ids are in + # capture-after semantics and keep the +1.) + layer_ids = list(dspark_layer_ids) + else: + layer_ids = [i + 1 for i in dspark_layer_ids] if not layer_ids: # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics. target_layer_ids = getattr(hf_config, "target_layer_ids", None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d6b8ed025c..2730416f76 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -813,6 +813,9 @@ class GPUModelRunner( # Persistent buffers for CUDA graphs. self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) + # [max_num_reqs, depth] prompt token ids preceding each request's chunk + # start, for models exposing `token_lookback_depth`. Sized after load. + self.lookback_token_ids: CpuGpuBuffer | None = None self.positions = torch.zeros( self.max_num_tokens, dtype=torch.int64, device=self.device ) @@ -1065,9 +1068,36 @@ class GPUModelRunner( ) return self._mamba_bufs - def _init_model_kwargs(self): + def _prepare_lookback_token_ids(self, num_reqs: int) -> torch.Tensor: + """Gather, per request, the `depth` prompt token ids preceding its + first scheduled token (column j is position start - 1 - j); -1 where + the position is before the prompt or already past it. Generated + positions are left to the model: under async scheduling the CPU token + table holds placeholders for them.""" + buf = self.lookback_token_ids + assert buf is not None + buf.np.fill(-1) + if num_reqs > 0: + depth = buf.np.shape[1] + starts = self.input_batch.num_computed_tokens_cpu[:num_reqs, None] + pos = starts - np.arange(1, depth + 1) + num_prompt = self.input_batch.num_prompt_tokens[:num_reqs, None] + valid = (pos >= 0) & (pos < num_prompt) + rows = np.arange(num_reqs)[:, None] + ids = self.input_batch.token_ids_cpu[rows, np.clip(pos, 0, None)] + buf.np[:num_reqs] = np.where(valid, ids, -1) + return buf.copy_to_gpu() + + def _init_model_kwargs(self, num_reqs: int | None = None): model_kwargs = dict[str, Any]() + if self.lookback_token_ids is not None: + if num_reqs is None: + num_reqs = self.input_batch.num_reqs + model_kwargs["lookback_token_ids"] = self._prepare_lookback_token_ids( + num_reqs + ) + if not self.is_pooling_model: return model_kwargs @@ -5314,6 +5344,11 @@ class GPUModelRunner( self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.model_config ) + lookback_depth = getattr(self.model, "token_lookback_depth", 0) + if lookback_depth > 0: + self.lookback_token_ids = self._make_buffer( + self.max_num_reqs, lookback_depth, dtype=torch.int32 + ) if self.lora_config: self.model = self.load_lora_model( self.model, self.vllm_config, self.device @@ -6076,7 +6111,7 @@ class GPUModelRunner( ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - model_kwargs = self._init_model_kwargs() + model_kwargs = self._init_model_kwargs(num_reqs=0) if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) @@ -6087,7 +6122,7 @@ class GPUModelRunner( elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - model_kwargs = self._init_model_kwargs() + model_kwargs = self._init_model_kwargs(num_reqs=0) else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None