From 64ed3f60806d91b0d8c78a885a04a3f801d7ef06 Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Fri, 3 Jul 2026 02:09:33 +0000 Subject: [PATCH 1/8] fix(dspark): reduce V2 working VRAM --- .../kernels/attention/test_flashmla_sparse.py | 41 ++++++++++++++++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 47 +++++++++++++++++++ vllm/models/deepseek_v4/attention.py | 12 ++++- .../nvidia/flashinfer_sm120_decode.py | 40 ++++++++++------ vllm/models/deepseek_v4/nvidia/model.py | 27 +++++++---- vllm/models/deepseek_v4/sparse_mla.py | 34 ++++++++------ vllm/v1/worker/gpu/model_runner.py | 18 +++++-- .../gpu/spec_decode/dflash/speculator.py | 18 ++++--- .../gpu/spec_decode/dspark/speculator.py | 7 ++- 9 files changed, 195 insertions(+), 49 deletions(-) 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/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/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..4ed23d037dcb 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -948,6 +948,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, @@ -1198,11 +1207,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, @@ -1288,9 +1298,10 @@ def forward( if not get_pp_group().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: + # 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, 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..1fcd15c5e148 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, @@ -1539,9 +1544,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/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..20094bdd0c5a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -43,7 +43,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 From 65d769b84c43f51cdb201574d7186fa475414921 Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Fri, 3 Jul 2026 08:16:27 +0000 Subject: [PATCH 2/8] fix(dspark): reduce MHC working buffers --- tests/kernels/test_mhc_kernels.py | 220 ++++++++++++ vllm/model_executor/kernels/mhc/__init__.py | 3 + vllm/model_executor/kernels/mhc/tilelang.py | 281 ++++++++++++++- .../kernels/mhc/tilelang_kernels.py | 328 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/dspark.py | 14 +- vllm/models/deepseek_v4/nvidia/model.py | 103 ++++-- 6 files changed, 921 insertions(+), 28 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 067e7cfd0a43..140628fbf4a3 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,75 @@ 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", +) +def test_mhc_fused_post_pre_reuses_residual_buffer(): + torch.set_default_device(DEVICE) + set_random_seed(0) + + num_tokens = 128 + 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, + ) + + assert actual[0].data_ptr() == residual_reuse.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 +441,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/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..7a2e1b6ea4c8 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,7 @@ 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, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Run one MHC post block followed by the next MHC pre block. @@ -547,7 +583,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, @@ -664,6 +704,84 @@ 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, + ) + + +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]: + 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, + ) + + def _mhc_fused_post_pre_tilelang_fake( x: torch.Tensor, residual: torch.Tensor, @@ -720,6 +838,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 +878,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 +990,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 +1034,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 +1054,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/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/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 4ed23d037dcb..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, @@ -1076,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, @@ -1100,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, @@ -1267,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, @@ -1280,37 +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}) 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 From 56523cca9ccafb671805cc4468fa465fc3ed555d Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Fri, 3 Jul 2026 21:06:55 +0000 Subject: [PATCH 3/8] fix(dspark): reuse current draft logits for rejection --- tests/v1/spec_decode/test_dspark_config.py | 56 ++++++++++++ .../test_rejection_sampler_utils.py | 86 +++++++++++++++++++ vllm/config/speculative.py | 3 + vllm/v1/worker/gpu/model_runner.py | 36 ++++++-- vllm/v1/worker/gpu/sample/gumbel.py | 18 ++++ .../gpu/spec_decode/dspark/speculator.py | 42 +++++++-- .../gpu/spec_decode/rejection_sampler.py | 7 ++ .../spec_decode/rejection_sampler_utils.py | 62 ++++++++++++- vllm/v1/worker/gpu/spec_decode/speculator.py | 5 +- 9 files changed, 297 insertions(+), 18 deletions(-) create mode 100644 tests/v1/spec_decode/test_dspark_config.py 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..6c230550101e --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_config.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +from vllm.config.speculative import SpeculativeConfig + + +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, + 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" 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/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/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1fcd15c5e148..5dd78ba9ba7c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -659,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] @@ -780,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] @@ -1149,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 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/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 20094bdd0c5a..3ae5125d4858 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" @@ -61,15 +64,24 @@ 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.draft_logits_index_mapping: torch.Tensor | None = None + if self.speculative_config.draft_sample_method == "probabilistic": + # DSpark keeps only the current step's dense [req, step, vocab] + # logits. Rejection kernels still index by req_state_idx, so this + # maps req_state_idx -> dense row for the active request batch. + self.draft_logits_index_mapping = torch.empty( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self._dense_req_indices = torch.arange( + self.max_num_reqs, dtype=torch.int32, device=device + ) + logger.info("Using DSpark current-step draft logits for rejection.") + def load_draft_model( self, target_model: torch.nn.Module, @@ -77,6 +89,10 @@ def load_draft_model( ) -> torch.nn.Module: return load_dspark_model(target_model, self.vllm_config) + def clear_runtime_draft_logits(self) -> None: + if self.draft_logits_index_mapping is not None: + self.draft_logits = None + 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 @@ -89,6 +105,13 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) + if self.draft_logits_index_mapping is not None: + self.draft_logits = base_logits + # idx_map is repeated for all speculative steps for a request, so + # the first column identifies the request state for each dense row. + self.draft_logits_index_mapping[idx_map[:, 0].long()] = ( + self._dense_req_indices[:num_reqs] + ) # Anchor (bonus) token per request = the input id at query offset 0, # read via the precomputed persistent index (fixed buffer for capture). @@ -98,8 +121,12 @@ 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 - if self.draft_logits is not None: + 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_index_mapping 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. draft_i = gumbel_sample( @@ -109,8 +136,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: self.seeds, sample_pos[:, i] - 1, apply_temperature=True, - output_processed_logits=self.draft_logits, - output_processed_logits_col=self._step_cols[i], + output_processed_logits_inplace=True, use_fp64=self.use_fp64_gumbel, ) else: 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_{ Date: Thu, 9 Jul 2026 23:40:08 +0000 Subject: [PATCH 4/8] fix(dspark): keep rejection logits graph-stable --- tests/v1/spec_decode/test_dspark_config.py | 101 ++++++++++++++++++ .../gpu/spec_decode/dspark/speculator.py | 41 ++++--- 2 files changed, 121 insertions(+), 21 deletions(-) diff --git a/tests/v1/spec_decode/test_dspark_config.py b/tests/v1/spec_decode/test_dspark_config.py index 6c230550101e..488a4bece1d7 100644 --- a/tests/v1/spec_decode/test_dspark_config.py +++ b/tests/v1/spec_decode/test_dspark_config.py @@ -3,7 +3,10 @@ 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): @@ -54,3 +57,101 @@ def test_dspark_standard_rejection_uses_probabilistic_draft_sampling(monkeypatch ) 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/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 3ae5125d4858..3115dd212471 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -69,18 +69,22 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): * self.num_query_per_req ) - self.draft_logits_index_mapping: torch.Tensor | None = None + self._step_cols = torch.arange( + self.num_speculative_steps, dtype=torch.int32, device=device + ) if self.speculative_config.draft_sample_method == "probabilistic": - # DSpark keeps only the current step's dense [req, step, vocab] - # logits. Rejection kernels still index by req_state_idx, so this - # maps req_state_idx -> dense row for the active request batch. - self.draft_logits_index_mapping = torch.empty( - self.max_num_reqs, dtype=torch.int32, device=device - ) - self._dense_req_indices = torch.arange( - self.max_num_reqs, dtype=torch.int32, device=device + # 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 current-step draft logits for rejection.") + logger.info("Using DSpark preallocated draft logits for rejection.") def load_draft_model( self, @@ -90,8 +94,9 @@ def load_draft_model( return load_dspark_model(target_model, self.vllm_config) def clear_runtime_draft_logits(self) -> None: - if self.draft_logits_index_mapping is not None: - self.draft_logits = 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. @@ -105,13 +110,6 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) - if self.draft_logits_index_mapping is not None: - self.draft_logits = base_logits - # idx_map is repeated for all speculative steps for a request, so - # the first column identifies the request state for each dense row. - self.draft_logits_index_mapping[idx_map[:, 0].long()] = ( - self._dense_req_indices[:num_reqs] - ) # Anchor (bonus) token per request = the input id at query offset 0, # read via the precomputed persistent index (fixed buffer for capture). @@ -126,7 +124,7 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: logits_i.add_(bias) else: logits_i.copy_(logits_i + bias) - if self.draft_logits_index_mapping is not None: + 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. draft_i = gumbel_sample( @@ -136,7 +134,8 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: self.seeds, sample_pos[:, i] - 1, apply_temperature=True, - output_processed_logits_inplace=True, + output_processed_logits=self.draft_logits, + output_processed_logits_col=self._step_cols[i], use_fp64=self.use_fp64_gumbel, ) else: From 6f386f874a92c8ca915a7324fd8bc83e26d334b5 Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Fri, 10 Jul 2026 19:08:23 +0000 Subject: [PATCH 5/8] fix(dspark): reuse dead MHC input buffer --- tests/kernels/test_mhc_kernels.py | 10 +++++++--- vllm/model_executor/kernels/mhc/tilelang.py | 18 +++++++++++++----- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 140628fbf4a3..5cdba79be026 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -305,11 +305,11 @@ def run_ref(): not HAS_TILELANG_MHC, reason="TileLang MHC support required", ) -def test_mhc_fused_post_pre_reuses_residual_buffer(): +@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) - num_tokens = 128 hidden_size = 4096 hc_mult = 4 x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) @@ -363,7 +363,11 @@ def test_mhc_fused_post_pre_reuses_residual_buffer(): norm_eps=norm_eps, ) - assert actual[0].data_ptr() == residual_reuse.data_ptr() + 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 diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 7a2e1b6ea4c8..23eeda1dc842 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -488,6 +488,7 @@ def _mhc_fused_post_pre_tilelang_impl( 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. @@ -600,11 +601,15 @@ def _mhc_fused_post_pre_tilelang_impl( 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: @@ -740,6 +745,7 @@ def mhc_fused_post_pre_tilelang( norm_weight=norm_weight, norm_eps=norm_eps, reuse_residual_buffer=False, + reuse_input_buffer=False, ) @@ -761,6 +767,7 @@ def mhc_fused_post_pre_tilelang_reuse_residual( 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, @@ -779,6 +786,7 @@ def mhc_fused_post_pre_tilelang_reuse_residual( norm_weight=norm_weight, norm_eps=norm_eps, reuse_residual_buffer=True, + reuse_input_buffer=True, ) From 47efda7e408a832412c2daedab66c161e36e99b8 Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Sun, 12 Jul 2026 00:16:39 +0000 Subject: [PATCH 6/8] test(dspark): update config fixture after upstream sync --- tests/v1/spec_decode/test_dspark_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/spec_decode/test_dspark_config.py b/tests/v1/spec_decode/test_dspark_config.py index 488a4bece1d7..074a44d31264 100644 --- a/tests/v1/spec_decode/test_dspark_config.py +++ b/tests/v1/spec_decode/test_dspark_config.py @@ -28,6 +28,7 @@ def test_dspark_standard_rejection_uses_probabilistic_draft_sampling(monkeypatch dtype="auto", seed=0, tokenizer_revision=None, + hf_overrides={}, max_model_len=4096, enforce_eager=False, max_logprobs=20, From f056ff7105ca25a20cd2ecb9f6228ccb096af2c9 Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Tue, 14 Jul 2026 01:22:41 +0000 Subject: [PATCH 7/8] chore(dspark): remove dead inv-rope-layout and greedy-argmax kernels dspark_inv_rope_bf16_layout and dspark_markov_greedy_argmax (and the two Triton kernels backing the latter) have no production callers anywhere in the tree; the live draft path uses the BF16 WO_A projection and the fused probabilistic Markov sampler (dspark_markov_probs_sample). Drop the dead wrappers, kernels, and their standalone tests. --- tests/v1/spec_decode/test_dspark.py | 87 ------ .../deepseek_v4/nvidia/dspark_triton.py | 282 ------------------ 2 files changed, 369 deletions(-) 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/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, From f02be76748dea1d7dbf8b859015ecafd8b94d87b Mon Sep 17 00:00:00 2001 From: alexbi29 Date: Mon, 27 Jul 2026 08:56:32 +0000 Subject: [PATCH 8/8] fix(config): keep sequence-parallel MoE enabled at DP=1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream #48849 re-added `and self.data_parallel_size > 1` to `ParallelConfig.use_sequence_parallel_moe`, which disables sequence-parallel MoE for every DP=1 deployment — including DeepSeek-V4-Flash served at TP=2 with expert parallelism, where it costs ~5% decode throughput. The rationale comment directly above the property is about replication across the *tensor* parallel group: with the all_reduce at the end of attention, every TP rank holds the same tokens, so under EP the experts redo the same work on every rank. That argument does not depend on data_parallel_size, and the code had drifted away from its own comment. #48849 was motivated by memory — ~5.6 GiB of SP buffers on a pure-TP Nemotron 550B — but it measured only memory, never throughput, and concluded "no benefit for any model" at DP=1. That does not hold here. Measured on DeepSeek-V4-Flash, 2x RTX PRO 6000 Blackwell (SM120), TP=2 + EP, DP=1, FP8 KV, DSpark MTP, at a fixed 250 W cap: C=1 decode (2 sets) 274.6 / 276.7 -> 288.2 / 289.8 tok/s (+5%) prefill 1K 5,948 -> 6,259 tok/s (+5%) prefill 4K 7,518 -> 7,536 tok/s (flat) available KV cache 11.44 GiB -> 11.44 GiB (identical) GPU KV cache tokens 1,189,790 -> 1,189,790 (identical) peak activation 1.89 GiB -> 1.13 GiB i.e. on this configuration the SP buffers cost no KV capacity whatsoever, so the trade #48849 describes is not present and only the slowdown is. Draft acceptance was flat at 80.7% across both configurations, and GPU utilization stayed at 90-92%, so this is expert-path work and not a scheduling or speculation artifact. If the pure-TP memory regression still needs addressing, the gate should be made config-aware (or user-selectable) rather than switching SP-MoE off for all DP=1 deployments. --- vllm/config/parallel.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) 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