diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 010c44797665..e448a335979f 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/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 diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index bf2e9ceaee31..ff3524d4ccf4 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/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 ───────────────────────────── diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 067e7cfd0a43..5cdba79be026 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/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) diff --git a/tests/v1/spec_decode/test_dspark.py b/tests/v1/spec_decode/test_dspark.py index ee324d3f7de5..d483737f4afb 100644 --- a/tests/v1/spec_decode/test_dspark.py +++ b/tests/v1/spec_decode/test_dspark.py @@ -13,8 +13,6 @@ ) from vllm.models.deepseek_v4.nvidia.dspark_triton import ( dspark_context_kv_store, - dspark_inv_rope_bf16_layout, - dspark_markov_greedy_argmax, dspark_qkv_postprocess, dspark_triton_attention, ) @@ -556,58 +554,6 @@ def test_dspark_triton_qkv_postprocess_matches_reference(): assert torch.allclose(kv_out.float(), kv_ref.float(), atol=2e-2, rtol=2e-2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_dspark_triton_inv_rope_bf16_layout_matches_reference(): - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - tokens = 5 - n_groups = 2 - heads_per_group = 3 - head_dim = 512 - rope_dim = 64 - nope_dim = head_dim - rope_dim - o = torch.randn( - tokens, - n_groups * heads_per_group, - head_dim, - dtype=dtype, - device=device, - ).contiguous() - positions = torch.tensor([0, 3, 7, 11, 15], dtype=torch.int64, device=device) - angles = torch.randn(32, rope_dim // 2, dtype=torch.float32, device=device) - cos_sin_cache = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1) - - out = dspark_inv_rope_bf16_layout( - o, - positions, - cos_sin_cache, - n_groups=n_groups, - heads_per_group=heads_per_group, - nope_dim=nope_dim, - rope_dim=rope_dim, - ) - - cs = cos_sin_cache.index_select(0, positions).to(torch.float32) - cos = cs[:, : rope_dim // 2].unsqueeze(1) - sin = cs[:, rope_dim // 2 :].unsqueeze(1) - ref = o.clone() - rope = ref[..., nope_dim:].float() - shape = rope.shape - rope = rope.reshape(*shape[:-1], rope_dim // 2, 2) - even = rope[..., 0] - odd = rope[..., 1] - inv_even = even * cos + odd * sin - inv_odd = odd * cos - even * sin - ref[..., nope_dim:] = ( - torch.stack((inv_even, inv_odd), dim=-1).reshape(shape).to(dtype) - ) - ref = ref.view(tokens, n_groups, heads_per_group, head_dim) - ref = ref.reshape(tokens, n_groups, heads_per_group * head_dim) - assert torch.allclose(out.float(), ref.float(), atol=2e-2, rtol=2e-2) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_dspark_triton_context_kv_store_matches_reference(): device = torch.device("cuda") @@ -648,36 +594,3 @@ def test_dspark_triton_context_kv_store_matches_reference(): assert torch.allclose(cache.float(), cache_ref.float(), atol=2e-2, rtol=2e-2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_dspark_triton_markov_greedy_argmax_matches_reference(): - device = torch.device("cuda") - torch.manual_seed(0) - batch_size = 2 - vocab_size = 137 - rank = 16 - dtype = torch.bfloat16 - - base_logits = torch.randn( - batch_size, vocab_size, dtype=torch.float32, device=device - ) - w1 = torch.randn(vocab_size, rank, dtype=dtype, device=device) - w2 = torch.randn(vocab_size, rank, dtype=dtype, device=device) - prev_token_ids = torch.tensor([3, 41], dtype=torch.int64, device=device) - num_blocks = (vocab_size + 31) // 32 - block_vals = torch.empty(batch_size, num_blocks, dtype=torch.float32, device=device) - block_ids = torch.empty(batch_size, num_blocks, dtype=torch.int64, device=device) - out = torch.empty(batch_size, dtype=torch.int64, device=device) - - dspark_markov_greedy_argmax( - base_logits, - prev_token_ids, - w1, - w2, - block_vals, - block_ids, - out, - block_v=32, - ) - - reference = (base_logits + w1[prev_token_ids].float() @ w2.float().T).argmax(dim=-1) - assert out.tolist() == reference.tolist() diff --git a/tests/v1/spec_decode/test_dspark_config.py b/tests/v1/spec_decode/test_dspark_config.py new file mode 100644 index 000000000000..074a44d31264 --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_config.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.config.speculative import SpeculativeConfig +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + +def test_dspark_standard_rejection_uses_probabilistic_draft_sampling(monkeypatch): + draft_model_config = SimpleNamespace( + model="DeepSeek-V4-Flash-DSpark-hybrid", + architectures=["Qwen3DSparkModel"], + hf_config=SimpleNamespace(model_type="deepseek_v4"), + max_model_len=4096, + verify_with_parallel_config=lambda parallel_config: None, + ) + target_model_config = SimpleNamespace( + model="deepseek-ai/DeepSeek-V4-Flash", + quantization=None, + tokenizer="deepseek-ai/DeepSeek-V4-Flash", + tokenizer_mode="deepseek_v4", + trust_remote_code=True, + allowed_local_media_path="", + allowed_media_domains=None, + dtype="auto", + seed=0, + tokenizer_revision=None, + hf_overrides={}, + max_model_len=4096, + enforce_eager=False, + max_logprobs=20, + config_format="auto", + ) + target_parallel_config = SimpleNamespace( + pipeline_parallel_size=1, + tensor_parallel_size=1, + distributed_executor_backend=None, + max_parallel_loading_workers=None, + disable_custom_all_reduce=False, + ray_workers_use_nsight=False, + placement_group=None, + ) + monkeypatch.setattr( + "vllm.config.speculative.ModelConfig", + lambda **kwargs: draft_model_config, + ) + + speculative_config = SpeculativeConfig( + method="dspark", + num_speculative_tokens=5, + draft_sample_method="greedy", + rejection_sample_method="standard", + target_model_config=target_model_config, + target_parallel_config=target_parallel_config, + ) + + assert speculative_config.draft_sample_method == "probabilistic" + + +def test_dspark_sequential_sampling_writes_persistent_draft_logits(monkeypatch): + num_reqs = 2 + num_speculative_steps = 3 + vocab_size = 4 + max_num_reqs = 5 + state_ids = torch.tensor([3, 1], dtype=torch.int32) + + speculator = object.__new__(DSparkSpeculator) + speculator.num_speculative_steps = num_speculative_steps + speculator.sample_indices = torch.arange(num_reqs * num_speculative_steps) + speculator.sample_idx_mapping = state_ids.repeat_interleave(num_speculative_steps) + speculator.sample_pos = torch.arange(10, 10 + num_reqs * num_speculative_steps) + speculator.input_buffers = SimpleNamespace( + input_ids=torch.arange(max_num_reqs * num_speculative_steps) + ) + speculator._anchor_idx = torch.tensor([0, num_speculative_steps]) + speculator.temperature = torch.ones(max_num_reqs) + speculator.seeds = torch.arange(max_num_reqs, dtype=torch.int64) + speculator.use_fp64_gumbel = False + speculator._step_cols = torch.arange(num_speculative_steps, dtype=torch.int32) + speculator.draft_tokens = torch.empty( + max_num_reqs, + num_speculative_steps, + dtype=torch.int64, + ) + speculator.draft_logits = torch.full( + (max_num_reqs, num_speculative_steps, vocab_size), + float("nan"), + ) + + class FakeModel: + + def compute_logits(self, hidden_states): + return torch.arange( + hidden_states.shape[0] * vocab_size, + dtype=torch.float32, + ).view(hidden_states.shape[0], vocab_size) + + def markov_embed(self, previous_tokens): + return previous_tokens.to(torch.float32).unsqueeze(-1) + + def markov_bias(self, markov_embed): + return torch.ones(markov_embed.shape[0], vocab_size) + + sampled_by_step: list[torch.Tensor] = [] + + def fake_gumbel_sample( + logits, + idx_mapping, + temperature, + seeds, + pos, + apply_temperature, + output_processed_logits=None, + output_processed_logits_col=None, + output_processed_logits_inplace=False, + use_fp64=False, + ): + assert output_processed_logits is speculator.draft_logits + assert output_processed_logits_col is not None + assert output_processed_logits_inplace is False + col = int(output_processed_logits_col.item()) + output_processed_logits[idx_mapping.long(), col] = logits + sampled = idx_mapping.to(torch.int64) * 10 + col + sampled_by_step.append(sampled) + return sampled + + monkeypatch.setattr( + "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample", + fake_gumbel_sample, + ) + speculator.model = FakeModel() + + DSparkSpeculator._sample_sequential( + speculator, + num_reqs, + torch.zeros(num_reqs * num_speculative_steps, 1), + ) + + for col in range(num_speculative_steps): + row_logits = torch.arange( + num_reqs * num_speculative_steps * vocab_size, + dtype=torch.float32, + ).view(num_reqs, num_speculative_steps, vocab_size)[:, col] + torch.testing.assert_close( + speculator.draft_logits[state_ids.long(), col], + row_logits + 1.0, + ) + torch.testing.assert_close( + speculator.draft_tokens[:num_reqs, col], + sampled_by_step[col], + ) + + draft_logits = speculator.draft_logits + DSparkSpeculator.clear_runtime_draft_logits(speculator) + assert speculator.draft_logits is draft_logits diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 982582ffaee8..e6f7f122a990 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -6,6 +6,7 @@ import pytest import torch +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( rejection_sample, ) @@ -129,6 +130,91 @@ def _assert_distribution_match( ) +def test_gumbel_sample_stores_processed_logits_inplace(): + logits = torch.tensor( + [[2.0, 4.0, 6.0], [3.0, 6.0, 9.0], [5.0, 10.0, 15.0]], + dtype=torch.float32, + device="cuda", + ) + original = logits.clone() + idx_mapping = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") + temperature = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32, device="cuda") + seeds = torch.tensor([123, 456, 789], dtype=torch.int64, device="cuda") + pos = torch.tensor([10, 11, 12], dtype=torch.int32, device="cuda") + + gumbel_sample( + logits, + idx_mapping, + temperature, + seeds, + pos, + apply_temperature=True, + output_processed_logits_inplace=True, + ) + + expected = original.clone() + expected[2] /= 2.0 + torch.testing.assert_close(logits, expected) + + +def test_dense_draft_logits_index_mapping_matches_state_indexed(): + torch.manual_seed(0) + device = "cuda" + num_speculative_steps = 3 + num_trials = 128 + target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32) + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32) + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + num_speculative_steps, + temperature=1.0, + num_trials=num_trials, + ) + + state_ids = torch.randperm(num_trials, device=device).to(torch.int32) + dense_rows = torch.arange(num_trials, dtype=torch.int32, device=device) + seed = torch.empty_like(inputs["seed"]) + seed[state_ids.long()] = inputs["seed"] + + state_indexed_draft_logits = torch.empty_like(inputs["draft_logits"]) + state_indexed_draft_logits[state_ids.long()] = inputs["draft_logits"] + state_indexed_inputs = dict(inputs) + state_indexed_inputs["idx_mapping"] = state_ids + state_indexed_inputs["expanded_idx_mapping"] = state_ids.repeat_interleave( + num_speculative_steps + 1 + ) + state_indexed_inputs["seed"] = seed + state_indexed_inputs["draft_logits"] = state_indexed_draft_logits + + dense_index_mapping = torch.empty( + num_trials, + dtype=torch.int32, + device=device, + ) + dense_index_mapping[state_ids.long()] = dense_rows + dense_inputs = dict(state_indexed_inputs) + dense_inputs["draft_logits"] = inputs["draft_logits"] + + state_sampled, state_num_sampled = rejection_sample( + **state_indexed_inputs, + num_speculative_steps=num_speculative_steps, + ) + dense_sampled, dense_num_sampled = rejection_sample( + **dense_inputs, + num_speculative_steps=num_speculative_steps, + draft_logits_index_mapping=dense_index_mapping, + ) + + torch.testing.assert_close(dense_num_sampled, state_num_sampled) + for req_idx in range(num_trials): + n = state_num_sampled[req_idx].item() + torch.testing.assert_close( + dense_sampled[req_idx, :n], + state_sampled[req_idx, :n], + ) + + @pytest.mark.parametrize( "num_speculative_steps,temperature", [ diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 949eb298a170..57dbe50c682a 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -669,6 +669,23 @@ def stateless_init_dp_group( # In this case, ensure the input to the experts is sequence parallel # to avoid the excess work. # + # Note this rationale is about replication across the *tensor* parallel + # group, so it applies at data_parallel_size == 1. Upstream #48849 added + # `and self.data_parallel_size > 1` here to recover ~5.6 GiB of SP buffers + # on a pure-TP Nemotron 550B config, but it measured only memory, not + # throughput. On DeepSeek-V4-Flash at TP=2 + EP (SM120, DP=1) the SP + # buffers cost no KV capacity at all, while turning SP-MoE off costs real + # decode throughput: + # + # C=1 decode 274.6-276.7 -> 288.2-289.8 tok/s (+5%) + # prefill 1K 5,948 -> 6,259 (+5%) + # prefill 4K 7,518 -> 7,536 (flat) + # available KV 11.44 GiB -> 11.44 GiB (identical) + # GPU KV tokens 1,189,790 -> 1,189,790 (identical) + # + # So the DP>1 requirement is not kept here. If the pure-TP memory + # regression needs addressing, the gate should be made config-aware + # rather than disabling SP-MoE for every DP=1 deployment. @property def use_sequence_parallel_moe(self) -> bool: return ( @@ -683,7 +700,6 @@ def use_sequence_parallel_moe(self) -> bool: ) and self.enable_expert_parallel and self.tensor_parallel_size > 1 - and self.data_parallel_size > 1 ) @property diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 5d56d6aabd3f..142b08fdd194 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1314,6 +1314,9 @@ def _verify_args(self) -> Self: "are only valid with rejection_sample_method='synthetic'." ) + if self.method == "dspark" and self.rejection_sample_method == "standard": + self.draft_sample_method = "probabilistic" + if self.draft_model_config: self.draft_model_config.verify_with_parallel_config( self.draft_parallel_config diff --git a/vllm/model_executor/kernels/mhc/__init__.py b/vllm/model_executor/kernels/mhc/__init__.py index 3bde02e28a31..4bae81d7cdfa 100644 --- a/vllm/model_executor/kernels/mhc/__init__.py +++ b/vllm/model_executor/kernels/mhc/__init__.py @@ -16,8 +16,11 @@ "hc_head_fused_aiter", "mhc_pre_tilelang", "mhc_post_tilelang", + "mhc_post_mean_tilelang", "mhc_fused_post_pre_tilelang", "hc_head_fused_tilelang", + "mhc_post_hc_head_tilelang", + "mhc_post_mean_hc_head_tilelang", "mhc_pre_torch", "mhc_post_torch", "mhc_fused_post_pre_torch", diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index a5a7d02d5cc8..23eeda1dc842 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -435,7 +435,42 @@ def mhc_post_tilelang( return out -def mhc_fused_post_pre_tilelang( +def mhc_post_mean_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, +) -> torch.Tensor: + """Apply MHC post and return only the mean across HC channels.""" + num_tokens, hc_mult, hidden_size = residual.shape + assert x.shape == (num_tokens, hidden_size) + assert post_layer_mix.shape in ( + (num_tokens, hc_mult, 1), + (num_tokens, hc_mult), + ) + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult) + + out = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=x.device) + if num_tokens == 0: + return out + + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_mean_tilelang as _mhc_post_mean_kernel, + ) + + _mhc_post_mean_kernel( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + out, + hc_mult, + hidden_size, + ) + return out + + +def _mhc_fused_post_pre_tilelang_impl( x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, @@ -452,6 +487,8 @@ def mhc_fused_post_pre_tilelang( tile_n: int = 1, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, + reuse_residual_buffer: bool = False, + reuse_input_buffer: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Run one MHC post block followed by the next MHC pre block. @@ -547,7 +584,11 @@ def mhc_fused_post_pre_tilelang( dtype=torch.float32, device=residual.device, ) - residual_cur = torch.empty_like(residual_flat) + residual_cur = ( + residual_flat + if reuse_residual_buffer and not use_small_fma + else torch.empty_like(residual_flat) + ) post_mix_cur = torch.empty( num_tokens, hc_mult, @@ -560,11 +601,15 @@ def mhc_fused_post_pre_tilelang( dtype=torch.float32, device=residual.device, ) - layer_input_cur = torch.empty( - num_tokens, - hidden_size, - dtype=torch.bfloat16, - device=residual.device, + layer_input_cur = ( + x_flat + if reuse_input_buffer + else torch.empty( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ) ) if use_small_fma: @@ -664,6 +709,87 @@ def mhc_fused_post_pre_tilelang( ) +def mhc_fused_post_pre_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return _mhc_fused_post_pre_tilelang_impl( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits=n_splits, + tile_n=tile_n, + norm_weight=norm_weight, + norm_eps=norm_eps, + reuse_residual_buffer=False, + reuse_input_buffer=False, + ) + + +def mhc_fused_post_pre_tilelang_reuse_residual( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Reuse dead residual and input storage for the next MHC transition.""" + return _mhc_fused_post_pre_tilelang_impl( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits=n_splits, + tile_n=tile_n, + norm_weight=norm_weight, + norm_eps=norm_eps, + reuse_residual_buffer=True, + reuse_input_buffer=True, + ) + + def _mhc_fused_post_pre_tilelang_fake( x: torch.Tensor, residual: torch.Tensor, @@ -720,6 +846,15 @@ def _mhc_post_tilelang_fake( return torch.empty_like(residual) +def _mhc_post_mean_tilelang_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, +) -> torch.Tensor: + return torch.empty_like(x) + + def hc_head_fused_kernel_tilelang( hs_flat: torch.Tensor, fn: torch.Tensor, @@ -751,6 +886,104 @@ def hc_head_fused_kernel_tilelang( return out +def mhc_post_hc_head_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, +) -> torch.Tensor: + """Apply final MHC post + hc_head without returning the post tensor.""" + num_tokens, hc_mult, hidden_size = residual.shape + assert x.shape == (num_tokens, hidden_size) + assert post_layer_mix.shape in ( + (num_tokens, hc_mult, 1), + (num_tokens, hc_mult), + ) + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult) + assert fn.shape == (hc_mult, hc_mult * hidden_size) + assert hc_scale.shape == (1,) + assert hc_base.shape == (hc_mult,) + + out = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=x.device) + if num_tokens == 0: + return out + + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_hc_head_tilelang as _mhc_post_hc_head_kernel, + ) + + _mhc_post_hc_head_kernel( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_eps, + hc_eps, + hc_mult, + ) + return out + + +def mhc_post_mean_hc_head_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply final MHC post + aux mean + hc_head without returning the post tensor.""" + num_tokens, hc_mult, hidden_size = residual.shape + assert x.shape == (num_tokens, hidden_size) + assert post_layer_mix.shape in ( + (num_tokens, hc_mult, 1), + (num_tokens, hc_mult), + ) + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult) + assert fn.shape == (hc_mult, hc_mult * hidden_size) + assert hc_scale.shape == (1,) + assert hc_base.shape == (hc_mult,) + + out = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=x.device) + mean_out = torch.empty_like(out) + if num_tokens == 0: + return out, mean_out + + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_mean_hc_head_tilelang as _mhc_post_mean_hc_head_kernel, + ) + + _mhc_post_mean_hc_head_kernel( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + fn, + hc_scale, + hc_base, + out, + mean_out, + hidden_size, + rms_eps, + hc_eps, + hc_mult, + ) + return out, mean_out + + def _hc_head_fused_kernel_tilelang_fake( hs_flat: torch.Tensor, fn: torch.Tensor, @@ -765,6 +998,37 @@ def _hc_head_fused_kernel_tilelang_fake( ) +def _mhc_post_hc_head_tilelang_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, +) -> torch.Tensor: + num_tokens, _, hidden_size = residual.shape + return torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=x.device) + + +def _mhc_post_mean_hc_head_tilelang_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens, _, hidden_size = residual.shape + out = torch.empty(num_tokens, hidden_size, dtype=torch.bfloat16, device=x.device) + return out, torch.empty_like(out) + + direct_register_custom_op( op_name="mhc_pre_tilelang", op_func=mhc_pre_tilelang, @@ -778,6 +1042,13 @@ def _hc_head_fused_kernel_tilelang_fake( fake_impl=_mhc_post_tilelang_fake, ) +direct_register_custom_op( + op_name="mhc_post_mean_tilelang", + op_func=mhc_post_mean_tilelang, + mutates_args=[], + fake_impl=_mhc_post_mean_tilelang_fake, +) + direct_register_custom_op( op_name="mhc_fused_post_pre_tilelang", op_func=mhc_fused_post_pre_tilelang, @@ -791,3 +1062,17 @@ def _hc_head_fused_kernel_tilelang_fake( mutates_args=[], fake_impl=_hc_head_fused_kernel_tilelang_fake, ) + +direct_register_custom_op( + op_name="mhc_post_hc_head_tilelang", + op_func=mhc_post_hc_head_tilelang, + mutates_args=[], + fake_impl=_mhc_post_hc_head_tilelang_fake, +) + +direct_register_custom_op( + op_name="mhc_post_mean_hc_head_tilelang", + op_func=mhc_post_mean_hc_head_tilelang, + mutates_args=[], + fake_impl=_mhc_post_mean_hc_head_tilelang_fake, +) diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 50abac3e0101..acfcd39212ef 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -695,6 +695,69 @@ def mhc_post_tilelang( T.pdl_trigger() +@tilelang.jit( + pass_configs=pass_configs, +) +def mhc_post_mean_tilelang( + a, + b, + c, + d, + out, + hc: int, + hidden: int, + n_thr: int = 128, + h_blk: int = 1024, +) -> tilelang.JITKernel: + n = T.dynamic("num_tokens") + h = hidden + h_blk = math.gcd(hidden, h_blk) + + a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type] + b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] + c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type] + d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type] + out: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type] + + with T.Kernel(n, threads=n_thr) as i_n: + b_shared = T.alloc_shared((hc, h_blk), T.bfloat16) + d_shared = T.alloc_shared(h_blk, T.bfloat16) + + b_local = T.alloc_fragment((hc, h_blk), T.float32) + d_local = T.alloc_fragment(h_blk, T.float32) + a_local = T.alloc_fragment((hc, hc), T.float32) + c_local = T.alloc_fragment(hc, T.float32) + + if ENABLE_PDL: + T.pdl_sync() + T.copy(a[i_n, 0, 0], a_local) + T.copy(c[i_n, 0], c_local) + + for i0_h in T.Serial(T.ceildiv(h, h_blk)): + mean_local = T.alloc_fragment(h_blk, T.float32) + T.clear(mean_local) + + T.copy(b[i_n, 0, i0_h * h_blk], b_shared) + T.copy(d[i_n, i0_h * h_blk], d_shared) + T.copy(b_shared, b_local) + T.copy(d_shared, d_local) + + for i_hco in T.serial(hc): + for i1_h in T.Parallel(h_blk): + val = T.alloc_var( + T.float32, + init=c_local[i_hco] * d_local[i1_h], + ) + for i_hci in T.unroll(hc): + val += a_local[i_hci, i_hco] * b_local[i_hci, i1_h] + mean_local[i1_h] += val * (1.0 / hc) + + T.copy(mean_local, out[i_n, i0_h * h_blk]) + + if ENABLE_PDL: + T.pdl_trigger() + + @tilelang.jit( pass_configs=pass_configs, ) @@ -973,3 +1036,268 @@ def hc_head_fuse_tilelang( if ENABLE_PDL: T.pdl_trigger() + + +@tilelang.jit( + pass_configs=pass_configs, +) +def mhc_post_hc_head_tilelang( + comb_mix, + residual, + post_mix, + layer_input, + fn, + hc_scale, + hc_base, + out, + hidden_size: int, + rms_eps: float, + hc_eps: float, + hc_mult: int = 4, + n_thr: int = 128, + h_blk: int = 1024, +): + """Fuse final MHC post reconstruction with hc_head. + + This computes the same result as: + hc_head_fuse_tilelang(mhc_post_tilelang(...)) + + It deliberately recomputes the post-mapped residual in two passes so the + [num_tokens, hc_mult, hidden_size] post tensor does not have to live in + global memory. + """ + num_tokens = T.dynamic("num_tokens") + hc_dim = hc_mult * hidden_size + h_block = math.gcd(h_blk, hidden_size) + n_h = hidden_size // h_block + + comb_mix: T.Tensor[[num_tokens, hc_mult, hc_mult], T.float32] # type: ignore[no-redef,valid-type] + residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef,valid-type] + layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type] + hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type] + hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type] + out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + + with T.Kernel(num_tokens, threads=n_thr) as i: + if ENABLE_PDL: + T.pdl_sync() + + comb_shared = T.alloc_shared((hc_mult, hc_mult), T.float32) + post_shared = T.alloc_shared(hc_mult, T.float32) + T.copy(comb_mix[i, 0, 0], comb_shared, disable_tma=True) + T.copy(post_mix[i, 0], post_shared, disable_tma=True) + + sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all") + mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all") + T.fill(sqrsum_r, 0.0) + T.fill(mixes_r, 0.0) + + for i_h in T.serial(n_h): + residual_shared = T.alloc_shared((hc_mult, h_block), T.bfloat16) + layer_shared = T.alloc_shared(h_block, T.bfloat16) + residual_local = T.alloc_fragment((hc_mult, h_block), T.float32) + layer_local = T.alloc_fragment(h_block, T.float32) + post_local = T.alloc_fragment((hc_mult, h_block), T.float32) + + T.copy(residual[i, 0, i_h * h_block], residual_shared, disable_tma=True) + T.copy(layer_input[i, i_h * h_block], layer_shared, disable_tma=True) + T.copy(residual_shared, residual_local) + T.copy(layer_shared, layer_local) + + for out_hc, h in T.Parallel(hc_mult, h_block): + val = T.alloc_var( + T.float32, init=post_shared[out_hc] * layer_local[h] + ) + for in_hc in T.unroll(hc_mult): + val += comb_shared[in_hc, out_hc] * residual_local[in_hc, h] + post_local[out_hc, h] = val + sqrsum_r[0] += val * val + + for mix_hc in T.unroll(hc_mult): + for out_hc in T.serial(hc_mult): + fn_local = T.alloc_fragment(h_block, T.float32) + T.copy( + fn[mix_hc, out_hc * hidden_size + i_h * h_block], + fn_local, + disable_tma=True, + ) + for h in T.Parallel(h_block): + mixes_r[mix_hc] += post_local[out_hc, h] * fn_local[h] + + T.finalize_reducer(sqrsum_r) + T.finalize_reducer(mixes_r) + + pre_mix_shared = T.alloc_shared(hc_mult, T.float32) + rsqrt_val = T.alloc_fragment(1, T.float32) + rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps) + for mix_hc in T.Parallel(hc_mult): + pre_mix_shared[mix_hc] = ( + T.sigmoid( + mixes_r[mix_hc] * rsqrt_val[0] * hc_scale[0] + hc_base[mix_hc] + ) + + hc_eps + ) + + for i_h in T.Pipelined(n_h, num_stages=2): + residual_shared = T.alloc_shared((hc_mult, h_block), T.bfloat16) + layer_shared = T.alloc_shared(h_block, T.bfloat16) + residual_local = T.alloc_fragment((hc_mult, h_block), T.float32) + layer_local = T.alloc_fragment(h_block, T.float32) + out_local = T.alloc_fragment(h_block, T.float32) + T.clear(out_local) + + T.copy(residual[i, 0, i_h * h_block], residual_shared, disable_tma=True) + T.copy(layer_input[i, i_h * h_block], layer_shared, disable_tma=True) + T.copy(residual_shared, residual_local) + T.copy(layer_shared, layer_local) + + for out_hc in T.serial(hc_mult): + pre = pre_mix_shared[out_hc] + for h in T.Parallel(h_block): + val = T.alloc_var( + T.float32, init=post_shared[out_hc] * layer_local[h] + ) + for in_hc in T.unroll(hc_mult): + val += comb_shared[in_hc, out_hc] * residual_local[in_hc, h] + out_local[h] += pre * val + + T.copy(out_local, out[i, i_h * h_block], disable_tma=True) + + if ENABLE_PDL: + T.pdl_trigger() + + +@tilelang.jit( + pass_configs=pass_configs, +) +def mhc_post_mean_hc_head_tilelang( + comb_mix, + residual, + post_mix, + layer_input, + fn, + hc_scale, + hc_base, + out, + mean_out, + hidden_size: int, + rms_eps: float, + hc_eps: float, + hc_mult: int = 4, + n_thr: int = 128, + h_blk: int = 1024, +): + """Fuse final MHC post, aux mean, and hc_head. + + This is for Eagle/DSpark final aux-hidden capture where callers need + ``mhc_post(...).mean(dim=1)`` and the regular hc_head output, but not the + full post-mapped [num_tokens, hc_mult, hidden_size] tensor. + """ + num_tokens = T.dynamic("num_tokens") + hc_dim = hc_mult * hidden_size + h_block = math.gcd(h_blk, hidden_size) + n_h = hidden_size // h_block + + comb_mix: T.Tensor[[num_tokens, hc_mult, hc_mult], T.float32] # type: ignore[no-redef,valid-type] + residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef,valid-type] + layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type] + hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type] + hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type] + out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + mean_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + + with T.Kernel(num_tokens, threads=n_thr) as i: + if ENABLE_PDL: + T.pdl_sync() + + comb_shared = T.alloc_shared((hc_mult, hc_mult), T.float32) + post_shared = T.alloc_shared(hc_mult, T.float32) + T.copy(comb_mix[i, 0, 0], comb_shared, disable_tma=True) + T.copy(post_mix[i, 0], post_shared, disable_tma=True) + + sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all") + mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all") + T.fill(sqrsum_r, 0.0) + T.fill(mixes_r, 0.0) + + for i_h in T.serial(n_h): + residual_shared = T.alloc_shared((hc_mult, h_block), T.bfloat16) + layer_shared = T.alloc_shared(h_block, T.bfloat16) + residual_local = T.alloc_fragment((hc_mult, h_block), T.float32) + layer_local = T.alloc_fragment(h_block, T.float32) + post_local = T.alloc_fragment((hc_mult, h_block), T.float32) + + T.copy(residual[i, 0, i_h * h_block], residual_shared, disable_tma=True) + T.copy(layer_input[i, i_h * h_block], layer_shared, disable_tma=True) + T.copy(residual_shared, residual_local) + T.copy(layer_shared, layer_local) + + for out_hc, h in T.Parallel(hc_mult, h_block): + val = T.alloc_var( + T.float32, init=post_shared[out_hc] * layer_local[h] + ) + for in_hc in T.unroll(hc_mult): + val += comb_shared[in_hc, out_hc] * residual_local[in_hc, h] + post_local[out_hc, h] = val + sqrsum_r[0] += val * val + + for mix_hc in T.unroll(hc_mult): + for out_hc in T.serial(hc_mult): + fn_local = T.alloc_fragment(h_block, T.float32) + T.copy( + fn[mix_hc, out_hc * hidden_size + i_h * h_block], + fn_local, + disable_tma=True, + ) + for h in T.Parallel(h_block): + mixes_r[mix_hc] += post_local[out_hc, h] * fn_local[h] + + T.finalize_reducer(sqrsum_r) + T.finalize_reducer(mixes_r) + + pre_mix_shared = T.alloc_shared(hc_mult, T.float32) + rsqrt_val = T.alloc_fragment(1, T.float32) + rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps) + for mix_hc in T.Parallel(hc_mult): + pre_mix_shared[mix_hc] = ( + T.sigmoid( + mixes_r[mix_hc] * rsqrt_val[0] * hc_scale[0] + hc_base[mix_hc] + ) + + hc_eps + ) + + for i_h in T.Pipelined(n_h, num_stages=2): + residual_shared = T.alloc_shared((hc_mult, h_block), T.bfloat16) + layer_shared = T.alloc_shared(h_block, T.bfloat16) + residual_local = T.alloc_fragment((hc_mult, h_block), T.float32) + layer_local = T.alloc_fragment(h_block, T.float32) + out_local = T.alloc_fragment(h_block, T.float32) + mean_local = T.alloc_fragment(h_block, T.float32) + T.clear(out_local) + T.clear(mean_local) + + T.copy(residual[i, 0, i_h * h_block], residual_shared, disable_tma=True) + T.copy(layer_input[i, i_h * h_block], layer_shared, disable_tma=True) + T.copy(residual_shared, residual_local) + T.copy(layer_shared, layer_local) + + for out_hc in T.serial(hc_mult): + pre = pre_mix_shared[out_hc] + for h in T.Parallel(h_block): + val = T.alloc_var( + T.float32, init=post_shared[out_hc] * layer_local[h] + ) + for in_hc in T.unroll(hc_mult): + val += comb_shared[in_hc, out_hc] * residual_local[in_hc, h] + out_local[h] += pre * val + mean_local[h] += val * (1.0 / hc_mult) + + T.copy(out_local, out[i, i_h * h_block], disable_tma=True) + T.copy(mean_local, mean_out[i, i_h * h_block], disable_tma=True) + + if ENABLE_PDL: + T.pdl_trigger() diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 98feee842ec2..6efe06c93a75 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -422,9 +422,14 @@ def _get_q_padded_scratch(self, q: torch.Tensor) -> torch.Tensor: ) return scratch[:num_tokens] + def _qnorm_rope_can_write_inplace(self) -> bool: + return self.n_local_heads == self.padded_heads + def reserve_profile_scratch(self) -> None: if self.kv_cache_torch_dtype != torch.uint8: return + if self._qnorm_rope_can_write_inplace(): + return device = self.q_norm.weight.device if device.type not in ("cuda", "xpu"): return @@ -640,6 +645,8 @@ def _fused_qnorm_rope_kv_insert( # Profile run: kernel doesn't fire; produce a padded tensor so # downstream FlashMLA gets the right shape. if self.kv_cache_torch_dtype == torch.uint8: + if self._qnorm_rope_can_write_inplace(): + return q return self._get_q_padded_scratch(q) if self.n_local_heads < self.padded_heads: return F.pad( @@ -669,7 +676,10 @@ def _fused_qnorm_rope_kv_insert( # the padding head slots into a reusable q buffer. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - q_out = self._get_q_padded_scratch(q) + if self._qnorm_rope_can_write_inplace(): + q_out = q + else: + q_out = self._get_q_padded_scratch(q) torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 69bb42c791cd..442f9840857e 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -22,9 +22,9 @@ ) from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( - hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, - mhc_post_tilelang, + mhc_fused_post_pre_tilelang_reuse_residual, + mhc_post_hc_head_tilelang, mhc_pre_tilelang, ) from vllm.model_executor.layers.fused_moe import ( @@ -627,7 +627,7 @@ def forward( norm_eps=attn_norm_eps, ) else: - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang_reuse_residual( x, residual, post_mix, @@ -650,7 +650,7 @@ def forward( ffn_norm_weight = self.ffn_norm.weight.data ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang_reuse_residual( x, residual, post_mix, @@ -922,9 +922,11 @@ def _forward_impl( res_mix, residual, ) - x = mhc_post_tilelang(x, residual, post_mix, res_mix) - x = hc_head_fused_kernel_tilelang( + x = mhc_post_hc_head_tilelang( x, + residual, + post_mix, + res_mix, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, diff --git a/vllm/models/deepseek_v4/nvidia/dspark_triton.py b/vllm/models/deepseek_v4/nvidia/dspark_triton.py index 41a4614f85ac..dd985d07dbea 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark_triton.py +++ b/vllm/models/deepseek_v4/nvidia/dspark_triton.py @@ -130,140 +130,6 @@ def dspark_qkv_postprocess( return q_out, kv_out -@triton.jit -def _dspark_inv_rope_bf16_layout_kernel( - o_ptr, - positions_ptr, - cos_sin_ptr, - out_ptr, - num_tokens, - heads_per_group: tl.constexpr, - o_stride_token, - o_stride_head, - cache_stride_pos, - out_stride_token, - out_stride_group, - out_stride_hidden, - head_dim: tl.constexpr, - nope_dim: tl.constexpr, - rope_dim: tl.constexpr, - block_d: tl.constexpr, -): - token_pid = tl.program_id(0).to(tl.int64) - head_pid = tl.program_id(1).to(tl.int64) - group = head_pid // heads_per_group - head_in_group = head_pid - group * heads_per_group - - offs = tl.arange(0, block_d) - mask = offs < head_dim - input_base = o_ptr + token_pid * o_stride_token + head_pid * o_stride_head - x = tl.load(input_base + offs, mask=mask, other=0.0).to(tl.float32) - - rope_offsets = offs - nope_dim - is_rope = (offs >= nope_dim) & mask - pair_offsets = tl.maximum(rope_offsets, 0) ^ 1 - x_partner = tl.load( - input_base + nope_dim + pair_offsets, - mask=is_rope, - other=0.0, - ).to(tl.float32) - - pos = tl.load(positions_ptr + token_pid).to(tl.int64) - pair_idx = tl.maximum(rope_offsets // 2, 0) - rope_half: tl.constexpr = rope_dim // 2 - cos = tl.load( - cos_sin_ptr + pos * cache_stride_pos + pair_idx, - mask=is_rope & (pair_idx < rope_half), - other=1.0, - ).to(tl.float32) - sin = tl.load( - cos_sin_ptr + pos * cache_stride_pos + rope_half + pair_idx, - mask=is_rope & (pair_idx < rope_half), - other=0.0, - ).to(tl.float32) - x_add = x * cos + x_partner * sin - x_sub = x * cos - x_partner * sin - rotated = tl.where((rope_offsets & 1) == 0, x_add, x_sub) - out = tl.where(is_rope, rotated, x) - - hidden_offsets = head_in_group * head_dim + offs - out_base = out_ptr + token_pid * out_stride_token + group * out_stride_group - tl.store( - out_base + hidden_offsets * out_stride_hidden, - out, - mask=mask, - ) - - -def dspark_inv_rope_bf16_layout( - o: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - *, - n_groups: int, - heads_per_group: int, - nope_dim: int, - rope_dim: int, -) -> torch.Tensor: - """Inverse-RoPE attention output into grouped BF16 WO_A layout. - - Output shape is [tokens, groups, heads_per_group * head_dim], matching the - activation layout consumed by DSpark WO_A. - """ - if o.dim() != 3: - raise ValueError(f"o must be [tokens, heads, dim], got {o.shape}") - if o.dtype is not torch.bfloat16: - raise ValueError("DSpark BF16 inverse-RoPE layout currently requires BF16") - if not o.is_contiguous(): - raise ValueError("o must be contiguous") - num_tokens, num_heads, head_dim = o.shape - if num_tokens == 0: - return torch.empty( - (0, n_groups, heads_per_group * head_dim), - dtype=o.dtype, - device=o.device, - ) - if num_heads != n_groups * heads_per_group: - raise ValueError( - f"head/group mismatch: heads={num_heads}, " - f"n_groups={n_groups}, heads_per_group={heads_per_group}" - ) - if head_dim != nope_dim + rope_dim: - raise ValueError( - f"head_dim={head_dim} does not match nope+rope={nope_dim + rope_dim}" - ) - if cos_sin_cache.shape[-1] != rope_dim: - raise ValueError( - f"cos_sin_cache last dim must be {rope_dim}, got {cos_sin_cache.shape}" - ) - out = torch.empty( - (num_tokens, n_groups, heads_per_group * head_dim), - dtype=o.dtype, - device=o.device, - ) - block_d = triton.next_power_of_2(head_dim) - _dspark_inv_rope_bf16_layout_kernel[(num_tokens, num_heads)]( - o, - positions.contiguous(), - cos_sin_cache, - out, - num_tokens, - heads_per_group=heads_per_group, - o_stride_token=o.stride(0), - o_stride_head=o.stride(1), - cache_stride_pos=cos_sin_cache.stride(0), - out_stride_token=out.stride(0), - out_stride_group=out.stride(1), - out_stride_hidden=out.stride(2), - head_dim=head_dim, - nope_dim=nope_dim, - rope_dim=rope_dim, - block_d=block_d, - num_warps=8, - ) - return out - - @triton.jit def _dspark_context_kv_store_kernel( kv_ptr, @@ -537,154 +403,6 @@ def dspark_triton_attention( return out -@triton.jit -def _dspark_markov_argmax_blocks_kernel( - base_logits_ptr, - prev_token_ids_ptr, - w1_ptr, - w2_ptr, - block_vals_ptr, - block_ids_ptr, - vocab_size: tl.constexpr, - rank: tl.constexpr, - base_batch_stride: tl.constexpr, - w1_vocab_stride: tl.constexpr, - w2_vocab_stride: tl.constexpr, - num_blocks: tl.constexpr, - block_v: tl.constexpr, - block_r: tl.constexpr, -): - batch_pid = tl.program_id(0) - block_pid = tl.program_id(1) - offs_v = block_pid * block_v + tl.arange(0, block_v) - offs_r = tl.arange(0, block_r) - v_mask = offs_v < vocab_size - r_mask = offs_r < rank - - prev_token_id = tl.load(prev_token_ids_ptr + batch_pid).to(tl.int64) - embed = tl.load( - w1_ptr + prev_token_id * w1_vocab_stride + offs_r, - mask=r_mask, - other=0.0, - ).to(tl.float32) - w2 = tl.load( - w2_ptr + offs_v[:, None] * w2_vocab_stride + offs_r[None, :], - mask=v_mask[:, None] & r_mask[None, :], - other=0.0, - ).to(tl.float32) - markov_bias = tl.sum(w2 * embed[None, :], axis=1) - base_logits = tl.load( - base_logits_ptr + batch_pid * base_batch_stride + offs_v, - mask=v_mask, - other=-float("inf"), - ).to(tl.float32) - scores = tl.where(v_mask, base_logits + markov_bias, -float("inf")) - max_val = tl.max(scores, axis=0) - max_id = tl.min(tl.where(scores == max_val, offs_v, vocab_size), axis=0) - tl.store(block_vals_ptr + batch_pid * num_blocks + block_pid, max_val) - tl.store(block_ids_ptr + batch_pid * num_blocks + block_pid, max_id) - - -@triton.jit -def _dspark_markov_argmax_reduce_kernel( - block_vals_ptr, - block_ids_ptr, - out_token_ids_ptr, - num_blocks: tl.constexpr, - block_nb: tl.constexpr, -): - batch_pid = tl.program_id(0) - offs = tl.arange(0, block_nb) - mask = offs < num_blocks - vals = tl.load( - block_vals_ptr + batch_pid * num_blocks + offs, - mask=mask, - other=-float("inf"), - ).to(tl.float32) - ids = tl.load( - block_ids_ptr + batch_pid * num_blocks + offs, - mask=mask, - other=2147483647, - ).to(tl.int64) - max_val = tl.max(vals, axis=0) - token_id = tl.min(tl.where(vals == max_val, ids, 2147483647), axis=0) - tl.store(out_token_ids_ptr + batch_pid, token_id) - - -def dspark_markov_greedy_argmax( - base_logits: torch.Tensor, - prev_token_ids: torch.Tensor, - w1_weight: torch.Tensor, - w2_weight: torch.Tensor, - block_vals: torch.Tensor, - block_ids: torch.Tensor, - out_token_ids: torch.Tensor, - *, - block_v: int = 64, -) -> None: - """Greedy argmax for base logits plus DSpark Markov bias. - - This is a small-shape greedy-only path. It avoids materializing the full - Markov-bias vector separately from the base-logit add and argmax. - """ - if base_logits.dim() != 2: - raise ValueError(f"base_logits must be [batch, vocab], got {base_logits.shape}") - if w1_weight.dim() != 2 or w2_weight.dim() != 2: - raise ValueError("w1_weight and w2_weight must be matrices") - batch_size, vocab_size = base_logits.shape - if w1_weight.shape[0] < vocab_size or w2_weight.shape[0] < vocab_size: - raise ValueError( - "Markov weights must cover the logits vocabulary: " - f"logits={base_logits.shape}, w1={w1_weight.shape}, w2={w2_weight.shape}" - ) - rank = w1_weight.shape[1] - if w2_weight.shape[1] != rank: - raise ValueError( - f"Markov rank mismatch: w1={w1_weight.shape}, w2={w2_weight.shape}" - ) - if prev_token_ids.shape[0] < batch_size or out_token_ids.shape[0] < batch_size: - raise ValueError("prev_token_ids and out_token_ids must cover batch_size") - - num_blocks = triton.cdiv(vocab_size, block_v) - if block_vals.shape[0] < batch_size or block_vals.shape[1] < num_blocks: - raise ValueError( - "block_vals too small: " - f"have {block_vals.shape}, need {(batch_size, num_blocks)}" - ) - if block_ids.shape[0] < batch_size or block_ids.shape[1] < num_blocks: - raise ValueError( - "block_ids too small: " - f"have {block_ids.shape}, need {(batch_size, num_blocks)}" - ) - - block_r = triton.next_power_of_2(rank) - _dspark_markov_argmax_blocks_kernel[(batch_size, num_blocks)]( - base_logits, - prev_token_ids, - w1_weight, - w2_weight, - block_vals, - block_ids, - vocab_size=vocab_size, - rank=rank, - base_batch_stride=base_logits.stride(0), - w1_vocab_stride=w1_weight.stride(0), - w2_vocab_stride=w2_weight.stride(0), - num_blocks=num_blocks, - block_v=block_v, - block_r=block_r, - num_warps=8, - ) - _dspark_markov_argmax_reduce_kernel[(batch_size,)]( - block_vals, - block_ids, - out_token_ids, - num_blocks=num_blocks, - block_nb=triton.next_power_of_2(num_blocks), - num_warps=8, - ) - - @triton.jit def _dspark_markov_probs_blocks_kernel( logits_ptr, diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sm120_decode.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sm120_decode.py index 04d4a3dfcd0b..3c72660c8941 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sm120_decode.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sm120_decode.py @@ -410,21 +410,31 @@ def _forward_prefill( topk_indices = global_indices.view(num_prefill_tokens, 1, -1) else: assert attn_metadata.c128a_prefill_topk_indices is not None - # c128a_prefill_topk_indices are per-request-local compressed - # block positions (0..n-1, -1 padded) -- the same basis the - # decode path maps through the block table. The paged packed - # runner has no block table of its own, so convert to global - # KV-cache slots (+ lens) here, mirroring the C4 branch above. - # C128A is already per-request-local, so unlike C4 it needs no - # cu_base rebase. - global_indices, topk_lens = compute_global_topk_indices_and_lens( - attn_metadata.c128a_prefill_topk_indices, - swa_metadata.token_to_req_indices[num_decode_tokens:num_tokens], - attn_metadata.block_table[:num_reqs], - block_size, - swa_metadata.is_valid_token[num_decode_tokens:num_tokens], - ) - topk_indices = global_indices.view(num_prefill_tokens, 1, -1) + if attn_metadata.c128a_global_prefill_topk_indices is None: + # c128a_prefill_topk_indices are per-request-local compressed + # block positions (0..n-1, -1 padded). The packed runner has + # no block table of its own, so the first C128A SM120 prefill + # layer maps them to global KV-cache slots in-place. Later + # layers reuse the same global view and avoid another 128 MiB + # full-context temporary plus a repeated conversion launch. + assert attn_metadata.c128a_prefill_topk_lens is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + attn_metadata.c128a_prefill_topk_indices, + swa_metadata.token_to_req_indices[num_decode_tokens:num_tokens], + attn_metadata.block_table[:num_reqs], + block_size, + swa_metadata.is_valid_token[num_decode_tokens:num_tokens], + global_topk_indices=attn_metadata.c128a_prefill_topk_indices, + topk_lens=attn_metadata.c128a_prefill_topk_lens, + ) + attn_metadata.c128a_global_prefill_topk_indices = ( + global_indices.view(num_prefill_tokens, 1, -1) + ) + attn_metadata.c128a_prefill_topk_lens = topk_lens + topk_indices = attn_metadata.c128a_global_prefill_topk_indices + topk_lens = attn_metadata.c128a_prefill_topk_lens + assert topk_indices is not None + assert topk_lens is not None topk_indices = topk_indices.contiguous() # --- Launch the packed prefill kernel via the runner. num_tokens > 64 diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 9a7c4e1c83ed..7a51b6778cd1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -23,6 +23,10 @@ from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, + mhc_fused_post_pre_tilelang_reuse_residual, + mhc_post_hc_head_tilelang, + mhc_post_mean_tilelang, + mhc_post_mean_hc_head_tilelang, mhc_post_tilelang, mhc_pre_broadcast_tilelang, mhc_pre_tilelang, @@ -948,6 +952,15 @@ class is used; on SM12x it runs through the stock Triton sparse-MLA route, so return DeepseekV4FlashMLAAttention +def _needs_mtp_target_hidden_buffer(vllm_config: VllmConfig) -> bool: + spec_config = vllm_config.speculative_config + if spec_config is None or spec_config.method != "mtp": + return False + draft_config = spec_config.draft_model_config + hf_config = getattr(draft_config, "hf_config", None) + return hasattr(hf_config, "compress_ratios") and hasattr(hf_config, "hc_mult") + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -1067,7 +1080,7 @@ def forward( norm_eps=attn_norm_eps, ) else: - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang_reuse_residual( x, residual, post_mix, @@ -1091,7 +1104,7 @@ def forward( ffn_norm_weight = self.ffn_norm.weight.data ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang_reuse_residual( x, residual, post_mix, @@ -1198,11 +1211,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) # Pre-hc_head residual stream buffer for the MTP draft. Stable # address (outside the cudagraph pool) so the copy_ in forward() - # refreshes it correctly across captured shapes. - # refreshes it correctly across captured shapes. Only allocated on - # the last PP rank — that's where MTP target hidden states are - # produced. - if get_pp_group().is_last_rank: + # refreshes it correctly across captured shapes. Only allocate it + # when an MTP drafter can consume it; DSpark/DFlash use aux hidden + # states instead. + if get_pp_group().is_last_rank and _needs_mtp_target_hidden_buffer( + vllm_config + ): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, @@ -1257,10 +1271,14 @@ def forward( residual, post_mix, res_mix = None, None, None aux_hidden_states: list[torch.Tensor] = [] final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call + final_aux_mean_pending = False + is_last_rank = get_pp_group().is_last_rank + ran_decoder_layer = False for idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): + ran_decoder_layer = True hidden_states, residual, post_mix, res_mix = layer( hidden_states, positions, @@ -1270,36 +1288,92 @@ def forward( residual, ) if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) - aux_hidden_states.append(aux_recon.mean(dim=1)) - final_aux_recon = aux_recon - if layer is not None: + if ( + idx + 1 == self.end_layer + and is_last_rank + and self._mtp_hidden_buffer is None + ): + final_aux_mean_pending = True + final_aux_recon = None + elif idx + 1 == self.end_layer: + # The final local layer may need the full post tensor for + # PP handoff or the MTP target hidden buffer. + aux_recon = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) + aux_hidden_states.append(aux_recon.mean(dim=1)) + final_aux_recon = aux_recon + else: + aux_hidden_states.append( + mhc_post_mean_tilelang( + hidden_states, + residual, + post_mix, + res_mix, + ) + ) + final_aux_recon = None + + final_post_materialized = not ran_decoder_layer + if ran_decoder_layer: # Reuse if the last layer was captured as an aux hidden state - if self.end_layer in self.aux_hidden_state_layers: + if final_aux_mean_pending: + pass + elif self.end_layer in self.aux_hidden_state_layers: hidden_states = final_aux_recon - else: + final_post_materialized = True + elif ( + not is_last_rank + or self._mtp_hidden_buffer is not None + ): hidden_states = mhc_post_tilelang( hidden_states, residual, post_mix, res_mix ) + final_post_materialized = True - if not get_pp_group().is_last_rank: + if not is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) - # Stash pre-hc_head residual for the MTP draft (captured copy_). - num_tokens = hidden_states.shape[0] - self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + if self._mtp_hidden_buffer is not None: + assert final_post_materialized + # Stash pre-hc_head residual for the MTP draft (captured copy_). + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) - hidden_states = hc_head_fused_kernel_tilelang( - hidden_states, - self.hc_head_fn, - self.hc_head_scale, - self.hc_head_base, - self.rms_norm_eps, - self.hc_eps, - ) + if final_aux_mean_pending: + hidden_states, aux_mean = mhc_post_mean_hc_head_tilelang( + hidden_states, + residual, + post_mix, + res_mix, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + aux_hidden_states.append(aux_mean) + elif ran_decoder_layer and not final_post_materialized: + hidden_states = mhc_post_hc_head_tilelang( + hidden_states, + residual, + post_mix, + res_mix, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + else: + hidden_states = hc_head_fused_kernel_tilelang( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) hidden_states = self.norm(hidden_states) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index c287fba7eb29..dad2e4358003 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -129,6 +129,10 @@ class DeepseekV4FlashMLAMetadata(AttentionMetadata): c128a_decode_topk_lens: torch.Tensor | None = None # Prefill: local topk indices (used by combine_topk_swa_indices). c128a_prefill_topk_indices: torch.Tensor | None = None + # SM120 packed prefill lazily maps local C128A indices to global slot ids + # once per step and reuses them across layers. + c128a_global_prefill_topk_indices: torch.Tensor | None = None + c128a_prefill_topk_lens: torch.Tensor | None = None class DeepseekV4FlashMLAMetadataBuilder( @@ -180,7 +184,11 @@ def __init__( # into adjacent rows (present in both decode and prefill branches of # _build_c128a_topk_metadata_kernel). self.c128a_max_compressed = c128a_max_compressed - self.c128a_global_decode_buffer = torch.empty( + # Decode rows and prefill rows partition the same per-step token + # batch, so a single backing matrix is enough. Decode writes the + # leading num_decode_tokens rows and prefill writes the following + # num_prefill_tokens rows. + self.c128a_topk_buffer = torch.empty( (max_num_batched_tokens, c128a_max_compressed), dtype=torch.int32, device=device, @@ -188,11 +196,6 @@ def __init__( self.c128a_decode_lens_buffer = torch.empty( max_num_batched_tokens, dtype=torch.int32, device=device ) - self.c128a_prefill_buffer = torch.empty( - (max_num_batched_tokens, c128a_max_compressed), - dtype=torch.int32, - device=device, - ) def build( self, @@ -235,6 +238,7 @@ def build( ), c128a_decode_topk_lens=c128a_fields.get("c128a_decode_topk_lens"), c128a_prefill_topk_indices=c128a_fields.get("c128a_prefill_topk_indices"), + c128a_prefill_topk_lens=c128a_fields.get("c128a_prefill_topk_lens"), ) def _build_c128a_metadata( @@ -270,9 +274,8 @@ def _build_c128a_metadata( cm.block_table_tensor[:num_decodes], block_size, cm.slot_mapping, - self.c128a_global_decode_buffer, + self.c128a_topk_buffer, self.c128a_decode_lens_buffer, - self.c128a_prefill_buffer, max_compressed_tokens=self.c128a_max_compressed, max_seq_len=cm.max_seq_len, ) @@ -285,6 +288,9 @@ def _build_c128a_metadata( result["c128a_decode_topk_lens"] = decode_lens if num_prefill_tokens > 0: result["c128a_prefill_topk_indices"] = prefill_local + result["c128a_prefill_topk_lens"] = self.c128a_decode_lens_buffer[ + num_decode_tokens:num_total + ] return result @@ -296,9 +302,8 @@ def build_c128a_topk_metadata( block_table: torch.Tensor, block_size: int, slot_mapping: torch.Tensor, - global_decode_buffer: torch.Tensor, + topk_buffer: torch.Tensor, decode_lens_buffer: torch.Tensor, - prefill_buffer: torch.Tensor, max_compressed_tokens: int = 8192, max_seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -312,9 +317,10 @@ def build_c128a_topk_metadata( """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens + prefill_buffer = topk_buffer[num_decode_tokens:num_tokens] if num_tokens == 0: - global_decode = global_decode_buffer[:num_decode_tokens, :0] + global_decode = topk_buffer[:num_decode_tokens, :0] decode_lens = decode_lens_buffer[:num_decode_tokens] prefill_local = prefill_buffer[:num_prefill_tokens, :0] return global_decode, decode_lens, prefill_local @@ -339,7 +345,7 @@ def build_c128a_topk_metadata( alignment=_C128A_TOPK_ALIGNMENT, ) - global_decode = global_decode_buffer[:num_decode_tokens, :effective_topk] + global_decode = topk_buffer[:num_decode_tokens, :effective_topk] decode_lens = decode_lens_buffer[:num_decode_tokens] prefill_local = prefill_buffer[:num_prefill_tokens, :effective_topk] @@ -353,8 +359,8 @@ def build_c128a_topk_metadata( return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( - global_decode_buffer, - global_decode_buffer.stride(0), + topk_buffer, + topk_buffer.stride(0), decode_lens_buffer, prefill_buffer, prefill_buffer.stride(0), diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9bdcecdead1f..5dd78ba9ba7c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -630,9 +630,14 @@ def _dummy_run( # target returns a persistent buffer sized at max_num_batched_tokens; # slice to the active token count that propose() expects. spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): + if ( + self.speculative_config is not None + and self.speculative_config.method == "mtp" + and hasattr(self.model, "get_mtp_target_hidden_states") + ): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + if pre_hc_hidden_states is not None: + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] self.speculator.propose( input_batch=input_batch, attn_metadata=attn_metadata, @@ -654,6 +659,13 @@ def _dummy_run( mm_inputs=mm_inputs, is_profile=is_profile, ) + clear_runtime_draft_logits = getattr( + self.speculator, + "clear_runtime_draft_logits", + None, + ) + if clear_runtime_draft_logits is not None: + clear_runtime_draft_logits() assert hidden_states is not None # Last PP rank always has hidden_states sample_hidden_states = hidden_states[input_batch.logits_indices] @@ -775,6 +787,13 @@ def capture_model(self) -> int: ) if self.speculator is not None: self.speculator.capture() + clear_runtime_draft_logits = getattr( + self.speculator, + "clear_runtime_draft_logits", + None, + ) + if clear_runtime_draft_logits is not None: + clear_runtime_draft_logits() end_time = time.perf_counter() end_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -1144,12 +1163,22 @@ def sample( # Rejection sampling for spec decoding. assert self.rejection_sampler is not None assert self.speculator is not None - sampler_output = self.rejection_sampler( - logits, - input_batch, - # Draft logits are needed for probabilistic rejection sampling. - self.speculator.draft_logits, - ) + try: + sampler_output = self.rejection_sampler( + logits, + input_batch, + # Draft logits are needed for probabilistic rejection sampling. + self.speculator.draft_logits, + getattr(self.speculator, "draft_logits_index_mapping", None), + ) + finally: + clear_runtime_draft_logits = getattr( + self.speculator, + "clear_runtime_draft_logits", + None, + ) + if clear_runtime_draft_logits is not None: + clear_runtime_draft_logits() return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected @@ -1539,9 +1568,14 @@ def sample_tokens( # target returns a persistent buffer sized at max_num_batched_tokens; # slice to the active token count that propose() expects. spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): + if ( + self.speculative_config is not None + and self.speculative_config.method == "mtp" + and hasattr(self.model, "get_mtp_target_hidden_states") + ): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + if pre_hc_hidden_states is not None: + spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] draft_tokens = self.speculator.propose( input_batch, attn_metadata, diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 190307d5e75a..98678fcf478c 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -94,6 +94,8 @@ def gumbel_block_argmax( processed_logits_ptr, processed_logits_stride, processed_logits_col_ptr, + inplace_processed_logits_ptr, + inplace_processed_logits_stride, vocab_size, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, @@ -128,6 +130,15 @@ def gumbel_block_argmax( mask=mask & is_valid_req, ) + if inplace_processed_logits_ptr is not None: + tl.store( + inplace_processed_logits_ptr + + token_idx * inplace_processed_logits_stride + + block, + logits, + mask=mask & is_valid_req, + ) + # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max. if USE_FP64: @@ -167,6 +178,8 @@ def _gumbel_sample_kernel( processed_logits_ptr, processed_logits_stride, processed_logits_col_ptr, + inplace_processed_logits_ptr, + inplace_processed_logits_stride, logits_ptr, logits_stride, expanded_idx_mapping_ptr, @@ -202,6 +215,8 @@ def _gumbel_sample_kernel( processed_logits_ptr, processed_logits_stride, processed_logits_col_ptr, + inplace_processed_logits_ptr, + inplace_processed_logits_stride, vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, USE_FP64=USE_FP64, @@ -221,6 +236,7 @@ def gumbel_sample( apply_temperature: bool, output_processed_logits: torch.Tensor | None = None, output_processed_logits_col: torch.Tensor | None = None, + output_processed_logits_inplace: bool = False, use_fp64: bool = False, ) -> torch.Tensor: # Enforce contiguity on non-strided input tensors @@ -246,6 +262,8 @@ def gumbel_sample( output_processed_logits, output_processed_logits.stride(0) if output_processed_logits is not None else 0, output_processed_logits_col, + logits if output_processed_logits_inplace else None, + logits.stride(0) if output_processed_logits_inplace else 0, logits, logits.stride(0), expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index cc274cdaa4c9..05e44306c8d5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -36,13 +36,16 @@ def __init__( vllm_config: VllmConfig, device: torch.device, hidden_states_size: int | None = None, + allocate_hidden_states: bool = True, ): super().__init__(vllm_config, device) hidden_states_size = hidden_states_size or self.hidden_size - self.hidden_states = torch.zeros( - self.max_num_tokens, hidden_states_size, dtype=self.dtype, device=device - ) + self.hidden_states: torch.Tensor | None = None + if allocate_hidden_states: + self.hidden_states = torch.zeros( + self.max_num_tokens, hidden_states_size, dtype=self.dtype, device=device + ) # Multimodal inputs not currently supported. self.supports_mm_inputs = False @@ -348,7 +351,10 @@ def propose( ) else: hidden_states = last_hidden_states - self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens]) + context_hidden_states = hidden_states[:num_target_tokens] + if self.hidden_states is not None: + self.hidden_states[:num_target_tokens].copy_(context_hidden_states) + context_hidden_states = self.hidden_states[:num_target_tokens] self._copy_request_inputs( num_reqs, @@ -369,7 +375,7 @@ def propose( "num_rejected_tokens": num_rejected, } self.model.precompute_and_store_context_kv( - self.hidden_states[:num_target_tokens], + context_hidden_states, self.context_positions[:num_target_tokens], **precompute_kwargs, ) @@ -436,7 +442,7 @@ def propose( "num_rejected_tokens": num_rejected, } self.model.precompute_and_store_context_kv( - self.hidden_states[:num_target_tokens], + context_hidden_states, self.context_positions[:num_target_tokens], context_slots, **precompute_kwargs, diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 54110b848bed..3115dd212471 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -27,10 +27,13 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model +logger = init_logger(__name__) + class DSparkSpeculator(DFlashSpeculator): _speculator_name = "DSpark" @@ -43,7 +46,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): draft_hidden = ( vllm_config.speculative_config.draft_model_config.get_hidden_size() ) - super().__init__(vllm_config, device, hidden_states_size=draft_hidden) + super().__init__( + vllm_config, + device, + hidden_states_size=draft_hidden, + allocate_hidden_states=False, + ) # Whether to sample from the anchor position. When True, uses anchor-as-first # (N slots, each position predicts the next token). When False, uses 1+N @@ -56,15 +64,28 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): else: self.num_query_per_req = 1 + self.num_speculative_steps - self._step_cols = torch.arange( - self.num_speculative_steps, dtype=torch.int32, device=device - ) - self._anchor_idx = ( torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) * self.num_query_per_req ) + self._step_cols = torch.arange( + self.num_speculative_steps, dtype=torch.int32, device=device + ) + if self.speculative_config.draft_sample_method == "probabilistic": + # DSpark draft generation is CUDA-graph replayed, so Python-side + # reassignment of self.draft_logits to an intermediate base_logits + # tensor would not run per replay. Keep a persistent graph-written + # buffer keyed by request-state index for probabilistic rejection. + self.draft_logits = torch.zeros( + self.max_num_reqs, + self.num_speculative_steps, + self.vocab_size, + dtype=torch.float32, + device=device, + ) + logger.info("Using DSpark preallocated draft logits for rejection.") + def load_draft_model( self, target_model: torch.nn.Module, @@ -72,6 +93,11 @@ def load_draft_model( ) -> torch.nn.Module: return load_dspark_model(target_model, self.vllm_config) + def clear_runtime_draft_logits(self) -> None: + # The persistent draft_logits buffer is graph-written every draft step. + # Do not clear it; rejection only reads rows selected by idx_map. + pass + def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: # Sequential Markov sampling over the backbone's output hidden states. n_spec = self.num_speculative_steps @@ -93,7 +119,11 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: # Sequential stage: Markov bias from the previously sampled token. markov_embed = self.model.markov_embed(prev) bias = self.model.markov_bias(markov_embed) - logits_i = base_logits[:, i] + bias + logits_i = base_logits[:, i] + if bias.dtype == logits_i.dtype: + logits_i.add_(bias) + else: + logits_i.copy_(logits_i + bias) if self.draft_logits is not None: # sample_pos is the predicted token's position Q; the target # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1. diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 63d479efa0b2..3fdefa964e3e 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -137,6 +137,7 @@ def _verify( idx_mapping_np: np.ndarray, expanded_idx_mapping: torch.Tensor, expanded_local_pos: torch.Tensor, + draft_logits_index_mapping: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: processed_logits = self.sampler.apply_sampling_params( logits, @@ -158,6 +159,7 @@ def _verify( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, self.num_speculative_steps, + draft_logits_index_mapping, self.synthetic_conditional_rates, use_fp64=self.sampler.use_fp64_gumbel, use_block_verification=self.use_block_verification, @@ -173,6 +175,7 @@ def _verify_in_chunks( pos: torch.Tensor, max_chunk_logits: int, max_num_logprobs: int, + draft_logits_index_mapping: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]: cu_num_logits_np = input_batch.cu_num_logits_np use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES @@ -196,6 +199,8 @@ def _verify_in_chunks( input_batch.idx_mapping_np[start:end], input_batch.expanded_idx_mapping[lo:hi], input_batch.expanded_local_pos[lo:hi], + # Indexed by persistent request-state index, so it stays global. + draft_logits_index_mapping, ) chunk_logprobs = self._get_logprobs_tensors( sampled, @@ -234,6 +239,7 @@ def __call__( logits: torch.Tensor, input_batch: InputBatch, draft_logits: torch.Tensor | None = None, + draft_logits_index_mapping: torch.Tensor | None = None, ) -> SamplerOutput: # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear # that num_nans is computed before applying penalties and temperature. @@ -254,6 +260,7 @@ def __call__( pos, max_chunk_logits, max_num_logprobs, + draft_logits_index_mapping, ) num_sampled, num_rejected = get_num_sampled_and_rejected( diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 0c59a1c65eb8..9ed0075375e3 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -133,6 +133,8 @@ def _compute_global_logprobs_and_logsumexp( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + # [max_num_reqs] + draft_logits_index_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -141,6 +143,7 @@ def _compute_global_logprobs_and_logsumexp( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, + DRAFT_LOGITS_USE_INDEX_MAPPING: tl.constexpr, ): target_logit = tl.load( target_logits_ptr + logit_idx * target_logits_stride + token, @@ -158,9 +161,14 @@ def _compute_global_logprobs_and_logsumexp( ) target_log_prob = target_logit - target_lse if HAS_DRAFT_LOGITS: + draft_logits_idx = req_state_idx + if DRAFT_LOGITS_USE_INDEX_MAPPING: + draft_logits_idx = tl.load( + draft_logits_index_mapping_ptr + req_state_idx + ).to(tl.int64) draft_logit = tl.load( draft_logits_ptr - + req_state_idx * draft_logits_stride_0 + + draft_logits_idx * draft_logits_stride_0 + draft_step * draft_logits_stride_1 + token, mask=mask, @@ -207,6 +215,8 @@ def _compute_local_logits_stats_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + # [max_num_reqs] + draft_logits_index_mapping_ptr, # [num_logits] expanded_idx_mapping_ptr, # [num_logits] @@ -217,6 +227,7 @@ def _compute_local_logits_stats_kernel( num_speculative_steps, BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, + DRAFT_LOGITS_USE_INDEX_MAPPING: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) @@ -270,10 +281,15 @@ def _compute_local_logits_stats_kernel( target_sumexp, ) if HAS_DRAFT_LOGITS: + draft_logits_idx = req_state_idx + if DRAFT_LOGITS_USE_INDEX_MAPPING: + draft_logits_idx = tl.load( + draft_logits_index_mapping_ptr + req_state_idx + ).to(tl.int64) # Get local draft max and summed exponentials. draft_logits = tl.load( draft_logits_ptr - + req_state_idx * draft_logits_stride_0 + + draft_logits_idx * draft_logits_stride_0 + draft_step_idx * draft_logits_stride_1 + block_offsets, mask=mask, @@ -311,6 +327,8 @@ def _compute_cumulative_log_p_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + # [max_num_reqs] + draft_logits_index_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -326,6 +344,7 @@ def _compute_cumulative_log_p_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, + DRAFT_LOGITS_USE_INDEX_MAPPING: tl.constexpr, ): req_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64) @@ -355,6 +374,7 @@ def _compute_cumulative_log_p_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_index_mapping_ptr, draft_local_max_ptr, draft_local_max_stride, draft_local_sumexp_ptr, @@ -362,6 +382,7 @@ def _compute_cumulative_log_p_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS, HAS_DRAFT_LOGITS, + DRAFT_LOGITS_USE_INDEX_MAPPING, ) log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0) tl.store(cumulative_log_p_ptr + logit_idx, log_p) @@ -387,6 +408,8 @@ def _compute_local_residual_mass_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + # [max_num_reqs] + draft_logits_index_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -404,6 +427,7 @@ def _compute_local_residual_mass_kernel( vocab_num_blocks, BLOCK_SIZE: tl.constexpr, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, + DRAFT_LOGITS_USE_INDEX_MAPPING: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) @@ -436,6 +460,7 @@ def _compute_local_residual_mass_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_index_mapping_ptr, draft_local_max_ptr, draft_local_max_stride, draft_local_sumexp_ptr, @@ -443,6 +468,7 @@ def _compute_local_residual_mass_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS, True, # HAS_DRAFT_LOGITS + DRAFT_LOGITS_USE_INDEX_MAPPING, ) # Compute the residual mass: max(p_i * M_b(x|x_{