fix(dspark): reduce V2 working VRAM

tests/kernels/attention/test_flashmla_sparse.py

@@ -4,6 +4,47 @@
 import torch
 
 
+def test_compute_global_topk_indices_and_lens_allows_inplace_output():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the Triton sparse index kernel")
+
+    from vllm.models.deepseek_v4.common.ops import compute_global_topk_indices_and_lens
+
+    device = torch.device("cuda")
+    local_indices = torch.tensor(
+        [[0, 1, -1, 5], [2, -1, 3, 4]],
+        dtype=torch.int32,
+        device=device,
+    )
+    token_to_req_indices = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    block_table = torch.tensor([[10, 11], [20, 21]], dtype=torch.int32, device=device)
+    is_valid_token = torch.tensor([True, True], dtype=torch.bool, device=device)
+
+    expected_indices, expected_lens = compute_global_topk_indices_and_lens(
+        local_indices.clone(),
+        token_to_req_indices,
+        block_table,
+        block_size=4,
+        is_valid_token=is_valid_token,
+    )
+
+    lens = torch.empty((local_indices.shape[0],), dtype=torch.int32, device=device)
+    actual_indices, actual_lens = compute_global_topk_indices_and_lens(
+        local_indices,
+        token_to_req_indices,
+        block_table,
+        block_size=4,
+        is_valid_token=is_valid_token,
+        global_topk_indices=local_indices,
+        topk_lens=lens,
+    )
+
+    assert actual_indices.data_ptr() == local_indices.data_ptr()
+    assert actual_lens.data_ptr() == lens.data_ptr()
+    torch.testing.assert_close(actual_indices.cpu(), expected_indices.cpu())
+    torch.testing.assert_close(actual_lens.cpu(), expected_lens.cpu())
+
+
 def test_sparse_flashmla_metadata_smoke():
     import vllm.v1.attention.ops.flashmla as fm
 

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -340,6 +340,53 @@ def test_quant_insert_writes_caller_owned_q_out():
     assert q_out[:, n_heads:padded_heads].abs().max().item() == 0.0
 
 
+def test_quant_insert_allows_inplace_q_when_unpadded():
+    torch.manual_seed(5)
+    device = "cuda"
+    dtype = torch.bfloat16
+    eps = 1e-6
+    num_tokens = 17
+    n_heads = 32
+    block_size = 16
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device)
+    k_cache = torch.zeros(
+        2, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+    slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
+
+    q_out = _call_fused(
+        q,
+        n_heads,
+        kv,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        eps,
+        block_size,
+    )
+
+    q_inplace = q.clone()
+    returned = _call_fused_out(
+        q_inplace,
+        q_inplace,
+        kv,
+        torch.zeros_like(k_cache),
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        eps,
+        block_size,
+    )
+
+    assert returned.data_ptr() == q_inplace.data_ptr()
+    torch.testing.assert_close(q_inplace, q_out, rtol=0, atol=0)
+
+
 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
 
 

tests/kernels/test_mhc_kernels.py

@@ -8,6 +8,8 @@
     _tilelang_hc_prenorm_gemm,
     _torch_hc_prenorm_gemm,
     _use_tf32_hc_prenorm_gemm,
+    mhc_fused_post_pre_tilelang,
+    mhc_fused_post_pre_tilelang_reuse_residual,
 )
 from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
 from vllm.platforms import current_platform
@@ -299,6 +301,79 @@ def run_ref():
     torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.skipif(
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
+)
+@pytest.mark.parametrize("num_tokens", [8, 128])
+def test_mhc_fused_post_pre_reuses_dead_buffers(num_tokens):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    hidden_size = 4096
+    hc_mult = 4
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
+    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+
+    hc_mult2 = hc_mult * hc_mult
+    hc_mult3 = hc_mult * 2 + hc_mult2
+    fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
+    hc_scale = torch.randn((3,), dtype=torch.float32) * 0.1
+    hc_base = torch.randn((hc_mult3,), dtype=torch.float32) * 0.1
+    norm_weight = torch.randn((hidden_size,), dtype=torch.bfloat16)
+
+    rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6
+    hc_post_alpha = 1.0
+    sinkhorn_repeat = 20
+
+    expected = mhc_fused_post_pre_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_alpha,
+        sinkhorn_repeat,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
+    )
+
+    residual_reuse = residual.clone()
+    actual = mhc_fused_post_pre_tilelang_reuse_residual(
+        x,
+        residual_reuse,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_alpha,
+        sinkhorn_repeat,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
+    )
+
+    if num_tokens > 16:
+        assert actual[0].data_ptr() == residual_reuse.data_ptr()
+    else:
+        assert actual[0].data_ptr() != residual_reuse.data_ptr()
+    assert actual[3].data_ptr() == x.data_ptr()
+    for actual_tensor, expected_tensor in zip(actual, expected):
+        torch.testing.assert_close(
+            actual_tensor, expected_tensor, atol=1e-2, rtol=1e-2
+        )
+
+
 @pytest.mark.skipif(
     not current_platform.is_rocm(),
     reason="ROCm required",
@@ -370,3 +445,152 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
 
     out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
     torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_post_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
+    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+    fn = torch.randn(
+        (hc_mult, hc_mult * hidden_size), dtype=torch.float32
+    ) * 1e-4
+    hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
+    hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
+    rms_eps = hc_eps = 1e-6
+
+    pre_hc_head = torch.ops.vllm.mhc_post_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+    expected = torch.ops.vllm.hc_head_fused_kernel_tilelang(
+        pre_hc_head,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_eps,
+    )
+    actual = torch.ops.vllm.mhc_post_hc_head_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_eps,
+    )
+
+    assert actual.shape == (num_tokens, hidden_size)
+    assert actual.dtype == torch.bfloat16
+    assert not torch.isnan(actual).any()
+    torch.testing.assert_close(actual, expected, atol=1.5e-1, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_post_mean_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
+    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+
+    pre_hc_head = torch.ops.vllm.mhc_post_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+    expected = pre_hc_head.mean(dim=1)
+    actual = torch.ops.vllm.mhc_post_mean_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+
+    assert actual.shape == (num_tokens, hidden_size)
+    assert actual.dtype == torch.bfloat16
+    assert not torch.isnan(actual).any()
+    torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not HAS_TILELANG_MHC,
+    reason="TileLang MHC support required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 8, 128])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_post_mean_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
+    torch.set_default_device(DEVICE)
+    set_random_seed(0)
+
+    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
+    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
+    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
+    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
+    fn = torch.randn(
+        (hc_mult, hc_mult * hidden_size), dtype=torch.float32
+    ) * 1e-4
+    hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
+    hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
+    rms_eps = hc_eps = 1e-6
+
+    pre_hc_head = torch.ops.vllm.mhc_post_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+    expected = torch.ops.vllm.hc_head_fused_kernel_tilelang(
+        pre_hc_head,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_eps,
+    )
+    expected_mean = pre_hc_head.mean(dim=1)
+    actual, actual_mean = torch.ops.vllm.mhc_post_mean_hc_head_tilelang(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_eps,
+    )
+
+    assert actual.shape == (num_tokens, hidden_size)
+    assert actual_mean.shape == (num_tokens, hidden_size)
+    assert actual.dtype == torch.bfloat16
+    assert actual_mean.dtype == torch.bfloat16
+    assert not torch.isnan(actual).any()
+    assert not torch.isnan(actual_mean).any()
+    torch.testing.assert_close(actual, expected, atol=1.5e-1, rtol=1e-2)
+    torch.testing.assert_close(actual_mean, expected_mean, atol=5e-2, rtol=1e-2)