From 5003e756e2cc75d6600790a249c8e89e59435a61 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 09:17:54 +0000 Subject: [PATCH] =?UTF-8?q?WIP:=20cross-check=202=20fix=20=E2=80=94=20bloc?= =?UTF-8?q?k-aligned=20compressed=20RoPE=20positions=20+=20compress=5Frope?= =?UTF-8?q?=5Ftheta=20support?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CRITICAL BUG FIX: comp_pos was using LAST position of each block (((bi+1)*r-1)) instead of FIRST position (bi*r). Off by r-1: 3 for CSA, 127 for HCA. vLLM uses (position // ratio) * ratio = block-aligned first position. - Added compress_rope_theta config support (vLLM uses separate theta for compressed) - Added comp_rope_cos/comp_rope_sin param to forward_layer (not yet wired through) Only single_shot_inference.py changed — no kernel code touched. Base commit: 572bdd2 --- TEMP/deepseek_v4_ref/deepseek_v4/__init__.py | 30 + .../deepseek_v4/amd/__init__.py | 2 + TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py | 972 +++++++++++ TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py | 509 ++++++ TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py | 856 ++++++++++ TEMP/deepseek_v4_ref/deepseek_v4/attention.py | 806 +++++++++ .../deepseek_v4/common/__init__.py | 2 + .../deepseek_v4/common/ops/__init__.py | 28 + .../deepseek_v4/common/ops/cache_utils.py | 594 +++++++ .../common/ops/fused_compress_quant_cache.py | 666 ++++++++ .../deepseek_v4/common/ops/fused_indexer_q.py | 438 +++++ .../common/ops/fused_inv_rope_fp8_quant.py | 318 ++++ .../common/ops/fused_mtp_input_rmsnorm.py | 203 +++ .../common/ops/fused_qk_rmsnorm.py | 96 ++ .../common/ops/save_partial_states.py | 101 ++ .../deepseek_v4/common/rope.py | 36 + .../deepseek_v4_ref/deepseek_v4/compressor.py | 380 +++++ .../deepseek_v4/nvidia/__init__.py | 2 + .../deepseek_v4/nvidia/flashmla.py | 424 +++++ .../deepseek_v4/nvidia/model.py | 1476 +++++++++++++++++ .../deepseek_v4_ref/deepseek_v4/nvidia/mtp.py | 516 ++++++ .../deepseek_v4/nvidia/ops/__init__.py | 13 + .../nvidia/ops/dequant_gather_k_cutedsl.py | 331 ++++ .../nvidia/ops/fused_indexer_q_cutedsl.py | 610 +++++++ .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 173 ++ .../ops/sparse_attn_compress_cutedsl.py | 1380 +++++++++++++++ .../deepseek_v4/quant_config.py | 158 ++ single_shot_inference.py | 13 +- 28 files changed, 11132 insertions(+), 1 deletion(-) create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/attention.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/compressor.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py create mode 100644 TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py new file mode 100644 index 00000000..abaa794f --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4 model — hardware-isolated entry point. + +The actual implementation lives under ``nvidia/`` and ``amd/``; this module +picks the right one for the current platform and re-exports the public +classes used by the model registry and quantization config lookup. +""" + +from typing import TYPE_CHECKING + +from vllm.platforms import current_platform + +from .quant_config import DeepseekV4FP8Config + +# Pick the per-platform implementation. The NVIDIA branch is the static +# default that mypy sees; the ROCm branch overrides it at runtime and is +# kept type-compatible via ``# type: ignore[assignment]``. +if TYPE_CHECKING or not current_platform.is_rocm(): + from .nvidia.model import DeepseekV4ForCausalLM + from .nvidia.mtp import DeepSeekV4MTP +else: + from .amd.model import DeepseekV4ForCausalLM # type: ignore[assignment] + from .amd.mtp import DeepSeekV4MTP # type: ignore[assignment] + +__all__ = [ + "DeepSeekV4MTP", + "DeepseekV4FP8Config", + "DeepseekV4ForCausalLM", +] diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py new file mode 100644 index 00000000..208f01a7 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py new file mode 100644 index 00000000..fb724fbe --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py @@ -0,0 +1,972 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import ( + HCHeadOp, + MHCFusedPostPreOp, + MHCPostOp, + MHCPreOp, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.models.deepseek_v4.attention import ( + DeepseekV4Indexer, + DeepseekV4MLA, +) +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.utils.import_utils import has_tilelang + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + swiglu_limit: float | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + + # If is_sequence_parallel, the input and output tensors are sharded + # across the ranks within the tp_group. In this case the weights are + # replicated and no collective ops are needed. + # Otherwise we use standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + if swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.prefix = prefix + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hidden_size = config.hidden_size + + self.n_routed_experts = config.n_routed_experts + self.n_activated_experts = config.num_experts_per_tok + self.moe_intermediate_size = config.moe_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + + self.gate = GateLinear( + input_size=config.hidden_size, + output_size=config.n_routed_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers + self.hash_indices_dtype = torch.int32 + if is_hash_moe: + # hash MoE doesn't use e_score_correction_bias + # Use randint instead of empty to avoid garbage values causing + # invalid memory access in dummy mode (--load-format="dummy") + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=self.hash_indices_dtype, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + + if config.n_shared_experts is None: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = DeepseekV4MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + swiglu_limit=self.swiglu_limit, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + + self.tp_rank = get_tensor_model_parallel_rank() + assert config.n_routed_experts % self.tp_size == 0 + + self.n_local_experts = config.n_routed_experts // self.tp_size + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + swiglu_limit=self.swiglu_limit, + router_logits_dtype=torch.float32, + ) + + def forward( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + + org_shape = hidden_states.shape + if self.experts.is_internal_router: + # In this case, the gate/router runs inside the FusedMoE class + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=hidden_states, + input_ids=input_ids, + ) + else: + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + + return final_hidden_states.view(org_shape) + + +class DeepseekV4Attention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + layer_id = extract_layer_index(prefix) + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0 + + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.max_position_embeddings = config.max_position_embeddings + + # Padded to min 64 heads for FlashMLA, initialized to -inf + # (no sink effect). Weight loading fills the first n_local_heads slots. + padded_heads = max(self.n_local_heads, 64) + self.attn_sink = nn.Parameter( + torch.full((padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + self.softmax_scale = self.head_dim**-0.5 + self.scale_fmt = config.quantization_config["scale_fmt"] + + self.rope_parameters = config.rope_scaling + + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + compress_ratio=self.compress_ratio, + ) + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=vllm_config.cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + ) + + self.mla_attn = DeepseekV4MLA( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, + vllm_config=vllm_config, + fused_wqa_wkv=self.fused_wqa_wkv, + q_norm=self.q_norm, + wq_b=self.wq_b, + kv_norm=self.kv_norm, + wo_a=self.wo_a, + wo_b=self.wo_b, + attn_sink=self.attn_sink, + rotary_emb=self.rotary_emb, + indexer=self.indexer, + indexer_rotary_emb=self.rotary_emb, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + window_size=self.window_size, + compress_ratio=self.compress_ratio, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ): + return self.mla_attn(positions, hidden_states, llama_4_scaling) + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + + # Lazy import to avoid top-level tilelang dependency. + # Registers both torch.ops.vllm.mhc_pre and mhc_post + import vllm.model_executor.layers.mhc # noqa: F401 + + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + + self.rms_norm_eps = config.rms_norm_eps + self.attn = DeepseekV4Attention( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.mhc_pre = MHCPreOp() + self.mhc_post = MHCPostOp() + self.mhc_fused_post_pre = MHCFusedPostPreOp() + self.has_tilelang = has_tilelang() + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ): + post_mix, res_mix, layer_input = self.mhc_pre( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + ) + return layer_input, post_mix, res_mix + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ): + return self.mhc_post(x, residual, post, comb) + + def _forward_fused_post_pre( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if residual is None: + # Run standalone hc_pre on first layer + residual = x + x, post_mix, res_mix = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + else: + residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + ) + + x = self.attn_norm(x) + x = self.attn(positions, x, None) + + residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + return x, residual, post_mix, res_mix + + def _forward_unfused_post_pre( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn(positions, x, None) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + x = self.hc_post(x, residual, post, comb) + return x, None, None, None + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + if not self.has_tilelang: + return self._forward_unfused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) + return self._forward_fused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) + + +class DeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # Three aux streams: one per non-default input GEMM in + # DeepseekV4MLA.attn_gemm_parallel_execute + # (compressor kv_score, indexer.weights_proj, indexer.compressor + # kv_score). fused_wqa_wkv stays on the default stream. + # Disable them on ROCm because of hang issues. + aux_stream_list = ( + None + if current_platform.is_rocm() + else [torch.cuda.Stream() for _ in range(3)] + ) + + self.device = current_platform.device_type + # Reserved topk indices buffer for all Indexer layers to reuse. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_dim, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty( + self.hc_mult, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_op = HCHeadOp() + self.has_tilelang = has_tilelang() + # 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: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + else: + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + # PP intermediate tensors carry the multi-stream hidden_states + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the + # token embedding to hc_mult streams before the first decoder + # layer and keeps that shape until hc_head() collapses it. + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, self.config.hidden_size), + dtype=dtype, + device=device, + ), + } + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + residual, post_mix, res_mix = None, None, None + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + if layer is not None and self.has_tilelang: + hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) + + 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)) + + hidden_states = self.hc_head_op( + 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) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = self.get_expert_mapping() + + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name, self): + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + if ".experts." in name: + # E8M0 scales are stored as float8_e8m0fnu in + # checkpoints but the MoE param is uint8. copy_() + # would do a numeric conversion (e.g. 2^-7 → 0), + # destroying the raw exponent bytes. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + loaded_params.add(name_mapped) + continue + elif "attn_sink" in name: + if is_pp_missing_parameter(name, self): + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + + +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: + if expert_dtype == "fp4": + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as + # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and + # shared experts use Fp8LinearMethod's block scales, which + # register as ``weight_scale_inv``. + scale_regex = { + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", + re.compile(r"\.scale$"): ".weight_scale_inv", + } + else: + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys + # there. + scale_regex = { + re.compile(r"\.scale$"): ".weight_scale_inv", + } + return WeightsMapper( + orig_to_new_prefix={ + "layers.": "model.layers.", + "embed.": "model.embed.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", + }, + orig_to_new_regex=scale_regex, + orig_to_new_suffix={ + "head.weight": "lm_head.weight", + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".attn.compressor.": ".attn.mla_attn.compressor.", + ".shared_experts.w2": ".shared_experts.down_proj", + }, + ) + + +class DeepseekV4ForCausalLM(nn.Module, SupportsPP): + model_cls = DeepseekV4Model + + # Default mapper assumes the original FP4-expert checkpoint layout. + # Overridden per-instance in __init__ when expert_dtype != "fp4". + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + self.config = config + expert_dtype = getattr(config, "expert_dtype", "fp4") + if expert_dtype != "fp4": + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) + + self.model = self.model_cls( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer (max_num_batched_tokens, + hc_mult * hidden_size) for the MTP draft model. Populated by + forward(); valid after each target step.""" + return getattr(self.model, "_mtp_hidden_buffer", None) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py new file mode 100644 index 00000000..5938cde6 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py @@ -0,0 +1,509 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4). + +Split from ``deepseek_mtp.py`` because the V4 architecture introduces several +pieces that have no analogue in V3/V32: + * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of + the fused ``eh_proj``); + * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``; + * ``DeepseekV4DecoderLayer`` with its own aux-stream management; + * V4-specific checkpoint weight-name remapping in ``load_weights``. +""" + +import typing +from collections.abc import Callable, Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name +from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, + mtp_shared_head_rmsnorm, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.utils.import_utils import has_tilelang + +from .model import DeepseekV4DecoderLayer + +logger = init_logger(__name__) + +# MoE expert scales are fused into per-layer w13/w2 tensors. The exact +# parameter suffix depends on which FusedMoE method handles the experts: +# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; +# - fp8 experts (Fp8MoEMethod with block_quant=True) register +# ``w{1,2,3}_weight_scale_inv``. +# Other FP8 linear scales (including shared experts) always use +# ``.weight_scale_inv``. Mirrors the per-instance mapper built by +# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py. +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DeepSeekV4MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + topk_indices_buffer: torch.Tensor, + prefix: str, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ) -> None: + super().__init__() + + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + self.rms_norm_eps = config.rms_norm_eps + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than + # fusing them the way V3 does with eh_proj. + self.e_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + return_bias=False, + quant_config=quant_config, + ) + self.h_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + return_bias=False, + quant_config=quant_config, + ) + + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + self.mtp_block = DeepseekV4DecoderLayer( + vllm_config, + prefix, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + + self.hc_head_op = HCHeadOp() + self.has_tilelang = has_tilelang() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + # Target stashes pre-hc_head residual as flat (T, hc_mult * D); + # reshape to (T, hc_mult, D) — the training-time layout — before + # the fused norm pass so both inputs are 3D-friendly. + previous_hidden_states = previous_hidden_states.view( + -1, self.hc_mult, self.config.hidden_size + ) + # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. + inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( + inputs_embeds, + positions, + previous_hidden_states, + self.enorm.weight.data, + self.hnorm.weight.data, + self.enorm.variance_epsilon, + self.hc_mult, + ) + hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( + inputs_embeds + ).unsqueeze(-2) + hidden_states, residual, post_mix, res_mix = self.mtp_block( + positions=positions, x=hidden_states, input_ids=None + ) + if self.has_tilelang: + hidden_states = self.mtp_block.hc_post( + hidden_states, residual, post_mix, res_mix + ) + # Return the flat pre-hc_head residual so it can be re-fed as the + # next spec step's `previous_hidden_states` when + # num_speculative_tokens > 1. hc_head is deferred to compute_logits. + return hidden_states.flatten(1) + + +class DeepSeekV4MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.device = current_platform.device_type + + topk_tokens = config.index_topk + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + topk_tokens, + dtype=torch.int32, + device=self.device, + ) + + # Three aux streams shared across all MTP layers, mirroring + # DeepseekV4Model. ROCm runs the same work serially for now. + aux_stream_list = ( + None + if current_platform.is_rocm() + else [torch.cuda.Stream() for _ in range(3)] + ) + + # to map the exact layer index from weights + self.layers = torch.nn.ModuleDict( + { + str(idx): DeepSeekV4MultiTokenPredictorLayer( + vllm_config, + self.topk_indices_buffer, + f"{prefix}.layers.{idx}", + aux_stream_list=aux_stream_list, + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply + # hc_head here so logits are computed from the dense hidden state. + hidden_states = hidden_states.view( + -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size + ) + hidden_states = mtp_layer.hc_head_op( + hidden_states, + mtp_layer.hc_head_fn, + mtp_layer.hc_head_scale, + mtp_layer.hc_head_base, + mtp_layer.rms_norm_eps, + mtp_layer.hc_eps, + ) + hidden_states = mtp_shared_head_rmsnorm( + hidden_states, + mtp_layer.shared_head.norm.weight.data, + mtp_layer.shared_head.norm.variance_epsilon, + ) + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class DeepSeekV4MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = DeepSeekV4MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Weight name remapping for checkpoint compatibility. + # Maps checkpoint weight paths to model parameter paths. + WEIGHT_NAME_REMAPPING: dict[str, str] = { + ".emb.tok_emb.weight": ".embed_tokens.weight", + ".head.weight": ".shared_head.head.weight", + ".norm.weight": ".shared_head.norm.weight", + } + + def _remap_weight_name(name: str) -> str: + """Remap checkpoint weight names to model parameter names.""" + for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items(): + if old_pattern in name: + name = name.replace(old_pattern, new_pattern) + return name + + def _find_mtp_layer_idx(name: str) -> int: + subnames = name.split(".") + for subname in subnames: + try: + # we return the first encountered integer + return int(subname) + except ValueError: + continue + return 0 + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + + # FP8 experts register ``..._weight_scale_inv`` (block_quant) while + # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix + # for the rename below based on the model's expert dtype. + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + for name, loaded_weight in weights: + mtp_layer_idx = _find_mtp_layer_idx(name) + # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to + # `model.layers.{num_hidden_layers + i}.*` so that + # get_spec_layer_idx_from_weight_name can identify them. + name = name.replace( + f"mtp.{mtp_layer_idx}.", + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", + ) + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + + name = _remap_weight_name(name) + name = self._rewrite_spec_layer_name(spec_layer, name) + + if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name: + continue + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + if ".experts." in name: + # Reinterpret E8M0 scales as uint8 to preserve raw + # exponent bytes; numeric copy_() would zero them. + # Mirrors the main DeepseekV4 loader. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + loaded_params.add(name_mapped) + break + continue + elif "attn_sink" in name: + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if ".shared_experts.w2" in name: + name = name.replace( + ".shared_experts.w2", ".shared_experts.down_proj" + ) + if name.endswith(".ffn.gate.bias"): + name = name.replace(".bias", ".e_score_correction_bias") + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may have " + f"been quantized without including the MTP layers. " + f"Use a checkpoint that includes MTP layer weights, " + f"or disable speculative decoding." + ) + logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """ + Rewrite the weight name to match the format of the original model. + Add .mtp_block for modules in transformer layer block for spec layer + and rename shared layer weights to be top level. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "h_proj", + "e_proj", + "shared_head", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + # treat rest weights as weights for transformer layer block + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + ) + elif shared_weight: + # treat shared weights as top level weights + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py b/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py new file mode 100644 index 00000000..2af93fba --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py @@ -0,0 +1,856 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache +from vllm.models.deepseek_v4.nvidia.flashmla import ( + DeepseekV4FlashMLASparseBackend, + DeepseekV4SparseMLAAttentionImpl, +) +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import ( + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + FlashMLASparseMetadata, + FlashMLASparseMetadataBuilder, +) +from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + DeepseekSparseSWAMetadataBuilder, +) +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_ragged_indices_from_dense, + rocm_sparse_attn_decode, + rocm_sparse_attn_prefill, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.models.deepseek_v4.attention import ( + DeepseekV4MLAAttention, + ) + + +def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: + lengths = lengths.to(dtype=torch.int32).contiguous() + indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) + torch.cumsum(lengths, dim=0, out=indptr[1:]) + return indptr + + +# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes +# do not touch the shared DeepSeek V4 cache utilities. +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 + + +@triton.jit +def _combine_topk_swa_indices_kernel( + combined_indices_ptr, + combined_indices_stride, + combined_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + query_start_loc_ptr, + seq_lens_ptr, + gather_lens_ptr, + M, + N, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + TOPK_WIDTH: tl.constexpr, + PADDED_TOP_K: tl.constexpr, +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + gather_len = tl.load(gather_lens_ptr + batch_idx) + start_pos = seq_len - query_len + gather_start = seq_len - gather_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + + topk_offset = tl.arange(0, PADDED_TOP_K) + topk_mask = topk_offset < topk_len + safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0) + topk_indices = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset, + mask=topk_mask, + other=-1, + ) + valid_topk = (topk_indices >= 0) & (topk_indices < N) + topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1) + tl.store( + combined_indices_ptr + token_idx * combined_indices_stride + topk_offset, + topk_indices, + mask=topk_mask, + ) + + swa_offset = tl.arange(0, WINDOW_SIZE) + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + swa_offset, + M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, + mask=swa_offset < swa_len, + ) + + tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) + + +def combine_topk_swa_indices( + topk_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor, + window_size: int, + compress_ratio: int, + topk: int, + M: int, + N: int, +) -> tuple[torch.Tensor, torch.Tensor]: + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() + num_tokens = topk_indices.shape[0] + num_reqs = seq_lens.shape[0] + combined_topk = ( + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + // _SPARSE_PREFILL_TOPK_ALIGNMENT + * _SPARSE_PREFILL_TOPK_ALIGNMENT + ) + combined_indices = torch.full( + (num_tokens, combined_topk), + fill_value=-1, + dtype=torch.int32, + device=topk_indices.device, + ) + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + + num_workers = 128 + _combine_topk_swa_indices_kernel[(num_reqs, num_workers)]( + combined_indices, + combined_indices.stride(0), + combined_lens, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=window_size, + TOPK_WIDTH=topk_indices.shape[-1], + PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + ) + return combined_indices, combined_lens + + +@triton.jit +def _compute_topk_lens_kernel( + topk_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + topk, + is_valid_token_ptr, + TRITON_BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0) + is_valid_token = tl.load(is_valid_token_ptr + token_idx) + + count = tl.zeros((), dtype=tl.int32) + for i in range(0, topk, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + mask = offset < topk + local_idx = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + other=-1, + ) + count += tl.sum((local_idx >= 0).to(tl.int32), axis=0) + + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) + + +@triton.jit +def _pack_global_topk_ragged_kernel( + global_topk_ragged_ptr, + topk_indptr_ptr, + topk_indices_ptr, + topk_indices_stride, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + topk, + BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0) + block_idx = tl.program_id(1) + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + + out_start = tl.load(topk_indptr_ptr + token_idx) + out_end = tl.load(topk_indptr_ptr + token_idx + 1) + out_len = out_end - out_start + if block_idx * BLOCK_SIZE >= out_len: + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + mask = (offset < out_len) & (offset < topk) + local_idx = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + other=-1, + ) + valid = mask & (local_idx >= 0) + block_indices = local_idx // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=valid, + other=0, + ) + block_offsets = local_idx % block_size + slot_ids = tl.where(valid, block_numbers * block_size + block_offsets, -1) + tl.store(global_topk_ragged_ptr + out_start + offset, slot_ids, mask=mask) + + +def compute_global_topk_ragged_indices_and_indptr( + topk_indices: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + is_valid_token: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() + num_tokens = topk_indices.shape[0] + topk = topk_indices.shape[1] + + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + _compute_topk_lens_kernel[(num_tokens,)]( + topk_lens, + topk_indices, + topk_indices.stride(0), + topk, + is_valid_token, + TRITON_BLOCK_SIZE=1024, + ) + + topk_indptr = _build_indptr_from_lengths(topk_lens) + global_topk_ragged = torch.empty( + num_tokens * topk, + dtype=torch.int32, + device=topk_indices.device, + ) + if global_topk_ragged.numel() > 0: + block = 128 + _pack_global_topk_ragged_kernel[(num_tokens, triton.cdiv(topk, block))]( + global_topk_ragged, + topk_indptr, + topk_indices, + topk_indices.stride(0), + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + topk, + BLOCK_SIZE=block, + ) + return global_topk_ragged, topk_indptr, topk_lens + + +@triton.jit +def _compute_combined_lens_kernel( + combined_lens_ptr, + query_start_loc_ptr, + seq_lens_ptr, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + start_pos = seq_len - query_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) + + +@triton.jit +def _combine_topk_swa_indices_ragged_kernel( + combined_ragged_ptr, + combined_indptr_ptr, + topk_indices_ptr, + topk_indices_stride, + query_start_loc_ptr, + seq_lens_ptr, + gather_lens_ptr, + M, + N, + topk_width, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + block_idx = tl.program_id(2) + num_workers = tl.num_programs(1) + + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + gather_len = tl.load(gather_lens_ptr + batch_idx) + start_pos = seq_len - query_len + gather_start = seq_len - gather_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + combined_len = topk_len + swa_len + + offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + if block_idx * BLOCK_SIZE < combined_len: + out_start = tl.load(combined_indptr_ptr + token_idx) + topk_mask = (offset < topk_len) & (offset < topk_width) + topk_vals = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=topk_mask, + other=-1, + ) + tl.store( + combined_ragged_ptr + out_start + offset, + topk_vals + M * batch_idx, + mask=topk_mask, + ) + + swa_offset = offset - topk_len + swa_mask = (offset >= topk_len) & (swa_offset < swa_len) + tl.store( + combined_ragged_ptr + out_start + offset, + M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, + mask=swa_mask, + ) + + +def combine_topk_swa_indices_ragged( + topk_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor, + window_size: int, + compress_ratio: int, + topk: int, + M: int, + N: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() + num_tokens = topk_indices.shape[0] + num_reqs = seq_lens.shape[0] + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + + num_workers = 128 + _compute_combined_lens_kernel[(num_reqs, num_workers)]( + combined_lens, + query_start_loc, + seq_lens, + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=window_size, + ) + + combined_indptr = _build_indptr_from_lengths(combined_lens) + combined_ragged = torch.empty( + num_tokens * (topk + window_size), + dtype=torch.int32, + device=topk_indices.device, + ) + if combined_ragged.numel() > 0: + block = 128 + _combine_topk_swa_indices_ragged_kernel[ + (num_reqs, num_workers, triton.cdiv(topk + window_size, block)) + ]( + combined_ragged, + combined_indptr, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + topk_indices.shape[-1], + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=window_size, + BLOCK_SIZE=block, + ) + return combined_ragged, combined_indptr, combined_lens + + +def _copy_ragged_to_graph_buffers( + ragged_indices: torch.Tensor, + ragged_indptr: torch.Tensor, + ragged_indices_buffer: torch.Tensor, + ragged_indptr_buffer: torch.Tensor, + num_rows: int, + max_entries_per_row: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Copy dynamic ragged metadata into persistent CUDA graph buffers. + + FULL decode graphs capture kernel argument addresses. Keep the returned + tensors backed by stable storage, while indptr continues to bound reads. + """ + indptr_out = ragged_indptr_buffer[: num_rows + 1] + indptr_out.copy_(ragged_indptr, non_blocking=True) + + max_entries = max(num_rows * max_entries_per_row, 1) + ragged_out = ragged_indices_buffer[:max_entries] + nnz = ragged_indices.numel() + if nnz > 0: + ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) + return ragged_out, indptr_out + + +@dataclass +class DeepseekV4ROCMAiterMLASparseMetadata(FlashMLASparseMetadata): + """ROCm-specific DeepSeek V4 metadata carrying ragged decode topk.""" + + c128a_decode_topk_ragged_indices: torch.Tensor | None = None + c128a_decode_topk_ragged_indptr: torch.Tensor | None = None + + +@dataclass +class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): + decode_swa_ragged_indices: torch.Tensor | None = None + decode_swa_ragged_indptr: torch.Tensor | None = None + + +class DeepseekV4ROCMAiterMLASparseMetadataBuilder(FlashMLASparseMetadataBuilder): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None + self.c128a_decode_topk_ragged_indptr_buffer: torch.Tensor | None = None + if self.is_deepseek_v4 and self.compress_ratio == 128: + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + self.c128a_decode_topk_ragged_indices_buffer = torch.empty( + max_tokens * self.c128a_max_compressed, + dtype=torch.int32, + device=self.device, + ) + self.c128a_decode_topk_ragged_indptr_buffer = torch.empty( + max_tokens + 1, + dtype=torch.int32, + device=self.device, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV4ROCMAiterMLASparseMetadata: + base = super().build( + common_prefix_len=common_prefix_len, + common_attn_metadata=common_attn_metadata, + fast_build=fast_build, + ) + + ragged_indices = None + ragged_indptr = None + dense_decode = base.c128a_global_decode_topk_indices + decode_lens = base.c128a_decode_topk_lens + if dense_decode is not None and decode_lens is not None: + ragged_indices, ragged_indptr = build_ragged_indices_from_dense( + dense_decode.reshape(dense_decode.shape[0], -1), + decode_lens, + ) + assert self.c128a_decode_topk_ragged_indices_buffer is not None + assert self.c128a_decode_topk_ragged_indptr_buffer is not None + ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( + ragged_indices, + ragged_indptr, + self.c128a_decode_topk_ragged_indices_buffer, + self.c128a_decode_topk_ragged_indptr_buffer, + dense_decode.shape[0], + self.c128a_max_compressed, + ) + + return DeepseekV4ROCMAiterMLASparseMetadata( + **vars(base), + c128a_decode_topk_ragged_indices=ragged_indices, + c128a_decode_topk_ragged_indptr=ragged_indptr, + ) + + +class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + self.decode_swa_ragged_indices_buffer = torch.empty( + max_tokens * self.window_size, + dtype=torch.int32, + device=self.device, + ) + self.decode_swa_ragged_indptr_buffer = torch.empty( + max_tokens + 1, + dtype=torch.int32, + device=self.device, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV4ROCMAiterSparseSWAMetadata: + base = super().build( + common_prefix_len=common_prefix_len, + common_attn_metadata=common_attn_metadata, + fast_build=fast_build, + ) + + ragged_indices = None + ragged_indptr = None + if ( + base.num_decode_tokens > 0 + and base.decode_swa_indices is not None + and base.decode_swa_lens is not None + ): + ragged_indices, ragged_indptr = build_ragged_indices_from_dense( + base.decode_swa_indices.reshape(base.num_decode_tokens, -1), + base.decode_swa_lens, + ) + ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( + ragged_indices, + ragged_indptr, + self.decode_swa_ragged_indices_buffer, + self.decode_swa_ragged_indptr_buffer, + base.num_decode_tokens, + self.window_size, + ) + + return DeepseekV4ROCMAiterSparseSWAMetadata( + **vars(base), + decode_swa_ragged_indices=ragged_indices, + decode_swa_ragged_indptr=ragged_indptr, + ) + + +class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4FlashMLASparseBackend): + @staticmethod + def get_name() -> str: + return "ROCM_V4_FLASHMLA_SPARSE" + + @staticmethod + def get_builder_cls() -> type["DeepseekV4ROCMAiterMLASparseMetadataBuilder"]: + return DeepseekV4ROCMAiterMLASparseMetadataBuilder + + @staticmethod + def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]: + return DeepseekV4ROCMAiterMLASparseImpl + + +class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): + """ROCm sparse MLA implementation used by DeepSeek V4's custom MLA layer.""" + + backend_cls = DeepseekV4ROCMAiterMLASparseBackend + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + return num_heads + + @classmethod + def forward_mqa( # type: ignore[override] + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + if attn_metadata is None: + # Warmup dummy run: no real metadata. Reserve the same bf16 + # gather workspace _forward_prefill would; the dequantize / topk + # / sparse_fwd kernels are skipped this step. + swa_only = layer.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (layer.max_model_len + layer.compress_ratio - 1) + // layer.compress_ratio + ) + M = N + layer.window_size + layer.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) + output.zero_() + return + + assert isinstance(attn_metadata, dict) + rocm_metadata = cast( + DeepseekV4ROCMAiterMLASparseMetadata | None, + attn_metadata.get(layer.prefix), + ) + swa_metadata = cast( + DeepseekV4ROCMAiterSparseSWAMetadata | None, + attn_metadata.get(layer.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = layer.compress_ratio <= 1 + self_kv_cache = layer.kv_cache if not swa_only else None + swa_kv_cache = layer.swa_cache_layer.kv_cache + + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + cls._forward_prefill( + layer=layer, + q=q[num_decode_tokens:], + positions=positions[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=rocm_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + cls._forward_decode( + layer=layer, + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=rocm_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + @classmethod + def _forward_decode( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv_cache: torch.Tensor | None, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + topk_ragged_indices = None + topk_ragged_indptr = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // layer.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if layer.compress_ratio == 4: + assert layer.topk_indices_buffer is not None + ( + topk_ragged_indices, + topk_ragged_indptr, + topk_lens, + ) = compute_global_topk_ragged_indices_and_indptr( + layer.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + else: + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices + topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr + + rocm_sparse_attn_decode( + q=q, + kv_cache=kv_cache, + swa_k_cache=layer.swa_cache_layer.kv_cache, + swa_only=swa_only, + topk_indices=topk_indices, + topk_lens=topk_lens, + swa_indices=swa_metadata.decode_swa_indices, + swa_lens=swa_metadata.decode_swa_lens, + swa_ragged_indices=swa_metadata.decode_swa_ragged_indices, + swa_ragged_indptr=swa_metadata.decode_swa_ragged_indptr, + topk_ragged_indices=topk_ragged_indices, + topk_ragged_indptr=topk_ragged_indptr, + attn_sink=layer.attn_sink, + scale=layer.scale, + head_dim=layer.head_dim, + nope_head_dim=layer.nope_head_dim, + rope_head_dim=layer.rope_head_dim, + output=output, + ) + + @classmethod + def _forward_prefill( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + ) -> None: + swa_only = attn_metadata is None + + num_prefills = swa_metadata.num_prefills + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + assert seq_lens is not None + assert gather_lens is not None + + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert query_start_loc_cpu is not None + assert query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if layer.compress_ratio == 4: + assert layer.topk_indices_buffer is not None + topk_indices = layer.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + assert topk_indices is not None + top_k = topk_indices.shape[-1] + N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio + else: + assert layer.topk_indices_buffer is not None + topk_indices = layer.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + N = 0 + + M = N + layer.window_size + layer.max_num_batched_tokens + num_chunks = (num_prefills + cls.PREFILL_CHUNK_SIZE - 1) // ( + cls.PREFILL_CHUNK_SIZE + ) + + workspace_manager = current_workspace_manager() + kv = workspace_manager.get_simultaneous( + ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + )[0] + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE + chunk_end = min(chunk_start + cls.PREFILL_CHUNK_SIZE, num_prefills) + chunk_size = chunk_end - chunk_start + if not swa_only: + assert attn_metadata is not None + assert compressed_k_cache is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // layer.compress_ratio, + offset=0, + ) + + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=N, + ) + + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + layer.window_size, + layer.compress_ratio, + top_k, + M, + N, + ) + rocm_sparse_attn_prefill( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices, + topk_length=combined_lens, + scale=layer.scale, + head_dim=layer.head_dim, + nope_head_dim=layer.nope_head_dim, + rope_head_dim=layer.rope_head_dim, + attn_sink=layer.attn_sink, + output=output[query_start:query_end], + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/attention.py b/TEMP/deepseek_v4_ref/deepseek_v4/attention.py new file mode 100644 index 00000000..55cb3d94 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/attention.py @@ -0,0 +1,806 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DeepseekV4 MLA Attention Layer +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import DeepseekV2Config, DeepseekV3Config + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.model_executor.layers.linear import ( + ReplicatedLinear, +) +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.deepseek_v4.common.ops import ( + fused_indexer_q_rope_quant, + fused_inv_rope_fp8_quant, + fused_q_kv_rmsnorm, +) +from vllm.utils.deep_gemm import fp8_einsum +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + ) + +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.input_quant_fp8 import ( + QuantFP8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, +) +from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.platforms import current_platform +from vllm.utils.multi_stream_utils import ( + execute_in_parallel, + maybe_execute_in_parallel, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + FlashMLASparseBackend, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + get_max_prefill_buffer_size, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec + +if TYPE_CHECKING: + from vllm.models.deepseek_v4.nvidia.flashmla import ( + DeepseekV4SparseMLAAttentionImpl, + ) + +logger = init_logger(__name__) + + +def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]": + """Pick the platform-specific V4 sparse MLA impl class. Sole platform check.""" + if current_platform.is_rocm(): + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLASparseImpl, + ) + + return DeepseekV4ROCMAiterMLASparseImpl + from vllm.models.deepseek_v4.nvidia.flashmla import ( + DeepseekV4FlashMLASparseImpl, + ) + + return DeepseekV4FlashMLASparseImpl + + +class DeepseekV4MLA(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + head_dim: int, + scale: float, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + o_lora_rank: int | None, + vllm_config: VllmConfig, + fused_wqa_wkv: torch.nn.Module, + q_norm: torch.nn.Module, + wq_b: torch.nn.Module, + kv_norm: torch.nn.Module, + wo_a: torch.nn.Module, + wo_b: torch.nn.Module, + attn_sink: torch.nn.Module, + rotary_emb: torch.nn.Module, + indexer: torch.nn.Module | None, + indexer_rotary_emb: torch.nn.Module, + topk_indices_buffer: torch.Tensor | None, + aux_stream_list: list[torch.cuda.Stream] | None, + window_size: int, + compress_ratio: int | None, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.n_local_heads = num_heads + self.head_dim = head_dim + self.scale = scale + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.window_size = window_size + self.compress_ratio = compress_ratio if compress_ratio is not None else 1 + self.prefix = prefix + + # Extract config from vllm_config + config = vllm_config.model_config.hf_config + tp_size = get_tensor_model_parallel_world_size() + + # DeepseekV4-specific attributes (num_heads is already TP-adjusted) + self.eps = config.rms_norm_eps + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = head_dim - self.rope_head_dim + self.n_local_groups = config.o_groups // tp_size + self.o_lora_rank = config.o_lora_rank + + # Store projection modules + self.fused_wqa_wkv = fused_wqa_wkv + self.q_norm = q_norm + self.wq_b = wq_b + + self.kv_norm = kv_norm + self.wo_a = wo_a + + self._wo_a_act_quant = QuantFP8( + static=False, + group_shape=GroupShape(1, 128), + use_ue8m0=True, + ) + # Bypass packed-for-deepgemm path — we need FP32 scales (not packed + # INT32) so fp8_einsum can handle layout transform internally. + self._wo_a_act_quant.use_deep_gemm_supported = False + self.wo_b = wo_b + + # Pick fp8_einsum recipe based on GPU arch: + # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128 + # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1 + cap = current_platform.get_device_capability() + assert cap is not None, "DeepseekV4 attention requires a CUDA device" + self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) + self._tma_aligned_scales = cap.major >= 10 + + self.rotary_emb = rotary_emb + self.indexer_rotary_emb = indexer_rotary_emb + self.topk_indices_buffer = topk_indices_buffer + + self.indexer = indexer + + # Per-head RMS normalization for Q (no learnable weights) + self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False) + + # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic + head_bytes = ( + self.nope_head_dim # 448 fp8 NoPE + + self.rope_head_dim * 2 # 64 bf16 RoPE + + self.nope_head_dim // 64 # 7B scale factors + + 1 # 1B pad + ) + + # Will be None on ROCm for now. + self.aux_stream_list = aux_stream_list + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins + # before post-GEMM starts. + self.ln_events = [torch.cuda.Event() for _ in range(4)] + + assert cache_config is not None, "DeepseekV4 attention requires cache_config" + self.swa_cache_layer = DeepseekV4SWACache( + head_dim=self.head_dim, + window_size=self.window_size, + dtype=torch.uint8, + prefix=f"{prefix}.swa_cache", + cache_config=cache_config, + ) + + self.mla_attn = DeepseekV4MLAAttention( + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + compress_ratio=self.compress_ratio, + window_size=self.window_size, + head_bytes=head_bytes, + swa_cache_layer=self.swa_cache_layer, + attn_sink=attn_sink, # already padded with -inf + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + indexer=self.indexer, + topk_indices_buffer=self.topk_indices_buffer, + ) + # Mirror the inner layer's padded head count (single source of truth). + self.padded_heads = self.mla_attn.padded_heads + + # Create the compressor for layers with compress_ratio > 1; after + # creating the DeepseekV4MLAAttention layer to get its cache. + self.compressor = None + if self.compress_ratio > 1: + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=self.hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.mla_attn.prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # attention_impl is wrapped with @eager_break_during_capture: this is + # where the breakable cudagraph capture breaks (the attention op runs + # eagerly between captured graph segments). + self.attention_impl(hidden_states, positions, o_padded) + o = o_padded[:, : self.n_local_heads, :] + + # Keep ROCm on the BF16 reference wo_a path util kernel ready. + if current_platform.is_rocm(): + z = rocm_inv_rope_einsum( + self.rotary_emb, + o, + positions, + self.rope_head_dim, + self.n_local_groups, + self.o_lora_rank, + self.wo_a, + ) + return self.wo_b(z.flatten(1)) + + # O projection: inverse RoPE + FP8 quant + einsum + wo_b + o_fp8, o_scale = fused_inv_rope_fp8_quant( + o, + positions, + self.rotary_emb.cos_sin_cache, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + tma_aligned_scales=self._tma_aligned_scales, + ) + + wo_a_fp8 = self.wo_a.weight + wo_a_scale = self.wo_a.weight_scale_inv + + z = torch.empty( + (num_tokens, self.n_local_groups, self.o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + fp8_einsum( + "bhr,hdr->bhd", + (o_fp8, o_scale), + (wo_a_fp8, wo_a_scale), + z, + recipe=self._einsum_recipe, + ) + + return self.wo_b(z.flatten(1)) + + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: + aux_streams = self.aux_stream_list + if aux_streams is not None: + assert len(aux_streams) >= 3 + aux_streams = aux_streams[:3] + + # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs + # on aux streams 0..2 when their owning module exists. ln_events[0] + # is the fan-out start event; ln_events[1..3] are per-aux done events. + # On ROCm, aux_streams is None and execute_in_parallel runs serially. + aux_fns: list[Callable[[], Any] | None] = [None, None, None] + + if self.compressor is not None: + # Local ref so the closure keeps a non-None type for mypy. + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[0] = compressor_kv_score + + if self.indexer is not None: + indexer = self.indexer + + def indexer_weights_proj() -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + def indexer_compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[1] = indexer_weights_proj + aux_fns[2] = indexer_compressor_kv_score + + def fused_wqa_wkv() -> torch.Tensor: + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + + qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( + fused_wqa_wkv, + aux_fns, + self.ln_events[0], + self.ln_events[1:4], + aux_streams, + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, + ) + + return qr_kv, kv_score, indexer_kv_score, indexer_weights + + @eager_break_during_capture + def attention_impl( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place + ) -> None: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self.attn_gemm_parallel_execute(hidden_states) + ) + + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride + # on the default stream so q stays on its consumer stream (mla_attn + # downstream reads q on default). Indexer/compressor go on aux for + # overlap with default's GEMM + cache write. + if self.indexer is not None: + aux_streams = self.aux_stream_list + indexer = self.indexer + # Local ref so the closure keeps a non-None type for mypy. + assert self.compressor is not None + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs + # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the + # MLA compressor. Slot [2] is reserved for the indexer's inner + # overlap. ROCm (aux_streams is None) falls back to sequential. + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + self.ln_events[0], + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, + ) + elif self.compressor is not None: + # wq_b + kv_insert on default, compressor on aux. + aux_stream = ( + self.aux_stream_list[0] if self.aux_stream_list is not None else None + ) + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + # SWA-only layer: no compressor, no overlap. + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + # MLA attention writes into the pre-allocated `out` buffer + # ([num_tokens, padded_heads, head_dim]). + self.mla_attn(q, kv, positions, output=out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> torch.Tensor: + if not isinstance(attn_metadata, dict): + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + + # Horizontally fused: + # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE, + # with zero-fill for the padding head slots. The kernel + # allocates and returns the padded q tensor. + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert + # kv is unchanged; mla_attn reads kv solely via swa_kv_cache. + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions.to(torch.int64), + self.rotary_emb.cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + + +class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): + def __init__( + self, + num_heads: int, + head_dim: int, + scale: float, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + compress_ratio: int, + window_size: int, + head_bytes: int, + swa_cache_layer: DeepseekV4SWACache, + attn_sink: torch.Tensor, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + # Sparse MLA Args + indexer: object | None = None, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream: torch.cuda.Stream | None = None, + **extra_impl_args, + ) -> None: + super().__init__() + self.impl_cls = _select_v4_sparse_impl() + self.backend_cls = self.impl_cls.backend_cls + self.num_heads = num_heads + self.num_kv_heads = 1 + self.head_dim = head_dim + self.scale = scale + self.window_size = window_size + self.head_bytes = head_bytes + self.compress_ratio = compress_ratio + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.nope_head_dim = qk_nope_head_dim + self.rope_head_dim = qk_rope_head_dim + self.indexer = indexer + self.topk_indices_buffer = topk_indices_buffer + + self.prefix = prefix # Alias for compatibility with compressor + + self.aux_stream = aux_stream + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + + # Padded Q head count is dictated by the selected impl. + self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads) + + # Store attention sink + assert attn_sink is not None + self.attn_sink: torch.Tensor = attn_sink + # Store SWA cache + assert swa_cache_layer is not None + self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer + + # Get vllm config for cache setup + vllm_config = get_current_vllm_config() + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + # DeepseekV4 only supports fp8 kv-cache format for now. + kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8" + + assert kv_cache_dtype.startswith("fp8"), ( + f"DeepseekV4 only supports fp8 kv-cache format for now, " + f"got {kv_cache_dtype}" + ) + assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), ( + "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now" + ) + # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format + # Automatically convert fp8 kv-cache format to "fp8_ds_mla" + if ( + issubclass(self.get_attn_backend(), FlashMLASparseBackend) + and kv_cache_dtype.startswith("fp8") + and kv_cache_dtype != "fp8_ds_mla" + ): + assert cache_config is not None + cache_config.cache_dtype = "fp8_ds_mla" + kv_cache_dtype = "fp8_ds_mla" + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + + self.kv_cache_dtype = kv_cache_dtype + + # Register with compilation context for metadata lookup + compilation_config = vllm_config.compilation_config + if prefix and prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + if prefix: + compilation_config.static_forward_context[prefix] = self + + self.kv_cache = torch.tensor([]) + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.backend_cls + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + if ( + self.compress_ratio <= 1 + ): # SWA part. Allocated separately as DeepseekV4SWACache. + return None + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8, + compress_ratio=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, + alignment=576, # NOTE: FlashMLA requires 576B alignment + model_version="deepseek_v4", + ) + + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + self.impl_cls.forward_mqa(self, q, kv, positions, output) + + +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + compress_ratio: int = 1, + ): + super().__init__() + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + self.compress_ratio = compress_ratio + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # head_dim already carries the fp8 scale padding + # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with + # the indexer's compressor state cache. V3.2 keeps the legacy layout. + alignment=576, + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4IndexerBackend + + +class DeepseekV4Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + compress_ratio: int = 1, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.q_lora_rank = q_lora_rank # 1536 + self.compress_ratio = compress_ratio + self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + logger.info_once( + "Using %s indexer cache for Lightning Indexer.", + "MXFP4" if self.use_fp4_kv else "FP8", + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + bias=False, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) + self.softmax_scale = self.head_dim**-0.5 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + self.max_model_len = ( + vllm_config.model_config.max_model_len // self.compress_ratio + ) + self.prefix = prefix + + self.max_total_seq_len = ( + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio + ) + + assert cache_config is not None, "Deepseek V4 indexer requires cache_config" + # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + # For FP4 indexer cache, we still allocate the same amount of memory as FP8, + # but only use the first half of the memory. + k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 + self.k_cache = DeepseekV4IndexerCache( + head_dim=k_cache_head_dim, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + ) + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.k_cache.prefix, + use_fp4_cache=self.use_fp4_kv, + ) + + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + ) + + # None on ROCm — maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> torch.Tensor: + compressor = self.compressor + + def wq_b_and_q_quant(): + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + return fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + ) + + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), k = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + return self.indexer_op(hidden_states, q_quant, k, weights) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py new file mode 100644 index 00000000..208f01a7 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py new file mode 100644 index 00000000..dc6f3c60 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .cache_utils import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, + quantize_and_insert_k_cache, +) +from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant +from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant +from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm +from .fused_qk_rmsnorm import fused_q_kv_rmsnorm +from .save_partial_states import save_partial_states + +__all__ = [ + "MXFP4_BLOCK_SIZE", + "combine_topk_swa_indices", + "compute_global_topk_indices_and_lens", + "dequantize_and_gather_k_cache", + "fused_indexer_q_rope_quant", + "fused_inv_rope_fp8_quant", + "fused_mtp_input_rmsnorm", + "fused_q_kv_rmsnorm", + "mtp_shared_head_rmsnorm", + "quantize_and_insert_k_cache", + "save_partial_states", +] diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py new file mode 100644 index 00000000..ac66751e --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py @@ -0,0 +1,594 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index +preparation. + +- quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into + the paged cache. +- dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged + cache for sparse/SWA prefill. +- compute_global_topk_indices_and_lens: map local topk indices to global KV + cache slots and count valid entries. +- combine_topk_swa_indices: concatenate topk compressed indices with SWA + window indices for sparse prefill. +""" + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_cutedsl + + +@triton.jit +def quantize_and_insert_k_kernel( + # Input tensors + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + # Output tensor + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + # Dimensions + num_tokens, + input_dim: tl.constexpr, # 512 + fp8_dim: tl.constexpr, # 448 + bf16_dim: tl.constexpr, # 64 + scale_dim: tl.constexpr, # 8 + quant_block: tl.constexpr, # 64 (quantization block size) + cache_block_size: tl.constexpr, # 64 (paged cache block size) + token_data_size: tl.constexpr, # 576 bytes per token data + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) +): + """ + Quantize K tensor and insert into paged K cache. + + K Cache block layout (block_size=64 tokens): + - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 + - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales + - [64*576 + 64*8, block_stride): Padding + + One program per token. + """ + pid = tl.program_id(0) + + if pid >= num_tokens: + return + + # Get slot mapping + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # Input pointer for this token + input_row_ptr = k_ptr + pid * input_dim + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks + # (e.g. >= 57K at block_stride ~37K). Matches gather path below. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + + # Token data pointer: token data is stored contiguously at start of block + # Each token's data is at offset pos_in_block * token_data_size + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size + + # Scale pointer: scales are stored after ALL token data in the block + # Scale for this token is at offset (64 * 576) + pos_in_block * 8 + token_scale_ptr = ( + cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim + ) + + # Token data layout: [0:448] fp8, [448:576] bf16 + token_fp8_ptr = token_data_ptr + token_bf16_ptr = token_data_ptr + fp8_dim + + # ========== Quantize and store FP8 portion (first 448 elements) ========== + # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent) + for qblock_idx in tl.static_range(n_quant_blocks): + qblock_start = qblock_idx * quant_block + + if qblock_start < fp8_dim: + offsets = qblock_start + tl.arange(0, quant_block) + mask = offsets < fp8_dim + + # Load bf16 input + x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0) + + # Compute absmax scale (same as CUDA kernel) + abs_x = tl.abs(x) + block_max = tl.max(abs_x, axis=0) + block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4) + + # UE8M0: Round scale UP to next power of 2 + # scale = 2^ceil(log2(block_max / fp8_max)) + raw_scale = block_max / fp8_max + log_scale = tl.log2(raw_scale) + exponent = tl.ceil(log_scale) # Round UP to next integer exponent + scale = tl.exp2(exponent) # scale = 2^exponent (power of 2) + + # Quantize to fp8: fp8_value = bf16_value / scale + x_scaled = x / scale + x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max) + + # Convert to fp8, then bitcast to uint8 for storage + x_fp8 = x_clamped.to(tl.float8e4nv) + x_uint8 = x_fp8.to(tl.uint8, bitcast=True) + + # Store as uint8 (1 byte each) + tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask) + + # UE8M0 scale encoding: stored_value = exponent + 127 (bias) + # During dequant: scale = 2^(stored_value - 127) + encoded_scale = exponent + 127.0 + encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0) + tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8)) + + # Padding scale at index 7 + tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8)) + + # ========== Store BF16 portion (last 64 elements, no quantization) ========== + bf16_input_offset = fp8_dim + + # Process bf16 in chunks of 16 + bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) + for i in tl.static_range(bf16_dim // 16): + chunk_offsets = i * 16 + tl.arange(0, 16) + bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets) + tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) + + +def quantize_and_insert_k_cache( + k: torch.Tensor, # [num_tokens, 512] bf16 + k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 + slot_mapping: torch.Tensor, # [num_tokens] int64 + block_size: int = 64, + is_ue8m0: bool = True, +): + """ + Quantize K tensor and insert into paged K cache. + + K Cache block layout (block_size=64 tokens): + - First 64 * 576 = 36864 bytes: Token data + - Each token: 448 bytes (fp8) + 128 bytes (bf16) + - Next 64 * 8 = 512 bytes: Scales + - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) + - Padded to multiple of 576 + """ + assert k.dim() == 2 and k.shape[1] == 512, ( + f"K must be [num_tokens, 512], got {k.shape}" + ) + assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}" + assert is_ue8m0, "Only support ue8m0 quantization." + + # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to + # padding. Always use slot_mapping.shape[0] as the token count. + num_tokens = slot_mapping.shape[0] + block_stride = k_cache.stride(0) # bytes per block + + TOKEN_FP8_DIM = 448 + TOKEN_BF16_DIM = 64 + TOKEN_SCALE_DIM = 8 + QUANT_BLOCK_SIZE = 64 + FP8_MAX = 448.0 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 + + grid = (num_tokens,) + + quantize_and_insert_k_kernel[grid]( + k, + slot_mapping, + k_cache, + num_tokens, + input_dim=512, + fp8_dim=TOKEN_FP8_DIM, + bf16_dim=TOKEN_BF16_DIM, + scale_dim=TOKEN_SCALE_DIM, + quant_block=QUANT_BLOCK_SIZE, + cache_block_size=block_size, + token_data_size=TOKEN_DATA_SIZE, + block_stride=block_stride, + fp8_max=FP8_MAX, + n_quant_blocks=8, + ) + + +@triton.jit +def _dequantize_and_gather_k_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + # Constants + max_blocks_per_seq: tl.constexpr, + fp8_dim: tl.constexpr, # 448 + bf16_dim: tl.constexpr, # 64 + scale_dim: tl.constexpr, # 8 + quant_block: tl.constexpr, # 64 (quantization block size) + cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size) + token_data_size: tl.constexpr, # 576 bytes per token data + block_stride: tl.constexpr, # total bytes per block (padded) int32 + output_dim: tl.constexpr, # 512 + fp8_max: tl.constexpr, + n_quant_blocks: tl.constexpr, # 7 real blocks +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + # Gather all tokens + gather_len = seq_len + start_pos = seq_len - gather_len + + for i in range(worker_id, gather_len, num_workers): + # Calculate the actual token index in the sequence + pos = start_pos + i + + # Calculate which block and position within block + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + # Get physical block index from block table + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31 with many + # KV-cache blocks (e.g. >= 57K at block_stride ~37K). + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + + # Token data pointer + token_data_ptr = cache_block_ptr + pos_in_block * token_data_size + + # Scale pointer: after all token data + token_scale_ptr = ( + cache_block_ptr + + cache_block_size * token_data_size + + pos_in_block * scale_dim + ) + + # Token data layout: [0:448] fp8, [448:576] bf16 + token_fp8_ptr = token_data_ptr + token_bf16_ptr = token_data_ptr + fp8_dim + + # Output pointer for this token (flattened) + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + + # ========== Dequantize FP8 portion using UE8M0 ========== + for qblock_idx in tl.static_range(n_quant_blocks): + qblock_start = qblock_idx * quant_block + + if qblock_start < fp8_dim: + offsets = qblock_start + tl.arange(0, quant_block) + mask = offsets < fp8_dim + + # Load quantized fp8 values (stored as uint8) + x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0) + + # Bitcast uint8 back to fp8 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + + # Convert fp8 to float32 for computation + x_float = x_fp8.to(tl.float32) + + # Load and decode UE8M0 scale + # UE8M0: scale = 2^(stored_value - 127) + encoded_scale = tl.load(token_scale_ptr + qblock_idx) + exponent = encoded_scale.to(tl.float32) - 127.0 + scale = tl.exp2(exponent) + + # Dequantize: bf16_value = fp8_value * scale + x_dequant = x_float * scale + + # Store as bf16 + tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask) + + # ========== Copy BF16 portion directly ========== + bf16_output_offset = fp8_dim # After 448 elements in output + + # Read bf16 from cache + bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) + + # Process in chunks of 16 + for j in tl.static_range(bf16_dim // 16): + chunk_offsets = j * 16 + tl.arange(0, 16) + bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets) + tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) + + +def dequantize_and_gather_k_cache_triton( + # [num_reqs, max_num_tokens, head_size] + out: torch.Tensor, + # [num_blocks, block_size, head_bytes] + k_cache: torch.Tensor, + # [num_reqs] + seq_lens: torch.Tensor, + # [num_reqs] + gather_lens: torch.Tensor | None, + # [num_reqs, max_blocks_per_seq] + block_table: torch.Tensor, + block_size: int, + offset: int, +) -> None: + TOKEN_FP8_DIM = 448 + TOKEN_BF16_DIM = 64 + TOKEN_SCALE_DIM = 8 + QUANT_BLOCK_SIZE = 64 + FP8_MAX = 448.0 + TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 + + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + fp8_dim=TOKEN_FP8_DIM, + bf16_dim=TOKEN_BF16_DIM, + scale_dim=TOKEN_SCALE_DIM, + quant_block=QUANT_BLOCK_SIZE, + cache_block_size=block_size, + token_data_size=TOKEN_DATA_SIZE, + block_stride=k_cache.stride(0), + output_dim=512, + fp8_max=FP8_MAX, + n_quant_blocks=7, + ) + + +def dequantize_and_gather_k_cache( + # [num_reqs, max_num_tokens, head_size] + out: torch.Tensor, + # [num_blocks, block_size, head_bytes] + k_cache: torch.Tensor, + # [num_reqs] + seq_lens: torch.Tensor, + # [num_reqs] + gather_lens: torch.Tensor | None, + # [num_reqs, max_blocks_per_seq] + block_table: torch.Tensor, + block_size: int, + offset: int, +) -> None: + if has_cutedsl(): + # lazily import, otherwise some tests fail due to CUDA driver init failure. + from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( + dequantize_and_gather_k_cache_cutedsl, + ) + + dequantize_and_gather_k_cache_cutedsl( + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + ) + return + + dequantize_and_gather_k_cache_triton( + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + ) + + +def compute_global_topk_indices_and_lens( + topk_indices: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + is_valid_token: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map local topk indices to global KV cache slots and count valid entries. + + Fuses three operations into a single kernel: + 1. Block-table lookup (local index → global slot id) + 2. Valid-entry counting (topk_lens per token) + 3. Masking padding tokens to length 0 + """ + num_tokens = topk_indices.shape[0] + global_topk_indices = torch.empty_like(topk_indices) + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( + global_topk_indices, + global_topk_indices.stride(0), + topk_lens, + topk_indices, + topk_indices.stride(0), + topk_indices.shape[-1], + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + is_valid_token, + TRITON_BLOCK_SIZE=1024, + ) + return global_topk_indices, topk_lens + + +@triton.jit +def _compute_global_topk_indices_and_lens_kernel( + global_topk_indices_ptr, + global_topk_indices_stride, + topk_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + topk, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + is_valid_token_ptr, + TRITON_BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0) + is_valid_token = tl.load(is_valid_token_ptr + token_idx) + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + count = tl.zeros((), dtype=tl.int32) + for i in range(0, topk, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + mask = offset < topk + + local_idx = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + other=-1, + ) + is_valid = local_idx >= 0 + + block_indices = local_idx // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask & is_valid, + ) + block_offsets = local_idx % block_size + + slot_ids = block_numbers * block_size + block_offsets + slot_ids = tl.where(is_valid, slot_ids, -1) + tl.store( + global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset, + slot_ids, + mask=mask, + ) + count += tl.sum(is_valid.to(tl.int32), axis=0) + + # Zero out length for padding tokens. + tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) + + +# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see +# flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is +# 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both. +# The extra slots stay as -1 sentinels and `combined_lens` caps the valid +# range via `topk_length`, so padding is a no-op at kernel level. +_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 + + +def combine_topk_swa_indices( + topk_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor, + window_size: int, + compress_ratio: int, + topk: int, + M: int, + N: int, +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = topk_indices.shape[0] + num_reqs = seq_lens.shape[0] + combined_topk = ( + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + // _SPARSE_PREFILL_TOPK_ALIGNMENT + * _SPARSE_PREFILL_TOPK_ALIGNMENT + ) + combined_indices = torch.full( + (num_tokens, combined_topk), + fill_value=-1, + dtype=torch.int32, + device=topk_indices.device, + ) + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + + NUM_WORKERS = 128 + _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)]( + combined_indices, + combined_indices.stride(0), + combined_lens, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=window_size, + PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), + ) + return combined_indices, combined_lens + + +@triton.jit +def _combine_topk_swa_indices_kernel( + combined_indices_ptr, + combined_indices_stride, + combined_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + query_start_loc_ptr, + seq_lens_ptr, + gather_lens_ptr, + M, + N, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + PADDED_TOP_K: tl.constexpr, +): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + # query_start_loc is a global tensor; rebase to chunk-local offsets + # by subtracting the chunk's starting value. + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + gather_len = tl.load(gather_lens_ptr + batch_idx) + start_pos = seq_len - query_len + # The SWA portion of the gathered buffer starts from position + # (seq_len - gather_len), not position 0. We need this offset + # to correctly index into the gathered buffer. + gather_start = seq_len - gather_len + + for token_idx in range(query_start + worker_id, query_end, num_workers): + # topk_len is fully determined by the query token's absolute position: + # both the C4A indexer and the C128A metadata builder emit + # min((pos + 1) // compress_ratio, topk_tokens) valid entries. + # Caller passes TOP_K=0 for SWA-only layers to zero this out. + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + + offset = tl.arange(0, PADDED_TOP_K) + mask = offset < topk_len + topk_indices = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + ) + tl.store( + combined_indices_ptr + token_idx * combined_indices_stride + offset, + topk_indices + M * batch_idx, + mask=mask, + ) + offset = tl.arange(0, WINDOW_SIZE) + # Index into gathered buffer: N + (position - gather_start) + # For positions [pos - swa_len + 1, pos], the buffer indices are: + # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + offset, + M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, + mask=offset < swa_len, + ) + + combined_len = topk_len + swa_len + tl.store(combined_lens_ptr + token_idx, combined_len) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py new file mode 100644 index 00000000..9a5e478e --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -0,0 +1,666 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Fused compressor + FP8/MXFP4 UE8M0 quantization + KV cache insert kernels. + +Three specialized kernels: + - _fused_kv_compress_norm_rope_insert_sparse_attn: + head=512, nope=448 FP8 + rope=64 bf16 + - _fused_kv_compress_norm_rope_insert_indexer_attn: + head=128, all FP8, 1 block/token + - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn: + head=128, MXFP4 (block=32), 4 ue8m0 bytes + +RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the +even/odd halves are consumed directly for MXFP4, no interleave needed). +FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for +per-block absmax entirely in registers. MXFP4 does the same tiling on the +even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles +and N_QUANT_BLOCKS ue8m0 bytes. +""" + +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton + +from .fused_indexer_q import _fp32x2_to_fp4x2 + + +def compress_norm_rope_store_triton( + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, +) -> None: + """Shared triton launcher for the fused compress+norm+RoPE+insert path. + + Picks one of the three kernels in this module based on ``head_dim`` and + ``use_fp4_cache``. Identical launch signature for all three. + """ + if head_dim == 512: + kernel = _fused_kv_compress_norm_rope_insert_sparse_attn + num_warps = 4 + elif use_fp4_cache: + kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn + num_warps = 1 + else: + kernel = _fused_kv_compress_norm_rope_insert_indexer_attn + num_warps = 1 + + kernel[(num_actual,)]( + # state cache + state_cache, + state_cache.stride(0), + state_cache.stride(1), + # metadata + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_table.stride(0), + block_size, + # RMSNorm + rms_norm_weight, + rms_norm_eps, + # RoPE + cos_sin_cache, + cos_sin_cache.stride(0), + # KV cache + kv_cache, + k_cache_metadata.slot_mapping, + kv_cache.shape[1], # paged KV cache block size (tokens per block) + # constexprs + HEAD_SIZE=head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), + STATE_WIDTH=state_width, + COMPRESS_RATIO=compress_ratio, + OVERLAP=overlap, + ROPE_HEAD_DIM=rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=quant_block, + TOKEN_STRIDE=token_stride, + SCALE_DIM=scale_dim, + KV_BLOCK_STRIDE=kv_cache.stride(0), + num_warps=num_warps, + **pdl_kwargs, + ) + + +# ============================================================================= +# DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16) +# ============================================================================= +@triton.jit +def _fused_kv_compress_norm_rope_insert_sparse_attn( + # ── state cache (compressor internal state) ── + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + # ── metadata ── + token_to_req_indices_ptr, + positions_ptr, + slot_mapping_ptr, + block_table_ptr, + block_table_stride, + block_size, + # ── RMSNorm ── + rms_norm_weight_ptr, + rms_norm_eps, + # ── RoPE ── + cos_sin_cache_ptr, + cos_sin_stride, + # ── KV cache output ── + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + # ── constexprs ── + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + OVERLAP: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, # 448.0 + QUANT_BLOCK: tl.constexpr, # 64 for DeepseekV4 + TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4 + SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) + KV_BLOCK_STRIDE: tl.constexpr, +): + """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). + + One program per token; early-exits for non-boundary positions. + + Cache block layout (``block_size`` tokens): + [0, bs*576): token data (448 fp8 + 128 bf16 each) + [bs*576, +bs*8): uint8 UE8M0 scales (7 real + 1 pad each) + """ + token_idx = tl.program_id(0) + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + block_numbers_i64 = block_numbers.to(tl.int64) + + # Precomputed row base shared by score and kv loads + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + + combined_mask = mask_pos[:, None] & mask[None, :] + + # ── Softmax + weighted sum ─────────────────────────────────────── + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 + + # ── RMSNorm (fp32 throughout) ────────────────────────────────────── + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + # ── KV cache pointers ──────────────────────────────────────────── + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM # 448 + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 # 32 + + # FP8 UE8M0 quant: cast fp32 → bf16 → fp32 before quant to match reference. + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK # 7 + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + quant_input = normed.to(tl.bfloat16).to(tl.float32) + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) + abs_2d = tl.abs(quant_2d) + block_absmax = tl.max(abs_2d, axis=1) # [N_QUANT_BLOCKS] fp32 + block_absmax = tl.maximum(block_absmax, 1e-4) + + raw_scales = block_absmax * INV_FP8_MAX + exponents = tl.ceil(tl.log2(raw_scales)) + inv_scales = tl.exp2(-exponents) + inv_scales_col = tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) + x_scaled = quant_2d * inv_scales_col + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) + x_fp8 = x_clamped.to(tl.float8e4nv) + x_uint8 = x_fp8.to(tl.uint8, bitcast=True) + x_uint8_flat = tl.reshape(x_uint8, (TRITON_BLOCK_SIZE,)) + + nope_mask = block < NOPE_HEAD_DIM + tl.store(fp8_ptr + block, x_uint8_flat, mask=nope_mask) + + scale_idx = tl.arange(0, N_QUANT_BLOCKS) + encoded = exponents + 127.0 + encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0) + tl.store( + scale_ptr + scale_idx, + encoded.to(tl.uint8), + mask=scale_idx < N_NOPE_BLOCKS, + ) + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) + + # Register-based GPT-J RoPE in fp32. + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + + pair_2d = tl.reshape(normed, (NUM_PAIRS, 2)) + even, odd = tl.split(pair_2d) # each [NUM_PAIRS] fp32 + + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32 + + # Store rotated rope portion as bf16 into the cache's bf16 area. + bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) + rope_local = block - NOPE_HEAD_DIM + is_rope = (block >= NOPE_HEAD_DIM) & mask + tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) + + +# ============================================================================= +# Indexer path (head=128, all FP8, single quant block) +# ============================================================================= +@triton.jit +def _fused_kv_compress_norm_rope_insert_indexer_attn( + # ── state cache (compressor internal state) ── + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + # ── metadata ── + token_to_req_indices_ptr, + positions_ptr, + slot_mapping_ptr, + block_table_ptr, + block_table_stride, + block_size, + # ── RMSNorm ── + rms_norm_weight_ptr, + rms_norm_eps, + # ── RoPE ── + cos_sin_cache_ptr, + cos_sin_stride, + # ── KV cache output ── + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + # ── constexprs ── + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + OVERLAP: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, # 448.0 + QUANT_BLOCK: tl.constexpr, # 128 for indexer + TOKEN_STRIDE: tl.constexpr, # 128 for indexer + SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) + KV_BLOCK_STRIDE: tl.constexpr, +): + """Fused compress → RMSNorm → RoPE → FP8 quant → store. + + One program per token; early-exits for non-boundary positions. + + Cache block layout: + [0, bs*128): FP8 data (128 bytes/token) + [bs*128, +bs*4): float32 scales (4 bytes/token) + + For head_dim=128 we have exactly one quant block, so we skip the + [N_QUANT_BLOCKS, QUANT_BLOCK] reshape entirely and use a flat + ``tl.max`` reduction. + """ + token_idx = tl.program_id(0) + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + block_numbers_i64 = block_numbers.to(tl.int64) + + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + + combined_mask = mask_pos[:, None] & mask[None, :] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 + + # ── RMSNorm (fp32 throughout) ────────────────────────────────────── + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + # ── KV cache pointers ──────────────────────────────────────────── + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + + # ── Register-based GPT-J forward RoPE in fp32 ───────────────────── + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + + normed_2d = tl.reshape(normed, (NUM_PAIRS, 2)) + even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32 + + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + result = tl.interleave(new_even, new_odd) # fp32 + + # ── FP8 UE8M0 quant: single block, flat reduction ──────────────── + tl.static_assert( + TRITON_BLOCK_SIZE == QUANT_BLOCK, + "Indexer expects one quant block (QUANT_BLOCK == TRITON_BLOCK_SIZE)", + ) + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + result_bf16 = result.to(tl.bfloat16).to(tl.float32) + absmax = tl.max(tl.abs(result_bf16), axis=0) # scalar + absmax = tl.maximum(absmax, 1e-4) + raw_scale = absmax * INV_FP8_MAX + exponent = tl.ceil(tl.log2(raw_scale)) + inv_scale = tl.exp2(-exponent) + + x_scaled = result_bf16 * inv_scale + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) + x_fp8 = x_clamped.to(tl.float8e4nv) + x_uint8 = x_fp8.to(tl.uint8, bitcast=True) + + tl.store(fp8_ptr + block, x_uint8, mask=mask) + + # Single float32 scale + scale_val = tl.exp2(exponent) + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) + + +# ============================================================================= +# Indexer path (head=128, MXFP4: 2 nibbles/byte + ue8m0 per 32-elem block) +# ============================================================================= +@triton.jit +def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( + # ── state cache (compressor internal state) ── + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + # ── metadata ── + token_to_req_indices_ptr, + positions_ptr, + slot_mapping_ptr, + block_table_ptr, + block_table_stride, + block_size, + # ── RMSNorm ── + rms_norm_weight_ptr, + rms_norm_eps, + # ── RoPE ── + cos_sin_cache_ptr, + cos_sin_stride, + # ── KV cache output ── + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + # ── constexprs ── + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + OVERLAP: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, # unused for MXFP4 (kept for signature parity) + QUANT_BLOCK: tl.constexpr, # 32 for MXFP4 + TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token + SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token + KV_BLOCK_STRIDE: tl.constexpr, +): + """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. + + One program per token; early-exits for non-boundary positions. + + Cache block layout (``block_size`` tokens per cache block): + [0, bs*TOKEN_STRIDE): packed MXFP4 nibbles (2 values/byte) + [bs*TOKEN_STRIDE, +bs*SCALE_DIM): ue8m0 scale bytes (one per 32-elem block) + + MXFP4 format: + - E2M1 4-bit values packed two per byte (low nibble first, then high). + - Per-32-element block scale = 2^ceil(log2(amax / 6.0)), stored ue8m0 + (byte = exponent + 127). + - Max representable magnitude = 6.0. + """ + token_idx = tl.program_id(0) + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + # ── Gather state cache entries ──────────────────────────────────── + start = position - (1 + OVERLAP) * COMPRESS_RATIO + 1 + tokens = tl.arange(0, (1 + OVERLAP) * COMPRESS_RATIO) + pos = start + tokens + mask_pos = pos >= 0 + + block_indices = pos // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=mask_pos, + other=0, + ) + block_offsets = pos % block_size + head_offset = (tokens >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + block_numbers_i64 = block_numbers.to(tl.int64) + + row_base = ( + state_cache_ptr + + block_numbers_i64 * state_cache_stride0 + + block_offsets * state_cache_stride1 + + head_offset + ) + + combined_mask = mask_pos[:, None] & mask[None, :] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + block[None, :], + mask=combined_mask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + + kv = tl.load( + row_base[:, None] + block[None, :], + mask=combined_mask, + other=0.0, + ) + + compressed_kv = tl.sum(kv * score, axis=0) # [TRITON_BLOCK_SIZE] fp32 + + # ── RMSNorm (fp32 throughout) ────────────────────────────────────── + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + # ── KV cache pointers (segregated: values first, then scales) ──── + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + + # ── Register-based GPT-J forward RoPE in fp32 ───────────────────── + # We keep the even/odd halves (no tl.interleave afterwards) because the + # MXFP4 per-block absmax / pack naturally operates on (even, odd) pairs. + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + + normed_2d = tl.reshape(normed, (NUM_PAIRS, 2)) + even, odd = tl.split(normed_2d) # each [NUM_PAIRS] fp32 + + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + + # bf16 roundtrip for parity with reference / Q-side kernel numerics. + new_even = new_even.to(tl.bfloat16).to(tl.float32) + new_odd = new_odd.to(tl.bfloat16).to(tl.float32) + + # ── MXFP4 quant: tile even/odd halves into (N_BLOCKS, HALF_BLOCK) ── + # Each MXFP4 block of QUANT_BLOCK elements = HALF_BLOCK consecutive pairs, + # so (N_BLOCKS, HALF_BLOCK) rows of even/odd each land exactly one block. + N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // QUANT_BLOCK + HALF_BLOCK: tl.constexpr = QUANT_BLOCK // 2 + tl.static_assert(TRITON_BLOCK_SIZE == HEAD_SIZE) + tl.static_assert(HEAD_SIZE % QUANT_BLOCK == 0) + tl.static_assert(TOKEN_STRIDE == HEAD_SIZE // 2) + tl.static_assert(SCALE_DIM == N_QUANT_BLOCKS) + + even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK)) + odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK)) + + amax = tl.maximum( + tl.max(tl.abs(even_2d), axis=1), + tl.max(tl.abs(odd_2d), axis=1), + ) + amax = tl.maximum(amax, 6.0 * (2**-126)) + + # ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte. + log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0))) + log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0) + inv_scale = tl.exp2(-log2_ratio) + ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS] + + inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1)) + packed = _fp32x2_to_fp4x2( + even_2d * inv_scale_col, odd_2d * inv_scale_col + ) # (N_BLOCKS, HALF_BLOCK) uint8 + packed_flat = tl.reshape(packed, (TOKEN_STRIDE,)) + + tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat) + tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py new file mode 100644 index 00000000..d5aaf10f --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_cutedsl + +# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. +MXFP4_BLOCK_SIZE = 32 + + +@triton.jit +def _get_cos_sin( + cos_sin_cache_ptr, + cos_sin_cache_stride, + pos, + HALF_ROT_DIM: tl.constexpr, +): + block = tl.arange(0, HALF_ROT_DIM) + cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block) + cos = cos.to(tl.float32) + sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM) + sin = sin.to(tl.float32) + return cos, sin + + +@triton.jit +def _fp32x2_to_fp4x2(x_lo, x_hi): + # NOTE: $1 is high nibble, $2 is low nibble + return tl.inline_asm_elementwise( + """ + { + .reg .b8 tmp; + cvt.rn.satfinite.e2m1x2.f32 tmp, $1, $2; + cvt.u32.u8 $0, tmp; + } + """, + constraints="=r,f,f", + args=[x_hi, x_lo], + dtype=tl.uint32, + is_pure=True, + pack=1, + ).to(tl.uint8) + + +@triton.jit +def _quantize_mxfp4_pair(x_lo, x_hi): + """Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two + interleaved halves (x_lo = values at even positions in the block, + x_hi = values at odd positions). Returns: + - packed : uint8[BLOCK/2] (low nibble = quant(x_lo), high = quant(x_hi)) + - ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127)) + """ + amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi))) + # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163 + amax = tl.maximum(amax, 6.0 * (2**-126)) + # ue8m0 block scale: 2^ceil(log2(amax/6.0)). + log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0))) + log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0) + scale = tl.math.exp2(log2_ratio) + ue8m0 = (log2_ratio + 127.0).to(tl.uint8) + + inv_scale = 1.0 / scale + packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale) + return packed, ue8m0 + + +@triton.jit +def _fused_indexer_q_rope_quant_kernel( + pos_ptr, + # Index Q RoPE + index_q_ptr, + index_q_stride0, + index_q_stride1, + index_q_cos_sin_ptr, + index_q_cos_sin_stride, + INDEX_Q_HALF_ROT_DIM: tl.constexpr, + # Index Q Quantize + index_q_fp8_ptr, + index_q_fp8_stride0, + index_q_fp8_stride1, + INDEX_Q_HEAD_DIM: tl.constexpr, + # Index weights + index_weights_ptr, + index_weights_stride, + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out_ptr, + index_weights_out_stride, +): + # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding + # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the + # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed + # through unchanged. + INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM + INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM + tl.static_assert(INDEX_Q_NOPE_DIM >= 0) + + tok_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + pos = tl.load(pos_ptr + tok_idx) + cos, sin = _get_cos_sin( + index_q_cos_sin_ptr, + index_q_cos_sin_stride, + pos, + INDEX_Q_HALF_ROT_DIM, + ) + half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM) + base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 + + # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM): + # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1] + rot_base = base_ptr + INDEX_Q_NOPE_DIM + x_even = tl.load(rot_base + half_offset * 2).to(tl.float32) + x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32) + r_even = x_even * cos - x_odd * sin + r_odd = x_odd * cos + x_even * sin + + # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax. + # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py). + r_even = r_even.to(tl.bfloat16).to(tl.float32) + r_odd = r_odd.to(tl.bfloat16).to(tl.float32) + + amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd))) + if INDEX_Q_NOPE_DIM > 0: + nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM) + x_nope = tl.load(base_ptr + nope_offset).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(x_nope))) + index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0) + index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale))) + + # Store quantized values to index_q_fp8 + fp8_base_ptr = ( + index_q_fp8_ptr + tok_idx * index_q_fp8_stride0 + head_idx * index_q_fp8_stride1 + ) + if INDEX_Q_NOPE_DIM > 0: + tl.store( + fp8_base_ptr + nope_offset, + tl.div_rn(x_nope, index_q_scale).to(tl.float8e4nv), + ) + fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM + tl.store( + fp8_rot_base + half_offset * 2, + tl.div_rn(r_even, index_q_scale).to(tl.float8e4nv), + ) + tl.store( + fp8_rot_base + half_offset * 2 + 1, + tl.div_rn(r_odd, index_q_scale).to(tl.float8e4nv), + ) + + # FP8 weight-fold contract: + # index_weights_out = index_weights * q_scale * softmax_scale * head_scale + # The per-token-per-head q_scale (fp32) IS folded into the output weights + # here because FP8 Q is stored WITHOUT a companion scale tensor — the + # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use `weights` to + # apply per-token Q scale inline. See the MXFP4 kernel below for the + # contrasting convention (scales live with the Q values, weights are NOT + # q-scaled). + index_weights = tl.load( + index_weights_ptr + tok_idx * index_weights_stride + head_idx + ) + index_weights = index_weights.to(tl.float32) + index_weights *= index_q_scale + index_weights *= index_weights_softmax_scale + index_weights *= index_weights_head_scale + tl.store( + index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, + index_weights, + ) + + +@triton.jit +def _fused_indexer_q_rope_mxfp4_kernel( + pos_ptr, + # Index Q RoPE input (fp/bf16) + index_q_ptr, + index_q_stride0, + index_q_stride1, + index_q_cos_sin_ptr, + index_q_cos_sin_stride, + INDEX_Q_HALF_ROT_DIM: tl.constexpr, + # MXFP4 Q outputs + index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2) + index_q_mxfp4_stride0, + index_q_mxfp4_stride1, + index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK) + index_q_scale_stride0, + index_q_scale_stride1, + INDEX_Q_HEAD_DIM: tl.constexpr, + MXFP4_BLOCK: tl.constexpr, + # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay + # with the Q values in the output scale tensor). + index_weights_ptr, + index_weights_stride, + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out_ptr, + index_weights_out_stride, +): + INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM + INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM + NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK + NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK + HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2 + tl.static_assert(INDEX_Q_NOPE_DIM >= 0) + tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0) + tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0) + tl.static_assert(MXFP4_BLOCK % 2 == 0) + + tok_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + pos = tl.load(pos_ptr + tok_idx) + + q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1 + out_base = ( + index_q_mxfp4_ptr + + tok_idx * index_q_mxfp4_stride0 + + head_idx * index_q_mxfp4_stride1 + ) + scale_base = ( + index_q_scale_ptr + + tok_idx * index_q_scale_stride0 + + head_idx * index_q_scale_stride1 + ) + + half_off = tl.arange(0, HALF_BLOCK) + + # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ---- + for b in tl.static_range(NUM_NOPE_BLOCKS): + base = b * MXFP4_BLOCK + x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32) + x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32) + packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi) + tl.store(out_base + base // 2 + half_off, packed) + tl.store(scale_base + b, ue8m0) + + # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs, + # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ---- + rot_q_base = q_base + INDEX_Q_NOPE_DIM + for b in tl.static_range(NUM_ROPE_BLOCKS): + pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM) + cos_b = tl.load( + index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off + ).to(tl.float32) + sin_b = tl.load( + index_q_cos_sin_ptr + + pos * index_q_cos_sin_stride + + pair_off + + INDEX_Q_HALF_ROT_DIM + ).to(tl.float32) + x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32) + x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32) + r_even = x_even * cos_b - x_odd * sin_b + r_odd = x_odd * cos_b + x_even * sin_b + # bf16 roundtrip for parity with the FP8 kernel / reference numerics. + r_even = r_even.to(tl.bfloat16).to(tl.float32) + r_odd = r_odd.to(tl.bfloat16).to(tl.float32) + packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd) + rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2 + tl.store(out_base + rope_byte_off + half_off, packed) + tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0) + + # MXFP4 weight-fold contract: + # index_weights_out = index_weights * softmax_scale * head_scale + # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above). + # MXFP4 Q emits a separate ue8m0 scale tensor of shape + # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each + # per-block scale is applied by the downstream MXFP4 logits kernel when + # dequantizing Q — there is no per-token scalar to fold into `weights`. + index_weights = tl.load( + index_weights_ptr + tok_idx * index_weights_stride + head_idx + ).to(tl.float32) + index_weights *= index_weights_softmax_scale + index_weights *= index_weights_head_scale + tl.store( + index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx, + index_weights, + ) + + +def fused_indexer_q_rope_quant( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + # Index weights + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + use_fp4: bool = False, +) -> tuple[ + torch.Tensor | tuple[torch.Tensor, torch.Tensor], + torch.Tensor, +]: + """Fused RoPE + quantize Q for the sparse indexer. + + Weight-fold semantics (important — the two paths differ): + + FP8 path (use_fp4=False, default): + q_fp8 : (T, H, HEAD_DIM) float8_e4m3fn, per-token-per-head + scalar scale (NOT stored — folded into weights below) + weights_out = weights * q_scale * softmax_scale * head_scale + Rationale: a single per-token q_scale is a scalar the downstream FP8 + logits kernel would otherwise multiply in. Folding it into `weights` + avoids emitting a separate tensor and is free for the logits kernel. + + MXFP4 path (use_fp4=True): + q_packed : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte) + q_scale : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes + weights_out = weights * softmax_scale * head_scale + Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with + the Q values — they cannot be folded into a per-token weight + scalar, so `weights` carries only the softmax and head scales. + + Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or + a (values, scales) tuple (MXFP4). This matches the union type accepted + by `SparseAttnIndexer.forward_*`. + """ + assert positions.ndim == 1 + assert index_q.ndim == 3 + assert index_q_cos_sin_cache.ndim == 2 + + num_tokens = positions.shape[0] + num_index_q_heads = index_q.shape[1] + index_q_head_dim = index_q.shape[2] + + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + + if use_fp4: + assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( + f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block " + f"size {MXFP4_BLOCK_SIZE}" + ) + num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE + index_q_packed = torch.empty( + (num_tokens, num_index_q_heads, index_q_head_dim // 2), + dtype=torch.uint8, + device=index_q.device, + ) + index_q_scale = torch.empty( + (num_tokens, num_index_q_heads, num_scale_blocks), + dtype=torch.uint8, + device=index_q.device, + ) + if has_cutedsl(): + # lazily import, otherwise some tests fail due to CUDA driver init failure. + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( + fused_indexer_q_rope_quant_mxfp4_cutedsl, + ) + + fused_indexer_q_rope_quant_mxfp4_cutedsl( + positions, + index_q, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_packed, + index_q_scale, + index_weights_out, + ) + else: + _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)]( + positions, + index_q, + index_q.stride(0), + index_q.stride(1), + index_q_cos_sin_cache, + index_q_cos_sin_cache.stride(0), + index_q_cos_sin_cache.shape[-1] // 2, + index_q_packed, + index_q_packed.stride(0), + index_q_packed.stride(1), + index_q_scale, + index_q_scale.stride(0), + index_q_scale.stride(1), + index_q_head_dim, + MXFP4_BLOCK_SIZE, + index_weights, + index_weights.stride(0), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + index_weights_out.stride(0), + num_warps=1, # TODO: Tune this + ) + + # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0 + # bytes per (token, head) reinterpreted as one int32, then squeezed + # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank + # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to + # 3-D (batch, next_n, num_heads)). + return ( + index_q_packed, + index_q_scale.view(torch.int32).squeeze(-1), + ), index_weights_out + + index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) + if has_cutedsl(): + # lazily import, otherwise some tests fail due to CUDA driver init failure. + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( + fused_indexer_q_rope_quant_fp8_cutedsl, + ) + + fused_indexer_q_rope_quant_fp8_cutedsl( + positions, + index_q, + index_q_cos_sin_cache, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_q_fp8, + index_weights_out, + ) + else: + _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)]( + positions, + index_q, + index_q.stride(0), + index_q.stride(1), + index_q_cos_sin_cache, + index_q_cos_sin_cache.stride(0), + index_q_cos_sin_cache.shape[-1] // 2, + index_q_fp8, + index_q_fp8.stride(0), + index_q_fp8.stride(1), + index_q_head_dim, + index_weights, + index_weights.stride(0), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + index_weights_out.stride(0), + num_warps=1, # TODO: Tune this + ) + return index_q_fp8, index_weights_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py new file mode 100644 index 00000000..97fc0962 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Fused inverse RoPE + block-scaled FP8 quantization kernel for DeepseekV4 attention. + +Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90, +INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout. +""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _fused_inv_rope_fp8_quant_per_head( + o_ptr, + positions_ptr, + cos_sin_cache_ptr, + fp8_ptr, + scale_ptr, + num_tokens, + heads_per_group: tl.constexpr, + o_stride_token, + o_stride_head, + cache_stride_pos, + fp8_stride_group, + fp8_stride_token, + scale_stride_group, + scale_stride_k, + fp8_max: tl.constexpr, + eps: tl.constexpr, + QUANT_GROUP_SIZE: tl.constexpr, + CHUNKS_PER_HEAD: tl.constexpr, + ROPE_START: tl.constexpr, + HALF_ROPE: tl.constexpr, + TMA_ALIGNED_SCALES: tl.constexpr, +): + # int64: stride multiply overflows int32 past num_tokens=32768 (IMA). + pid_token = tl.program_id(0).to(tl.int64) + pid_gh = tl.program_id(1).to(tl.int64) + + g = pid_gh // heads_per_group + head_in_group = pid_gh % heads_per_group + global_head = pid_gh + qb_start = head_in_group * CHUNKS_PER_HEAD + + # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant. + if pid_token >= num_tokens: + if TMA_ALIGNED_SCALES: + scale_addr = ( + scale_ptr + + g * scale_stride_group + + pid_token + + head_in_group * scale_stride_k + ) + tl.store(scale_addr, tl.zeros((), dtype=tl.int32)) + else: + block_offsets = tl.arange(0, CHUNKS_PER_HEAD) + qb_indices = qb_start + block_offsets + scale_addrs = ( + scale_ptr + + g * scale_stride_group + + pid_token + + qb_indices * scale_stride_k + ) + tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32)) + return + + input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head + + HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE + offsets = tl.arange(0, HEAD_DIM) + x = tl.load(input_base + offsets).to(tl.float32) + + rope_abs_start: tl.constexpr = (CHUNKS_PER_HEAD - 1) * QUANT_GROUP_SIZE + ROPE_START + pos = tl.load(positions_ptr + pid_token) + cache_base = cos_sin_cache_ptr + pos * cache_stride_pos + is_rope = offsets >= rope_abs_start + rope_local = offsets - rope_abs_start + + x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to( + tl.float32 + ) + cs_idx = tl.maximum(rope_local >> 1, 0) + cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0) + x_add = x * cos_v + x_partner * sin_v + x_sub = x * cos_v - x_partner * sin_v + is_even = (rope_local & 1) == 0 + rotated = tl.where(is_even, x_add, x_sub) + x = tl.where(is_rope, rotated, x) + + x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE)) + block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps) + scale_raw = block_absmax * (1.0 / fp8_max) + scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) + + scales_exp = tl.reshape( + tl.broadcast_to( + tl.reshape(scales, (CHUNKS_PER_HEAD, 1)), + (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE), + ), + (HEAD_DIM,), + ) + x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv) + + fp8_base = ( + fp8_ptr + + g * fp8_stride_group + + pid_token * fp8_stride_token + + qb_start * QUANT_GROUP_SIZE + ) + tl.store(fp8_base + offsets, x_quant) + + block_offsets = tl.arange(0, CHUNKS_PER_HEAD) + qb_indices = qb_start + block_offsets + if TMA_ALIGNED_SCALES: + scale_bits = scales.to(tl.int32, bitcast=True) + ue8m0_bytes = (scale_bits >> 23) & 0xFF + packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8)) + scale_addr = ( + scale_ptr + + g * scale_stride_group + + pid_token + + head_in_group * scale_stride_k + ) + tl.store(scale_addr, packed_val) + else: + scale_addrs = ( + scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k + ) + tl.store(scale_addrs, scales) + + +def fused_inv_rope_fp8_quant( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int = 448, + rope_dim: int = 64, + quant_group_size: int = 128, + tma_aligned_scales: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused inverse RoPE + block-scaled FP8 quantization. + + Args: + o: Attention output [num_tokens, num_heads, head_dim] bf16. + positions: Token positions [num_tokens] int64. + cos_sin_cache: Precomputed [max_pos, rope_dim] with cos||sin. + n_groups: Number of output groups. + heads_per_group: Heads per group. + nope_dim: Non-RoPE dimensions per head (default 448). + rope_dim: RoPE dimensions per head (default 64). + quant_group_size: FP8 quantization block size (default 128). + tma_aligned_scales: Output INT32 packed UE8M0 for SM100 (True) + or FP32 for SM90 (False). + + Returns: + o_fp8: [T, G, D] float8_e4m3fn, strides (D, T*D, 1). + o_scale: Pre-transformed scale tensor for fp8_einsum. + """ + from vllm.utils.deep_gemm import get_tma_aligned_size + + num_tokens, num_heads, head_dim = o.shape + assert num_heads == n_groups * heads_per_group + assert head_dim == nope_dim + rope_dim + assert head_dim % quant_group_size == 0 + assert nope_dim % quant_group_size == (quant_group_size - rope_dim) + assert rope_dim % 2 == 0 + assert cos_sin_cache.shape[-1] == rope_dim + assert cos_sin_cache.dtype == torch.float32 + + d = heads_per_group * head_dim + num_scale_blocks = d // quant_group_size + chunks_per_head = head_dim // quant_group_size + + fp8_dtype = torch.float8_e4m3fn + fp8_max = torch.finfo(fp8_dtype).max + + tma_aligned_T = get_tma_aligned_size(num_tokens, 4) + if tma_aligned_scales: + packed_sf_k = (num_scale_blocks + 3) // 4 + scale_inner = packed_sf_k + else: + scale_inner = num_scale_blocks + + # Run kernel through a custom op so inductor sees an opaque boundary. + # It's a pytorch bug, see https://github.com/vllm-project/vllm/issues/41106 + fp8_buf, scale_buf = torch.ops.vllm.fused_inv_rope_fp8_quant_kernel( + o, + positions, + cos_sin_cache, + heads_per_group, + quant_group_size, + chunks_per_head, + nope_dim % quant_group_size, + rope_dim // 2, + tma_aligned_scales, + fp8_max, + tma_aligned_T, + num_tokens, + n_groups, + d, + scale_inner, + ) + return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1) + + +def _fused_inv_rope_fp8_quant_kernel_impl( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + heads_per_group: int, + quant_group_size: int, + chunks_per_head: int, + rope_start: int, + half_rope: int, + tma_aligned_scales: bool, + fp8_max: float, + tma_aligned_T: int, + num_tokens: int, + n_groups: int, + d: int, + scale_inner: int, +) -> tuple[torch.Tensor, torch.Tensor]: + fp8_buf = torch.empty( + (n_groups, num_tokens, d), + dtype=torch.float8_e4m3fn, + device=o.device, + ) + scale_dtype = torch.int32 if tma_aligned_scales else torch.float32 + scale_buf = torch.empty( + n_groups * scale_inner * tma_aligned_T, + dtype=scale_dtype, + device=o.device, + ).as_strided( + (n_groups, num_tokens, scale_inner), + (scale_inner * tma_aligned_T, 1, tma_aligned_T), + ) + grid = (tma_aligned_T, n_groups * heads_per_group) + pdl_kwargs = ( + {} + if current_platform.is_rocm() or current_platform.is_xpu() + else {"launch_pdl": False} + ) + _fused_inv_rope_fp8_quant_per_head[grid]( + o, + positions, + cos_sin_cache, + fp8_buf, + scale_buf, + 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), + fp8_stride_group=fp8_buf.stride(0), + fp8_stride_token=fp8_buf.stride(1), + scale_stride_group=scale_buf.stride(0), + scale_stride_k=scale_buf.stride(2), + fp8_max=fp8_max, + eps=1e-10, + QUANT_GROUP_SIZE=quant_group_size, + CHUNKS_PER_HEAD=chunks_per_head, + ROPE_START=rope_start, + HALF_ROPE=half_rope, + TMA_ALIGNED_SCALES=tma_aligned_scales, + num_stages=1, + **pdl_kwargs, + num_warps=1, + ) + return fp8_buf, scale_buf + + +def _fused_inv_rope_fp8_quant_kernel_fake( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + heads_per_group: int, + quant_group_size: int, + chunks_per_head: int, + rope_start: int, + half_rope: int, + tma_aligned_scales: bool, + fp8_max: float, + tma_aligned_T: int, + num_tokens: int, + n_groups: int, + d: int, + scale_inner: int, +) -> tuple[torch.Tensor, torch.Tensor]: + fp8_buf = torch.empty( + (n_groups, num_tokens, d), + dtype=torch.float8_e4m3fn, + device=o.device, + ) + scale_dtype = torch.int32 if tma_aligned_scales else torch.float32 + scale_buf = torch.empty( + n_groups * scale_inner * tma_aligned_T, + dtype=scale_dtype, + device=o.device, + ).as_strided( + (n_groups, num_tokens, scale_inner), + (scale_inner * tma_aligned_T, 1, tma_aligned_T), + ) + return fp8_buf, scale_buf + + +direct_register_custom_op( + op_name="fused_inv_rope_fp8_quant_kernel", + op_func=_fused_inv_rope_fp8_quant_kernel_impl, + fake_impl=_fused_inv_rope_fp8_quant_kernel_fake, +) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py new file mode 100644 index 00000000..874b4f60 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MTP-input RMSNorm: enorm (with mask-zero at position 0) + hnorm. + +Replaces the eager sequence at the top of the MTP draft forward: + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = previous_hidden_states.view(-1, hc_mult, H) + previous_hidden_states = self.hnorm(previous_hidden_states) + +which lowers to ~6 small kernels (CompareEq, where, Fill, enorm rms_norm, +hnorm rms_norm, plus aten elementwise helpers) on the breakable-cudagraph +path. Math is preserved: positions==0 → masked row → zero RMS output +regardless of weight. + +A single grid (T, hc_mult+1) drives both norms: task 0 is enorm on +inputs_embeds[token, :], task k+1 is hnorm on previous_hidden_states[token, k, :]. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rmsnorm_row( + x, + w_ptr, + out_row_ptr, + block, + mask, + eps, + HIDDEN: tl.constexpr, +): + x = x.to(tl.float32) + variance = tl.sum(x * x, axis=0) / HIDDEN + rrms = tl.rsqrt(variance + eps) + w = tl.load(w_ptr + block, mask=mask, other=0.0).to(tl.float32) + y = x * rrms * w + tl.store(out_row_ptr + block, y.to(out_row_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _fused_mtp_input_rmsnorm_kernel( + inputs_embeds_ptr, + positions_ptr, + prev_hidden_ptr, + enorm_weight_ptr, + hnorm_weight_ptr, + enorm_out_ptr, + hnorm_out_ptr, + eps, + HIDDEN: tl.constexpr, + HC_MULT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # int64 token index so per-token offsets don't overflow int32 at + # large num_tokens (matches the convention in fused_q_kv_rmsnorm). + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + block = tl.arange(0, BLOCK_SIZE) + mask = block < HIDDEN + + if pid_task == 0: + # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0. + # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0 + # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm. + pos = tl.load(positions_ptr + token_idx) + keep = pos != 0 + x = tl.load( + inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0 + ) + x = tl.where(keep, x, 0.0) + _rmsnorm_row( + x, + enorm_weight_ptr, + enorm_out_ptr + token_idx * HIDDEN, + block, + mask, + eps, + HIDDEN, + ) + else: + # hnorm path: load prev_hidden[token, slot, :]. + slot = pid_task - 1 + row_offset = (token_idx * HC_MULT + slot) * HIDDEN + x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0) + _rmsnorm_row( + x, + hnorm_weight_ptr, + hnorm_out_ptr + row_offset, + block, + mask, + eps, + HIDDEN, + ) + + +@triton.jit +def _mtp_shared_head_rmsnorm_kernel( + x_ptr, + weight_ptr, + out_ptr, + eps, + HIDDEN: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + block = tl.arange(0, BLOCK_SIZE) + mask = block < HIDDEN + x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0) + _rmsnorm_row( + x, + weight_ptr, + out_ptr + token_idx * HIDDEN, + block, + mask, + eps, + HIDDEN, + ) + + +def mtp_shared_head_rmsnorm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """RMSNorm for MTP's SharedHead.norm, on (T, H) bf16 input. + + Uses the same ``_rmsnorm_row`` body as ``fused_mtp_input_rmsnorm`` so the + MTP draft path runs one consistent RMSNorm implementation end to end. + """ + assert hidden_states.ndim == 2 + assert hidden_states.is_contiguous() + assert weight.is_contiguous() + num_tokens, hidden = hidden_states.shape + out = torch.empty_like(hidden_states) + if num_tokens == 0: + return out + block_size = triton.next_power_of_2(hidden) + _mtp_shared_head_rmsnorm_kernel[(num_tokens,)]( + hidden_states, + weight, + out, + eps, + HIDDEN=hidden, + BLOCK_SIZE=block_size, + ) + return out + + +def fused_mtp_input_rmsnorm( + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + hc_mult: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns (enorm_out, hnorm_out). + + enorm_out has the same shape as inputs_embeds (2D, [T, H]). + hnorm_out has the same shape as previous_hidden_states (3D, [T, hc_mult, H]). + previous_hidden_states must already be reshaped to 3D. + """ + assert inputs_embeds.ndim == 2 + assert previous_hidden_states.ndim == 3 + assert previous_hidden_states.shape[1] == hc_mult + assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], ( + "token dim mismatch" + ) + assert ( + inputs_embeds.shape[1] + == previous_hidden_states.shape[2] + == enorm_weight.shape[0] + == hnorm_weight.shape[0] + ) + assert inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous() + assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous() + + num_tokens, hidden = inputs_embeds.shape + enorm_out = torch.empty_like(inputs_embeds) + hnorm_out = torch.empty_like(previous_hidden_states) + if num_tokens == 0: + return enorm_out, hnorm_out + + block_size = triton.next_power_of_2(hidden) + _fused_mtp_input_rmsnorm_kernel[(num_tokens, hc_mult + 1)]( + inputs_embeds, + positions, + previous_hidden_states, + enorm_weight, + hnorm_weight, + enorm_out, + hnorm_out, + eps, + HIDDEN=hidden, + HC_MULT=hc_mult, + BLOCK_SIZE=block_size, + ) + return enorm_out, hnorm_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py new file mode 100644 index 00000000..0dd348a4 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_q_kv_rmsnorm_kernel( + q_ptr, + q_out_ptr, + q_weight_ptr, + q_in_stride, + q_out_stride, + kv_ptr, + kv_out_ptr, + kv_weight_ptr, + kv_in_stride, + kv_out_stride, + eps, + Q_SIZE: tl.constexpr, + KV_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y. + # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes + # the launch at max-num-batched-tokens >= 65536 with "invalid argument". + # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32 + # past num_tokens ~87K under large chunked prefill. + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + if pid_task == 0: + SIZE = Q_SIZE + row_in = q_ptr + token_idx * q_in_stride + weight_ptr = q_weight_ptr + row_out = q_out_ptr + token_idx * q_out_stride + else: + SIZE = KV_SIZE + row_in = kv_ptr + token_idx * kv_in_stride + weight_ptr = kv_weight_ptr + row_out = kv_out_ptr + token_idx * kv_out_stride + + # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which + # keep x, rrms, and w all in fp32 and perform a single cast at store. + block = tl.arange(0, BLOCK_SIZE) + mask = block < SIZE + x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / SIZE + rrms = tl.rsqrt(variance + eps) + w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32) + y = x * rrms * w + tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask) + + +def fused_q_kv_rmsnorm( + qr: torch.Tensor, + kv: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + assert qr.ndim == 2 and kv.ndim == 2 + assert qr.shape[0] == kv.shape[0], ( + f"token dim mismatch: qr={qr.shape}, kv={kv.shape}" + ) + assert qr.stride(-1) == 1 and kv.stride(-1) == 1 + assert q_weight.is_contiguous() and kv_weight.is_contiguous() + + q_size = qr.shape[1] + kv_size = kv.shape[1] + num_tokens = qr.shape[0] + qr_out = torch.empty_like(qr) + kv_out = torch.empty_like(kv) + if num_tokens == 0: + return qr_out, kv_out + + block_size = triton.next_power_of_2(max(q_size, kv_size)) + _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)]( + qr, + qr_out, + q_weight, + qr.stride(0), + qr_out.stride(0), + kv, + kv_out, + kv_weight, + kv.stride(0), + kv_out.stride(0), + eps, + Q_SIZE=q_size, + KV_SIZE=kv_size, + BLOCK_SIZE=block_size, + ) + return qr_out, kv_out diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py new file mode 100644 index 00000000..e3d7d38f --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +def save_partial_states( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + positions: torch.Tensor, + state_cache: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int, + state_width: int, + compress_ratio: int, + pdl_kwargs: dict | None = None, +) -> None: + """Write packed [kv, score+ape] partial states into the compressor cache. + + One program per token; pads (slot_id == -1) are skipped. + """ + num_actual = slot_mapping.shape[0] + head_size = kv.shape[-1] + _save_partial_states_kernel[(num_actual,)]( + kv, + kv.stride(0), + score, + score.stride(0), + ape, + ape.stride(0), + positions, + state_cache, + state_cache.stride(0), + state_cache.stride(1), + slot_mapping, + block_size, + HEAD_SIZE=head_size, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size), + STATE_WIDTH=state_width, + COMPRESS_RATIO=compress_ratio, + **(pdl_kwargs or {}), + ) + + +@triton.jit +def _save_partial_states_kernel( + kv_ptr, + kv_stride, + score_ptr, + score_stride, + ape_ptr, + ape_stride, + positions_ptr, + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + slot_mapping_ptr, + block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide. + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +): + token_idx = tl.program_id(0) + slot_id = tl.load(slot_mapping_ptr + token_idx) + + # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used + # by vLLM). During CUDA graph replay the batch may contain padding + # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an + # illegal memory access. + if slot_id < 0: + return + + block_idx = slot_id // block_size + pos_in_block = slot_id % block_size + base_ptr = ( + state_cache_ptr + + block_idx * state_cache_stride0 + + pos_in_block * state_cache_stride1 + ) + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + + kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask) + tl.store(base_ptr + block, kv, mask=mask) + + # Fused: score += ape[position % compress_ratio] + position = tl.load(positions_ptr + token_idx) + ape_row = position % COMPRESS_RATIO + ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask) + score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask) + tl.store( + base_ptr + STATE_WIDTH + block, + score + ape, + mask=mask, + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py b/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py new file mode 100644 index 00000000..44ae3286 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepseekV4 rotary embedding initialization.""" + +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding + + +def build_deepseek_v4_rope( + config, + *, + head_dim: int, + rope_head_dim: int, + max_position_embeddings: int, + compress_ratio: int, +) -> RotaryEmbedding: + rope_parameters = config.rope_parameters + rope_parameters["rope_theta"] = ( + config.compress_rope_theta if compress_ratio > 1 else config.rope_theta + ) + if rope_parameters["rope_type"] != "default": + rope_parameters["rope_type"] = ( + "deepseek_yarn" + if rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + rope_parameters["mscale"] = 0 # Disable mscale + rope_parameters["mscale_all_dim"] = 0 # Disable mscale + rope_parameters["is_deepseek_v4"] = True + rope_parameters["rope_dim"] = rope_head_dim + return get_rope( + head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py b/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py new file mode 100644 index 00000000..f36dc8f1 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from typing import Any, ClassVar, cast + +import torch +from torch import nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( + compress_norm_rope_store_triton, +) +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE +from vllm.models.deepseek_v4.common.ops.save_partial_states import ( + save_partial_states, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadataBuilder, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) + + +class CompressorBackend(AttentionBackend): + def __init__(self): + super().__init__() + + @staticmethod + def get_name() -> str: + return "CompressorBackend" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(1)] + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512, 1024] + + @staticmethod + def get_builder_cls() -> type["CompressorMetadataBuilder"]: + return CompressorMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + if include_num_layers_dimension: + return (0, 1, 2, 3) + return (0, 1, 2) + + +@dataclass +class CompressorMetadata: + block_table: torch.Tensor + slot_mapping: torch.Tensor + block_size: int + + token_to_req_indices: torch.Tensor | None = None # [num_tokens] + + +class CompressorMetadataBuilder(AttentionMetadataBuilder): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) + mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) + self.block_size = mla_spec.block_size + + self.token_to_req_indices = torch.zeros( + self.vllm_config.scheduler_config.max_num_batched_tokens, + dtype=torch.int32, + device=self.device, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> CompressorMetadata: + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + num_reqs = common_attn_metadata.num_reqs + query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() + token_to_req_indices = self.token_to_req_indices[: x.shape[0]] + token_to_req_indices.copy_(x, non_blocking=True) + return CompressorMetadata( + block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), + slot_mapping=common_attn_metadata.slot_mapping, + block_size=self.block_size, + token_to_req_indices=token_to_req_indices, + ) + + +class CompressorStateCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + state_dim: int, + dtype: torch.dtype, + compress_ratio: int, + prefix: str, + ): + super().__init__() + self.state_dim = state_dim + self.dtype = dtype + self.prefix = prefix + self.kv_cache = torch.tensor([]) + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + assert self.dtype == torch.float32 + assert compress_ratio in [4, 128] + coff = 1 + (compress_ratio == 4) + self.sliding_window = coff * compress_ratio + # Block size is constrained by tensor sharing between compressor states + # and KV blocks. Since compressor states share the same physical tensor + # as KV blocks, they must use the same page size. + # The KV block shape [256//4, head_dim] = [64, 584] determines: + # - C4 compressor block shape [4, 2*512*2*4] -> block_size = 4 + # - C128 compressor block shape [8, 512*2*4] -> block_size = 8 + # TODO(yifan): make block size automatically determined and configurable. + if compress_ratio == 4: + self.block_size = 4 + elif compress_ratio == 128: + self.block_size = 8 + else: + raise ValueError(f"Invalid compress ratio: {compress_ratio}") + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return SlidingWindowMLASpec( # only has one vector instead of K + V + block_size=self.block_size, + num_kv_heads=1, + head_size=self.state_dim, + dtype=self.dtype, + sliding_window=self.sliding_window, + alignment=576, # NOTE: FlashMLA requires 576B alignment + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return CompressorBackend + + +class DeepseekCompressor(nn.Module): + """DeepSeek V4 KV/score compressor. + + Owns the linear / norm / state-cache / ape state and the shared forward + prologue (kv/score split, save_partial_states launch). The + compress → norm → RoPE → store step is dispatched to a triton kernel + (``compress_norm_rope_store_triton``) by default, except for the NVIDIA + head_dim=128 indexer path which uses the cutedsl kernel + (``compress_norm_rope_store_cutedsl``) for better performance. + """ + + def __init__( + self, + vllm_config: VllmConfig, + compress_ratio: int, + hidden_size: int, + head_dim: int, + rotate: bool = False, + prefix: str = "", + k_cache_prefix="", + use_fp4_cache: bool = False, + ): + super().__init__() + self.compress_ratio = compress_ratio + self.hidden_size = hidden_size + self.head_dim = head_dim + self.rotate = rotate + self.prefix = prefix + self.k_cache_prefix = k_cache_prefix + self.use_fp4_cache = use_fp4_cache + + config = vllm_config.model_config.hf_config + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.rms_norm_eps = config.rms_norm_eps + self.device = current_platform.device_type + self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.max_model_len = vllm_config.model_config.max_model_len + + self.overlap = compress_ratio == 4 + self.coff = 1 + self.overlap + + state_dtype = torch.float32 + self.ape = nn.Parameter( + torch.empty( + (compress_ratio, self.coff * self.head_dim), + dtype=state_dtype, + device=self.device, + ), + requires_grad=False, + ) + + self.fused_wkv_wgate = MergedColumnParallelLinear( + self.hidden_size, + [self.coff * self.head_dim, self.coff * self.head_dim], + bias=False, + return_bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.fused_wkv_wgate", + ) + self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) + + self.state_cache = CompressorStateCache( + state_dim=2 * self.coff * self.head_dim, # kv_state + score_state + dtype=state_dtype, + compress_ratio=compress_ratio, + prefix=f"{prefix}.state_cache", + ) + + # Save reference to static_forward_context for forward-time KV cache lookup. + # get_current_vllm_config() is only available during __init__, not forward. + self._static_forward_context = ( + vllm_config.compilation_config.static_forward_context + ) + + if self.head_dim == 512: + assert not use_fp4_cache, ( + "MXFP4 cache is only supported for indexer (head=128)" + ) + self._quant_block = 64 + self._token_stride = self.nope_head_dim + self.rope_head_dim * 2 + self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad + elif self.head_dim == 128: + if use_fp4_cache: + self._quant_block = MXFP4_BLOCK_SIZE + self._token_stride = self.head_dim // 2 + self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE + else: + self._quant_block = 128 + self._token_stride = self.head_dim + self._scale_dim = 4 # single float32 scale + else: + raise ValueError( + f"Unsupported head_dim for fused quant+cache: {self.head_dim}" + ) + + def forward( + self, + # [num_tokens, 2 * self.coff * self.head_dim] + kv_score: torch.Tensor, + # [num_tokens] + positions: torch.Tensor, + rotary_emb, + ) -> None: + # Each of shape [num_tokens, coff * self.head_dim] + # input bf16, output are fp32 + kv, score = kv_score.split( + [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1 + ) + + # Get the metadata and handle dummy profiling run. + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return + + state_metadata = cast( + CompressorMetadata, attn_metadata[self.state_cache.prefix] + ) + token_to_req_indices = state_metadata.token_to_req_indices + slot_mapping = state_metadata.slot_mapping + num_actual = slot_mapping.shape[0] + block_table = state_metadata.block_table + block_size = state_metadata.block_size + + # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim + state_cache = self.state_cache.kv_cache + # kv_state stored in first half, score_state stored in second half + state_width = state_cache.shape[-1] // 2 + pdl_kwargs = ( + {} + if current_platform.is_rocm() or current_platform.is_xpu() + else {"launch_pdl": False} + ) + + # Store the KV and score (with fused APE addition) in the state. + # NOTE: PDL is disabled — both this kernel and the compress kernels + # below depend on preceding kernel outputs (kv/score from the cublas + # GEMM; state_cache from this kernel) but neither emits/waits on PDL + # grid dependency primitives, so launch_pdl=True caused a + # read-after-write race and non-deterministic output. + save_partial_states( + kv=kv, + score=score, + ape=self.ape, + positions=positions, + state_cache=state_cache, + slot_mapping=slot_mapping, + block_size=block_size, + state_width=state_width, + compress_ratio=self.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) + + # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write. + # RoPE requirements (kernel applies forward GPT-J style rotation): + # - is_neox_style=False (interleaved pairs, NOT split-half) + # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos, + # second half sin (per-pair, length rope_head_dim // 2 each) + # - applied to LAST rope_head_dim elements of head_dim + # - position used: (positions // compress_ratio) * compress_ratio + cos_sin_cache = rotary_emb.cos_sin_cache + k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) + kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache + + if current_platform.is_cuda(): + # NVIDIA GPUs. + if self.head_dim == 512: + from .nvidia.ops.sparse_attn_compress_cutedsl import ( + compress_norm_rope_store_cutedsl, + ) + + # Main compressor path. + # Use a cutedsl kernel for better performance. + compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl + else: + # Indexer path (head_dim == 128). + # Use a triton kernel. + compress_norm_rope_store_fn = compress_norm_rope_store_triton + else: + # AMD GPUs. + # Always use a triton kernel. + compress_norm_rope_store_fn = compress_norm_rope_store_triton + + compress_norm_rope_store_fn( + state_cache=state_cache, + num_actual=num_actual, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + state_width=state_width, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + k_cache_metadata=k_cache_metadata, + pdl_kwargs=pdl_kwargs, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + compress_ratio=self.compress_ratio, + overlap=self.overlap, + use_fp4_cache=self.use_fp4_cache, + rms_norm_weight=self.norm.weight, + rms_norm_eps=self.rms_norm_eps, + quant_block=self._quant_block, + token_stride=self._token_stride, + scale_dim=self._scale_dim, + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py new file mode 100644 index 00000000..208f01a7 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py new file mode 100644 index 00000000..5c8b08d4 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import abstractmethod +from typing import TYPE_CHECKING, ClassVar, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.common.ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + MultipleOf, + SparseMLAAttentionImpl, +) +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + FlashMLASparseBackend, + FlashMLASparseMetadata, +) +from vllm.v1.attention.ops.flashmla import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.models.deepseek_v4.attention import ( + DeepseekV4MLAAttention, + ) + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + +class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]): + """Abstract parent for DeepseekV4 sparse MLA impls. + + V4 sparse MLA is driven by the layer (``DeepseekV4MLAAttention.forward``) + rather than the v1 framework, so ``forward_mqa`` is overridden with a + classmethod that takes the layer as its first argument. This Liskov-broken + override is intentional: the grandparent's instance-method ``forward_mqa`` + is never called on V4 layers. + """ + + backend_cls: ClassVar[type[AttentionBackend]] + + # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather + # workspace allocated in _forward_prefill and is also read by the V4 layer's + # dummy-run path to pre-reserve that workspace. + PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + + @classmethod + @abstractmethod + def forward_mqa( # type: ignore[override] + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + raise NotImplementedError + + @classmethod + @abstractmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + """Q head count the backend wants q allocated at. + + The MLA wrapper allocates the q/output buffers at + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must + satisfy ``result >= num_heads``. Backends with no padding constraint + return ``num_heads``. + """ + raise NotImplementedError + + +class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend): + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [256] + + @staticmethod + def get_name() -> str: + return "V4_FLASHMLA_SPARSE" + + @staticmethod + def get_impl_cls() -> type["DeepseekV4SparseMLAAttentionImpl"]: + return DeepseekV4FlashMLASparseImpl + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512 (overrides the + # V3.2 default of 576 from FlashMLASparseBackend). + return [512] + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if cache_dtype_str == "fp8_ds_mla": + # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale). + # head_size passed in is the semantic head_dim (512). + return (num_blocks, block_size, 584) + else: + return (num_blocks, block_size, head_size) + + +class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl): + """FlashMLA sparse MLA implementation for DeepSeek V4's custom MLA layer.""" + + backend_cls = DeepseekV4FlashMLASparseBackend + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + # FP8 decode kernel only supports h_q = 64 or 128. + if num_heads > 128: + raise ValueError( + f"DeepseekV4 FlashMLA does not support {num_heads} heads " + "(FP8 decode kernel requires h_q in {64, 128})." + ) + return 64 if num_heads <= 64 else 128 + + @classmethod + def forward_mqa( # type: ignore[override] + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + # Get SWA and indexer metadata from forward context + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + if attn_metadata is None: + # Warmup dummy run: no real metadata. Reserve the same bf16 + # gather workspace _forward_prefill would; the dequantize / topk + # / sparse_fwd kernels are skipped this step. + swa_only = layer.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (layer.max_model_len + layer.compress_ratio - 1) + // layer.compress_ratio + ) + M = N + layer.window_size + layer.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) + output.zero_() + return + + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + FlashMLASparseMetadata | None, attn_metadata.get(layer.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(layer.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = layer.compress_ratio <= 1 + # SWA-only layers (compress_ratio <= 1) don't have their own KV cache + # allocation, so layer.kv_cache may be empty after profiling cleanup. + self_kv_cache = layer.kv_cache if not swa_only else None + swa_kv_cache = layer.swa_cache_layer.kv_cache + + # Split prefill and decode + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + cls._forward_prefill( + layer=layer, + q=q[num_decode_tokens:], + positions=positions[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + cls._forward_decode( + layer=layer, + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + @classmethod + def _forward_decode( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: FlashMLASparseMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // layer.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if layer.compress_ratio == 4: + # C4A: local indices differ per layer (filled by Indexer). + assert layer.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + layer.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + # C128A: pre-computed during metadata build. + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + + # We treat queries in the same seq as different queries + # and later we only attend by generated indices. + # q arrives pre-padded to layer.padded_heads by the outer wrapper. + q = q.unsqueeze(1) + + # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes) + # Use unsqueeze to preserve strides (handles padded blocks correctly) + swa_cache = layer.swa_cache_layer.kv_cache.unsqueeze(-2) + # Reshape KV cache to (num_blocks, block_size, 1, head_bytes) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + + # One FlashMLASchedMeta per layer type, shared across all same-type + # layers within this decode step. The first forward call per type + # triggers the in-kernel planner (allocating tile_scheduler_metadata + # and num_splits via PyTorch's graph-aware allocator so CUDA graph + # capture reuses the same addresses on replay); subsequent same-type + # layers see have_initialized=True and skip the planner. + if layer.compress_ratio <= 1: + tile_metadata = swa_metadata.tile_sched_swaonly + elif layer.compress_ratio == 4: + tile_metadata = swa_metadata.tile_sched_c4a + elif layer.compress_ratio == 128: + tile_metadata = swa_metadata.tile_sched_c128a + else: + raise ValueError( + f"Unsupported compress_ratio={layer.compress_ratio}; " + "expected 1, 4, or 128." + ) + assert tile_metadata is not None, ( + "swa_metadata missing tile_sched entry for " + f"compress_ratio={layer.compress_ratio}; " + "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not " + "allocate one for this layer type." + ) + + out, _ = flash_mla_with_kvcache( + q=q, + k_cache=swa_cache, + block_table=None, + head_dim_v=512, + tile_scheduler_metadata=tile_metadata, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_indices, + topk_length=swa_lens, + softmax_scale=layer.scale, + attn_sink=layer.attn_sink, + extra_k_cache=kv_cache if not swa_only else None, + extra_indices_in_kvcache=topk_indices, + extra_topk_length=topk_lens, + out=output.unsqueeze(1), + ) + + @classmethod + def _forward_prefill( + cls, + layer: "DeepseekV4MLAAttention", + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashMLASparseMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = attn_metadata is None + + num_prefills = swa_metadata.num_prefills + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + # Use pre-computed prefill metadata. + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + assert seq_lens is not None + assert gather_lens is not None + + # Derive prefill-local token offsets from the full query_start_loc_cpu. + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert query_start_loc_cpu is not None + assert query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if layer.compress_ratio == 4: + assert layer.topk_indices_buffer is not None + topk_indices = layer.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + # C128A: pre-computed during metadata build. + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + top_k = topk_indices.shape[-1] + # Compressed region must fit the full compressed pool (seq_len // + # compress_ratio), not just top_k. top_k bounds how many indices + # the indexer selects, not the pool size it indexes into. + N = (layer.max_model_len + layer.compress_ratio - 1) // layer.compress_ratio + else: + # NOTE(woosuk): topk_indices will not be used for SWA-only layers. + assert layer.topk_indices_buffer is not None + topk_indices = layer.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + N = 0 + + M = N + layer.window_size + layer.max_num_batched_tokens + chunk_size_const = cls.PREFILL_CHUNK_SIZE + num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const + + workspace_manager = current_workspace_manager() + kv = workspace_manager.get_simultaneous( + ((chunk_size_const, M, q.shape[-1]), torch.bfloat16), + )[0] + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size_const + chunk_end = min(chunk_start + chunk_size_const, num_prefills) + chunk_size = chunk_end - chunk_start + if not swa_only: + # Gather compressed KV + assert attn_metadata is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // layer.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // layer.compress_ratio, + offset=0, + ) + + # Gather SWA KV + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=N, + ) + + # Combine the topk indices and SWA indices for gathered KV cache + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + layer.window_size, + layer.compress_ratio, + top_k, + M, + N, + ) + flash_mla_sparse_fwd( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices.unsqueeze(1), + sm_scale=layer.scale, + attn_sink=layer.attn_sink, + topk_length=combined_lens, + out=output[query_start:query_end], + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py new file mode 100644 index 00000000..547048ab --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py @@ -0,0 +1,1476 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import typing +from collections.abc import Callable, Iterable, MutableSequence, Sequence +from itertools import islice + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.kernels.mhc.tilelang import ( + hc_head_fused_kernel_tilelang, + mhc_fused_post_pre_tilelang, + mhc_post_tilelang, + mhc_pre_tilelang, +) +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v4.attention import ( + DeepseekV4Indexer, + DeepseekV4MLA, +) +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.sequence import IntermediateTensors + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + swiglu_limit: float | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + + # If is_sequence_parallel, the input and output tensors are sharded + # across the ranks within the tp_group. In this case the weights are + # replicated and no collective ops are needed. + # Otherwise we use standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + if swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +def make_deepseek_v4_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + return [ + ( + "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_", + f"experts.{expert_id}.{weight_name}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id, weight_name in [ + ("w1", "w1"), + ("w2", "w2"), + ("w3", "w3"), + ] + ] + + +class DeepseekV4MegaMoEExperts(nn.Module): + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + + def __init__( + self, + vllm_config: VllmConfig, + *, + num_experts: int, + num_local_experts: int, + experts_start_idx: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + num_logical_experts: int | None = None, + ): + super().__init__() + self.prefix = prefix + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.experts_start_idx = experts_start_idx + self.experts_end_idx = experts_start_idx + num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + + self.num_logical_experts = ( + num_logical_experts if num_logical_experts is not None else num_experts + ) + + self.eplb_state = EplbLayerState() + + weight_attrs = {"weight_loader": self.weight_loader} + self.w13_weight = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight, weight_attrs) + + self.w13_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 32, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight_scale, weight_attrs) + self.w13_weight_scale.quant_method = "block" + + self.w2_weight = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight, weight_attrs) + + self.w2_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 32, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight_scale, weight_attrs) + self.w2_weight_scale.quant_method = "block" + + self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None + self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None + + # Register in the static forward context so the custom-op wrapper + # can look up this module by name from within a torch.compile graph. + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _map_global_expert_id(self, expert_id: int) -> list[int]: + """Return local (per-rank) slot offsets where logical expert + `expert_id` should land on this rank. + """ + physical_ids: list[int] = [] + for p in range(self.experts_start_idx, self.experts_end_idx): + if p % self.num_logical_experts == expert_id: + physical_ids.append(p - self.experts_start_idx) + return physical_ids + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ) -> bool | None: + local_expert_ids = self._map_global_expert_id(expert_id) + if not local_expert_ids: + return False if return_success else None + + loaded_any = False + for local_expert_id in local_expert_ids: + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if "w13_" not in weight_name: + continue + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow( + 0, shard_offset, self.intermediate_size + ) + elif shard_id == "w2": + if "w2_" not in weight_name: + continue + else: + raise ValueError(f"Unsupported expert shard id: {shard_id}") + + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{weight_name}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + loaded_any = True + + if return_success: + return loaded_any + return None + + @staticmethod + def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: + return (sf.to(torch.int32) << 23).view(torch.float32) + + def _check_runtime_supported(self) -> None: + device = self.w13_weight.device + if torch.cuda.get_device_capability(device)[0] != 10: + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, 32), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, 32), + self.num_local_experts, + ) + self._transformed_l1_weights, self._transformed_l2_weights = ( + deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + ) + ) + # Drop the original loader-side parameters: the MegaMoE kernels only + # consume the transformed views above. transform_weights_for_mega_moe + # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights) + # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that + # aliases the original storage, and _transformed_l2_weights still holds + # it, so the storage stays live after we drop the Parameter. + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + def get_symm_buffer(self): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + symm_buffer = self._symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + self._symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, + ) + + def get_expert_weights(self) -> list[torch.Tensor]: + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + + def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: + """Return a (num_local_experts, -1) view with contiguous memory layout.""" + assert t.shape[0] == self.num_local_experts + if t.is_contiguous(): + return t.view(self.num_local_experts, -1) + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: + # scales have shape (E, M, N) with memory layout (E, N, M) + back = torch.transpose(t, 1, 2) + assert back.is_contiguous() + return back.view(self.num_local_experts, -1) + + raise AssertionError( + f"DSv4 EPLB {name}: non-contiguous expert tensor with " + f"unexpected layout shape={tuple(t.shape)} " + f"stride={tuple(t.stride())} dtype={t.dtype}" + ) + + return [ + _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), + _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), + _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), + _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), + ] + + def update_expert_map(self) -> None: + pass + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but the symmetric buffer was sized for {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + + # EPLB: map logical expert IDs to physical replicas and record load. + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + ) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + ) + + # This method must have been already called during the weight loading phase. + # We call it again here to cover the dummy weight loading case. + self.finalize_weights() + + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) + return y + + +DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.prefix = prefix + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hidden_size = config.hidden_size + + self.n_routed_experts = config.n_routed_experts + self.n_activated_experts = config.num_experts_per_tok + self.moe_intermediate_size = config.moe_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." + ) + if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": + raise NotImplementedError( + "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" + f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" + "deep_gemm_mega_moe for this checkpoint." + ) + + self.gate = GateLinear( + input_size=config.hidden_size, + output_size=config.n_routed_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers + self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + if is_hash_moe: + # hash MoE doesn't use e_score_correction_bias + # Use randint instead of empty to avoid garbage values causing + # invalid memory access in dummy mode (--load-format="dummy") + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=self.hash_indices_dtype, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + + if config.n_shared_experts is None: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = DeepseekV4MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + swiglu_limit=self.swiglu_limit, + quant_config=quant_config, + reduce_results=self.use_mega_moe, + prefix=f"{prefix}.shared_experts", + ) + + if self.use_mega_moe: + self._init_mega_moe_experts(vllm_config, config, prefix) + else: + self._init_fused_moe_experts(config, quant_config, prefix) + + def _init_mega_moe_experts( + self, + vllm_config: VllmConfig, + config, + prefix: str, + ) -> None: + self.ep_group = get_ep_group() + self.ep_size = self.ep_group.world_size + self.ep_rank = self.ep_group.rank_in_group + + eplb_config = vllm_config.parallel_config.eplb_config + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_routed_experts = config.n_routed_experts + self.n_shared_experts = config.n_shared_experts or 0 + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + assert self.n_physical_experts % self.ep_size == 0, ( + f"n_physical_experts={self.n_physical_experts} must be divisible by " + f"ep_size={self.ep_size}. Adjust num_redundant_experts." + ) + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + self.n_local_experts = self.n_local_physical_experts + self.experts_start_idx = self.physical_expert_start + self.experts_end_idx = self.physical_expert_end + + self.experts = DeepseekV4MegaMoEExperts( + vllm_config, + num_experts=self.n_physical_experts, + num_local_experts=self.n_local_physical_experts, + experts_start_idx=self.physical_expert_start, + num_logical_experts=self.n_logical_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + prefix=f"{prefix}.experts", + ) + + def _init_fused_moe_experts( + self, + config, + quant_config, + prefix: str, + ) -> None: + self.tp_rank = get_tensor_model_parallel_rank() + assert config.n_routed_experts % self.tp_size == 0 + + self.n_local_experts = config.n_routed_experts // self.tp_size + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.n_redundant_experts = 0 + self.n_shared_experts = config.n_shared_experts or 0 + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_local_physical_experts = self.n_local_experts + self.physical_expert_start = self.experts_start_idx + self.physical_expert_end = self.experts_end_idx + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + swiglu_limit=self.swiglu_limit, + router_logits_dtype=torch.float32, + ) + + def forward( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + + if not self.use_mega_moe: + return self._forward_fused_moe(hidden_states, input_ids) + + org_shape = hidden_states.shape + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + ) + activation_clamp = ( + float(self.swiglu_limit) if self.swiglu_limit is not None else None + ) + final_hidden_states = self.experts( + hidden_states, + topk_weights, + topk_ids, + activation_clamp=activation_clamp, + ) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + final_hidden_states += shared_output + + return final_hidden_states.view(org_shape) + + def _forward_fused_moe( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + org_shape = hidden_states.shape + if self.experts.is_internal_router: + # In this case, the gate/router runs inside the FusedMoE class + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=hidden_states, + input_ids=input_ids, + ) + else: + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + + return final_hidden_states.view(org_shape) + + def finalize_mega_moe_weights(self) -> None: + if self.use_mega_moe: + self.experts.finalize_weights() + + +class DeepseekV4Attention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + layer_id = extract_layer_index(prefix) + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0 + + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.max_position_embeddings = config.max_position_embeddings + + # Padded to min 64 heads for FlashMLA, initialized to -inf + # (no sink effect). Weight loading fills the first n_local_heads slots. + padded_heads = max(self.n_local_heads, 64) + self.attn_sink = nn.Parameter( + torch.full((padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + self.softmax_scale = self.head_dim**-0.5 + self.scale_fmt = config.quantization_config["scale_fmt"] + + self.rope_parameters = config.rope_scaling + + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + compress_ratio=self.compress_ratio, + ) + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is + # free here (outer GEMMs joined) for the inner overlap of + # wq_b+fused_indexer_q_rope_quant vs compressor. + indexer_aux_stream = ( + aux_stream_list[2] if aux_stream_list is not None else None + ) + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=vllm_config.cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + aux_stream=indexer_aux_stream, + ) + + self.mla_attn = DeepseekV4MLA( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, + vllm_config=vllm_config, + fused_wqa_wkv=self.fused_wqa_wkv, + q_norm=self.q_norm, + wq_b=self.wq_b, + kv_norm=self.kv_norm, + wo_a=self.wo_a, + wo_b=self.wo_b, + attn_sink=self.attn_sink, + rotary_emb=self.rotary_emb, + indexer=self.indexer, + indexer_rotary_emb=self.rotary_emb, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + window_size=self.window_size, + compress_ratio=self.compress_ratio, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ): + return self.mla_attn(positions, hidden_states, llama_4_scaling) + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + + self.rms_norm_eps = config.rms_norm_eps + self.attn = DeepseekV4Attention( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + attn_norm_weight = self.attn_norm.weight.data + attn_norm_eps = self.attn_norm.variance_epsilon + if residual is None: + # Run standalone mhc_pre on first layer + residual = x + post_mix, res_mix, x = mhc_pre_tilelang( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + else: + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=attn_norm_weight, + norm_eps=attn_norm_eps, + ) + + # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above. + x = self.attn(positions, x, None) + + 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( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + n_splits=1, + tile_n=1, + norm_weight=ffn_norm_weight, + norm_eps=ffn_norm_eps, + ) + + x = self.ffn(x, input_ids) + return x, residual, post_mix, res_mix + + +class DeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # Three aux streams: one per non-default input GEMM in + # DeepseekV4MLA.attn_gemm_parallel_execute + # (compressor kv_score, indexer.weights_proj, indexer.compressor + # kv_score). fused_wqa_wkv stays on the default stream. + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + + # Reserved topk indices buffer for all Indexer layers to reuse. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_dim, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty( + self.hc_mult, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + # 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: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + ) + else: + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + # PP intermediate tensors carry the multi-stream hidden_states + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the + # token embedding to hc_mult streams before the first decoder + # layer and keeps that shape until hc_head() collapses it. + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, self.config.hidden_size), + dtype=dtype, + device=device, + ), + } + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + if self.use_mega_moe: + input_ids = input_ids.to(torch.int64) + + residual, post_mix, res_mix = None, None, None + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + if layer is not None: + hidden_states = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) + + 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)) + + 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) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = self.get_expert_mapping() + + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name, self): + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + if ".experts." in name: + # E8M0 scales are stored as float8_e8m0fnu in + # checkpoints but the MoE param is uint8. copy_() + # would do a numeric conversion (e.g. 2^-7 → 0), + # destroying the raw exponent bytes. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + loaded_params.add(name_mapped) + continue + elif "attn_sink" in name: + if is_pp_missing_parameter(name, self): + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) + if first_layer.ffn.use_mega_moe: + return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + + def finalize_mega_moe_weights(self) -> None: + for layer in islice(self.layers, self.start_layer, self.end_layer): + layer.ffn.finalize_mega_moe_weights() + + +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: + if expert_dtype == "fp4": + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as + # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and + # shared experts use Fp8LinearMethod's block scales, which + # register as ``weight_scale_inv``. + scale_regex = { + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", + re.compile(r"\.scale$"): ".weight_scale_inv", + } + else: + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys + # there. + scale_regex = { + re.compile(r"\.scale$"): ".weight_scale_inv", + } + return WeightsMapper( + orig_to_new_prefix={ + "layers.": "model.layers.", + "embed.": "model.embed.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", + }, + orig_to_new_regex=scale_regex, + orig_to_new_suffix={ + "head.weight": "lm_head.weight", + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".attn.compressor.": ".attn.mla_attn.compressor.", + ".shared_experts.w2": ".shared_experts.down_proj", + }, + ) + + +class DeepseekV4MixtureOfExperts(MixtureOfExperts): + moe_mlp_layers: list["DeepseekV4MoE"] + + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None: + if example_moe is None: + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_shared_experts = 0 + self.num_redundant_experts = 0 + return + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_shared_experts = example_moe.n_shared_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts): + model_cls = DeepseekV4Model + + # Default mapper assumes the original FP4-expert checkpoint layout. + # Overridden per-instance in __init__ when expert_dtype != "fp4". + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + self.config = config + expert_dtype = getattr(config, "expert_dtype", "fp4") + if expert_dtype != "fp4": + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype) + + self.model = self.model_cls( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = [] + self.num_expert_groups = getattr(self.config, "n_group", 1) + self.num_moe_layers = self.config.num_hidden_layers + self.moe_layers: list[nn.Module] = [] + self.moe_mlp_layers: list[DeepseekV4MoE] = [] + example_moe: DeepseekV4MoE | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if not isinstance(layer, DeepseekV4DecoderLayer): + continue + if isinstance(layer.ffn, DeepseekV4MoE): + example_moe = layer.ffn + self.moe_mlp_layers.append(layer.ffn) + self.moe_layers.append(layer.ffn.experts) + + self.num_moe_layers = len(self.moe_layers) + self.extract_moe_parameters(example_moe) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer (max_num_batched_tokens, + hc_mult * hidden_size) for the MTP draft model. Populated by + forward(); valid after each target step.""" + return getattr(self.model, "_mtp_hidden_buffer", None) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.model.finalize_mega_moe_weights() + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py new file mode 100644 index 00000000..133a96e3 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py @@ -0,0 +1,516 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4). + +Split from ``deepseek_mtp.py`` because the V4 architecture introduces several +pieces that have no analogue in V3/V32: + * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of + the fused ``eh_proj``); + * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``; + * ``DeepseekV4DecoderLayer`` with its own aux-stream management; + * V4-specific checkpoint weight-name remapping in ``load_weights``. +""" + +import typing +from collections.abc import Callable, Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.kernels.mhc.tilelang import ( + hc_head_fused_kernel_tilelang, + mhc_post_tilelang, +) +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.deepseek_mtp import SharedHead +from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name +from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, + mtp_shared_head_rmsnorm, +) +from vllm.sequence import IntermediateTensors + +from .model import ( + DeepseekV4DecoderLayer, + make_deepseek_v4_expert_params_mapping, +) + +logger = init_logger(__name__) + +# MoE expert scales are fused into per-layer w13/w2 tensors. The exact +# parameter suffix depends on which FusedMoE method handles the experts: +# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``; +# - fp8 experts (Fp8MoEMethod with block_quant=True) register +# ``w{1,2,3}_weight_scale_inv``. +# Other FP8 linear scales (including shared experts) always use +# ``.weight_scale_inv``. Mirrors the per-instance mapper built by +# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py. +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DeepSeekV4MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + topk_indices_buffer: torch.Tensor, + prefix: str, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ) -> None: + super().__init__() + + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + self.rms_norm_eps = config.rms_norm_eps + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than + # fusing them the way V3 does with eh_proj. + self.e_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + return_bias=False, + quant_config=quant_config, + ) + self.h_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + return_bias=False, + quant_config=quant_config, + ) + + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + self.mtp_block = DeepseekV4DecoderLayer( + vllm_config, + prefix, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + # Target stashes pre-hc_head residual as flat (T, hc_mult * D); + # reshape to (T, hc_mult, D) — the training-time layout — before + # the fused norm pass so both inputs are 3D-friendly. + previous_hidden_states = previous_hidden_states.view( + -1, self.hc_mult, self.config.hidden_size + ) + # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. + inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( + inputs_embeds, + positions, + previous_hidden_states, + self.enorm.weight.data, + self.hnorm.weight.data, + self.enorm.variance_epsilon, + self.hc_mult, + ) + hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( + inputs_embeds + ).unsqueeze(-2) + hidden_states, residual, post_mix, res_mix = self.mtp_block( + positions=positions, x=hidden_states, input_ids=None + ) + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) + # Return the flat pre-hc_head residual so it can be re-fed as the + # next spec step's `previous_hidden_states` when + # num_speculative_tokens > 1. hc_head is deferred to compute_logits. + return hidden_states.flatten(1) + + +class DeepSeekV4MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + topk_tokens = config.index_topk + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + topk_tokens, + dtype=torch.int32, + ) + + # Three aux streams shared across all MTP layers, mirroring DeepseekV4Model. + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + + # to map the exact layer index from weights + self.layers = torch.nn.ModuleDict( + { + str(idx): DeepSeekV4MultiTokenPredictorLayer( + vllm_config, + self.topk_indices_buffer, + f"{prefix}.layers.{idx}", + aux_stream_list=aux_stream_list, + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply + # hc_head here so logits are computed from the dense hidden state. + hidden_states = hidden_states.view( + -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size + ) + hidden_states = hc_head_fused_kernel_tilelang( + hidden_states, + mtp_layer.hc_head_fn, + mtp_layer.hc_head_scale, + mtp_layer.hc_head_base, + mtp_layer.rms_norm_eps, + mtp_layer.hc_eps, + ) + hidden_states = mtp_shared_head_rmsnorm( + hidden_states, + mtp_layer.shared_head.norm.weight.data, + mtp_layer.shared_head.norm.variance_epsilon, + ) + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class DeepSeekV4MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = DeepSeekV4MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Weight name remapping for checkpoint compatibility. + # Maps checkpoint weight paths to model parameter paths. + WEIGHT_NAME_REMAPPING: dict[str, str] = { + ".emb.tok_emb.weight": ".embed_tokens.weight", + ".head.weight": ".shared_head.head.weight", + ".norm.weight": ".shared_head.norm.weight", + } + + def _remap_weight_name(name: str) -> str: + """Remap checkpoint weight names to model parameter names.""" + for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items(): + if old_pattern in name: + name = name.replace(old_pattern, new_pattern) + return name + + def _find_mtp_layer_idx(name: str) -> int: + subnames = name.split(".") + for subname in subnames: + try: + # we return the first encountered integer + return int(subname) + except ValueError: + continue + return 0 + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + first_layer = next(iter(self.model.layers.values())) + if first_layer.mtp_block.ffn.use_mega_moe: + expert_mapping = make_deepseek_v4_expert_params_mapping( + self.config.n_routed_experts + ) + else: + expert_mapping = FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + + # FP8 experts register ``..._weight_scale_inv`` (block_quant) while + # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix + # for the rename below based on the model's expert dtype. + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + for name, loaded_weight in weights: + mtp_layer_idx = _find_mtp_layer_idx(name) + # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to + # `model.layers.{num_hidden_layers + i}.*` so that + # get_spec_layer_idx_from_weight_name can identify them. + name = name.replace( + f"mtp.{mtp_layer_idx}.", + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", + ) + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + + name = _remap_weight_name(name) + name = self._rewrite_spec_layer_name(spec_layer, name) + + if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name: + continue + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + break + else: + if ".experts." in name: + # Reinterpret E8M0 scales as uint8 to preserve raw + # exponent bytes; numeric copy_() would zero them. + # Mirrors the main DeepseekV4 loader. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + loaded_params.add(name_mapped) + break + continue + elif "attn_sink" in name: + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if ".shared_experts.w2" in name: + name = name.replace( + ".shared_experts.w2", ".shared_experts.down_proj" + ) + if name.endswith(".ffn.gate.bias"): + # ``e_score_correction_bias`` lives on the gate + # under a different attribute name. + name = name.replace( + ".ffn.gate.bias", + ".ffn.gate.e_score_correction_bias", + ) + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may have " + f"been quantized without including the MTP layers. " + f"Use a checkpoint that includes MTP layer weights, " + f"or disable speculative decoding." + ) + self.finalize_mega_moe_weights() + logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def finalize_mega_moe_weights(self) -> None: + for layer in self.model.layers.values(): + layer.mtp_block.ffn.finalize_mega_moe_weights() + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """ + Rewrite the weight name to match the format of the original model. + Add .mtp_block for modules in transformer layer block for spec layer + and rename shared layer weights to be top level. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "h_proj", + "e_proj", + "shared_head", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + # treat rest weights as weights for transformer layer block + name = name.replace( + f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." + ) + elif shared_weight: + # treat shared weights as top level weights + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py new file mode 100644 index 00000000..20752bc1 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA-only (cutedsl/cutlass) kernels for DeepSeek V4. + +These modules import ``cutlass``/``cutedsl`` at module top level, so they must +not be imported on non-CUDA platforms. Callers should gate on +``vllm.utils.import_utils.has_cutedsl()`` before importing from here. + +This ``__init__`` deliberately imports nothing: re-exporting the cutedsl +modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for +anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses. +Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate. +""" diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py new file mode 100644 index 00000000..6c443246 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Int32, Uint8, Uint32 +from cutlass.cute.nvgpu import cpasync +from quack.compile_utils import make_fake_tensor + +from vllm.cute_utils import _bf16x2_mul, cvt + + +def dequantize_and_gather_k_cache_cutedsl( + out: torch.Tensor, + k_cache: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor | None, + block_table: torch.Tensor, + block_size: int, + offset: int, +) -> None: + DequantGatherKCacheKernel.compile( + block_size=block_size, + has_gather_lens=gather_lens is not None, + )(out, k_cache, seq_lens, gather_lens, block_table, offset) + + +class DequantGatherKCacheKernel: + # Hard-coded for DSv4. + head_dim = 512 + group_size = 64 # 1 scale per 64 elems + + def __init__(self, fp8_dim: int = 448, block_size: int = 64): + self.fp8_dim = fp8_dim + self.bf16_dim = self.head_dim - fp8_dim + self.data_dim = fp8_dim + self.bf16_dim * 2 + self.block_size = block_size + + self.num_warps = 4 + self.tb_size = self.num_warps * 32 + self.num_stages = 4 + + @cute.jit + def __call__( + self, + out: cute.Tensor, + k_cache: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + stream: CUstream, + ): + # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] + # block is actually a concat of + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + k_data = cute.make_tensor( + k_cache.iterator, + layout=cute.make_layout( + (k_cache.shape[0], self.block_size, self.data_dim), + stride=(k_cache.stride[0], self.data_dim, 1), + ), + ) + k_scale = cute.make_tensor( + k_cache.iterator + (self.block_size * self.data_dim), + layout=cute.make_layout( + (k_cache.shape[0], self.block_size, 8), + stride=(k_cache.stride[0], 8, 1), + ), + ) + + grid = (out.shape[0], 1024, 1) + self.kernel( + out, + k_data, + k_scale, + seq_lens, + gather_lens, + block_table, + offset, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.jit + def load_g2s( + self, + k_data_slice: cute.Tensor, + k_scale: cute.Tensor, + block_table: cute.Tensor, + s_kdata_slice: cute.Tensor, + s_kscale: cute.Tensor, + req_id, + pos, + lane_id, + stage_id, + ): + # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] + # s_kdata_slice: [(4, data_dim/16), num_stages] + + op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) + cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) + cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64) + page_id = block_table[req_id, pos // self.block_size] + block_offset = pos % self.block_size + + # Load the first 512 bytes (32x16B). + idx = lane_id + src = k_data_slice[page_id, block_offset, (None, idx)] + cute.copy( + cp16_atom, + cute.recast_tensor(src, Uint32), + s_kdata_slice[(None, idx), stage_id], + ) + + # Load the tail 64 bytes. + idx += 32 + if idx < cutlass.const_expr(self.data_dim // 16): + src = k_data_slice[page_id, block_offset, (None, idx)] + cute.copy( + cp16_atom, + cute.recast_tensor(src, Uint32), + s_kdata_slice[(None, idx), stage_id], + ) + elif idx == cutlass.const_expr(self.data_dim // 16): + cute.copy( + cp8_atom, + k_scale[page_id, block_offset, None], + s_kscale[None, stage_id], + ) + + @cute.kernel + def kernel( + self, + out: cute.Tensor, + k_data: cute.Tensor, + k_scale: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + ): + req_id, worker_id, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + + _, num_workers, _ = cute.arch.grid_dim() + + # Prepare smem. + smem = cutlass.utils.SmemAllocator() + s_kdata = smem.allocate_tensor( + Uint32, + cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)), + byte_alignment=16, + )[None, warp_id, None] + s_kscale = smem.allocate_tensor( + Uint8, + cute.make_layout((8, self.num_warps, self.num_stages)), + byte_alignment=8, + )[None, warp_id, None] + + # Prepare for 16B cp.async, also for BF16 smem loads later. + k_data_slice = cute.logical_divide(k_data, (None, None, 16)) + s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) + + # Load FP8 elems in 8B units, so once dequantized, they are 16B units. + s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) + + # 16B st.global. + out_slice = cute.logical_divide(out, (None, None, 8)) + + cp_op = cute.nvgpu.CopyUniversalOp() + cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + seq_len = seq_lens[req_id] + gather_len = seq_len + if cutlass.const_expr(gather_lens is not None): + gather_len = gather_lens[req_id] # type: ignore[index] + start_pos = seq_len - gather_len + + # Start prefetch. + for i in cutlass.range_constexpr(self.num_stages - 1): + next_pos = ( + start_pos + + worker_id * self.num_warps + + warp_id + + i * num_workers * self.num_warps + ) + if next_pos < seq_len: + self.load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + i, + ) + cute.arch.cp_async_commit_group() + prefetch_stage = self.num_stages - 1 + compute_stage = 0 + + # Main loop. + for i in range( + worker_id * self.num_warps + warp_id, + gather_len, + num_workers * self.num_warps, + ): + pos = start_pos + i + + # Prefetch next stage. + next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1) + if next_pos < seq_len: + self.load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + prefetch_stage, + ) + prefetch_stage = (prefetch_stage + 1) % self.num_stages + cute.arch.cp_async_commit_group() + + # Wait for gmem->smem to finish. + cute.arch.cp_async_wait_group(self.num_stages - 1) + cute.arch.sync_warp() + + # There are 512 elems per token. As a warp, data0 holds the first + # 256 elems and data1 holds the second 256 elems, i.e. each thread + # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as + # contiguous 16B global stores. On Blackwell, this might not be + # necessary as we have 32B global stores, but doing it this way + # does not seem to be slower. + data0 = cute.make_rmem_tensor((2,), Uint32) + data1 = cute.make_rmem_tensor((2,), Uint32) + cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0) + cute.copy( + cp8_atom, + s_kdata_8B_slice[(None, lane_id + 32), compute_stage], + data1, + ) + + # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 + # elements. An 8-element chunk advances the scale index by + # chunk_id * 8 // group_size. + scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage]) + scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) + scale1_u32 = Uint32( + s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage] + ) + scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) + + # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 + # (CUDA 13.2). + dequant0 = cute.make_rmem_tensor(4, Uint32) + dequant1 = cute.make_rmem_tensor(4, Uint32) + for j in cutlass.range_constexpr(2): + tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) + tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) + + # BF16 multiply is safe because the scales are exact powers of 2. + dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) + dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) + dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) + dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) + + # Last 64 elems are BF16 tail, corresponds to dequant1 of last + # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + + # 8x 16B for BF16. + if lane_id + 32 >= self.fp8_dim // 8: + idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) + + # Store two 16B BF16 chunks per lane: first half, then second half. + dst = out_slice[req_id, offset + i, (None, lane_id)] + cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) + + dst = out_slice[req_id, offset + i, (None, lane_id + 32)] + cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) + + compute_stage = (compute_stage + 1) % self.num_stages + + @cache + @staticmethod + def compile( + fp8_dim: int = 448, + block_size: int = 64, + has_gather_lens: bool = True, + ): + num_reqs = cute.sym_int() + head_dim = DequantGatherKCacheKernel.head_dim + head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 + + out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) + k_cache = cute.runtime.make_fake_tensor( + Uint8, + (cute.sym_int(), block_size, head_bytes), + stride=(cute.sym_int64(divisibility=32), head_bytes, 1), + assumed_align=32, + ) + seq_lens = make_fake_tensor(Int32, (num_reqs,)) + gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None + block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int())) + + kernel = DequantGatherKCacheKernel(fp8_dim, block_size) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + out, + k_cache, + seq_lens, + gather_lens, + block_table, + Int32(0), + stream, + options="--enable-tvm-ffi", + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py new file mode 100644 index 00000000..3bd5e1b5 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -0,0 +1,610 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr +from quack.compile_utils import make_fake_tensor + +from vllm.cute_utils import ( + _bf16x2_abs, + _bf16x2_max, + cvt, + recast_val, +) +from vllm.vllm_flash_attn.cute import utils as cute_utils + +# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. +MXFP4_BLOCK_SIZE = 32 + +_TORCH_TO_CUTE = { + torch.bfloat16: BFloat16, + torch.float32: Float32, +} + + +def fused_indexer_q_rope_quant_mxfp4_cutedsl( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_packed: torch.Tensor, + index_q_scale: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + num_tokens, num_heads, head_dim = index_q.shape + rope_dim = index_q_cos_sin_cache.shape[-1] + rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] + + # compile all variants at first invocation + for coarsen in (1, 4): + IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) + + # heuristic + coarsen = 1 if num_tokens < 512 else 4 + compiled = IndexerQMxFp4Kernel.compile( + head_dim, rope_dim, num_heads, rope_type, coarsen + ) + scale = float(index_weights_softmax_scale * index_weights_head_scale) + compiled( + positions, + index_q, + index_q_cos_sin_cache, + index_weights, + index_q_packed, + index_q_scale, + index_weights_out, + scale, + ) + + +def fused_indexer_q_rope_quant_fp8_cutedsl( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_q_fp8: torch.Tensor, + index_weights_out: torch.Tensor, +) -> None: + num_tokens, num_heads, head_dim = index_q.shape + rope_dim = index_q_cos_sin_cache.shape[-1] + rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] + + for coarsen in (1, 4): + IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) + + coarsen = 1 if num_tokens < 512 else 4 + compiled = IndexerQFp8Kernel.compile( + head_dim, rope_dim, num_heads, rope_type, coarsen + ) + scale = float(index_weights_softmax_scale * index_weights_head_scale) + # The cute kernel treats the FP8 buffer as raw bytes (Uint8). + compiled( + positions, + index_q, + index_q_cos_sin_cache, + index_weights, + index_q_fp8.view(torch.uint8), + index_weights_out, + scale, + ) + + +class IndexerQRopeQuantKernel: + """Shared infrastructure for indexer-Q RoPE+quant fused kernels. + + Subclasses implement ``kernel`` for a particular Q quantization scheme + (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the + common preamble: thread/token addressing, the BF16 Q load, and the + interleaved-RoPE pass over the trailing ``rope_dim`` lanes. + """ + + def __init__( + self, + head_dim: int = 128, + rope_dim: int = 64, + num_heads: int = 64, + cos_sin_dtype: type[cutlass.Numeric] = Float32, + coarsen: int = 4, + ): + self.head_dim = head_dim + self.rope_dim = rope_dim + self.nope_dim = head_dim - rope_dim + self.num_heads = num_heads + self.cos_sin_dtype = cos_sin_dtype + + # process multiple heads at the same time to armotize RoPE load costs + assert num_heads % coarsen == 0 + self.coarsen = coarsen + + # later we will use 32B load = 16 BF16 elems + # thus, head_dim=128 requires 8 threads to handle. + # let's call subwarp = 8 threads. + self.subwarp_size = head_dim // 16 + self.tb_size = 128 + self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size + + @cute.jit + def _load_q_and_rope( + self, + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + ): + """Compute thread indices, load Q (BF16), and apply interleaved RoPE. + + Returns a tuple + (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, + head_start, in_bounds, num_token_heads) + where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed + bf16x2 pairs covering the 16 BF16 lanes owned by this thread for + each of ``coarsen`` heads. RoPE is applied in place to the + trailing ``rope_dim`` lanes; the leading nope lanes pass through. + """ + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + + num_tokens = q.shape[0] + num_token_heads = num_tokens * self.num_heads + global_tid = block_id * self.tb_size + tid + + global_subwarp_id = global_tid // self.subwarp_size + sublane = tid % self.subwarp_size + + token_id = global_subwarp_id // (self.num_heads // self.coarsen) + head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen) + head_start = head_tile_id * self.coarsen + + # NOTE: token_id may exceed bounds, hence we need to add load/store guards + # we can't do early exit because CuteDSL doesn't support it. and we also need + # all threads in a warp to be active since we utilize warp shuffle later. + # must_in_bounds is constexpr, True when 1 threadblock fit within 1 token + # position. the compiler will remove bounds check when that happens. + must_in_bounds = cutlass.const_expr(self.tb_size % self.threads_per_token == 0) + in_bounds = must_in_bounds or (token_id < num_tokens) + + cp_op = cute.nvgpu.CopyUniversalOp() + + _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1)) + q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) + + if in_bounds: + # we can't do cute.copy() on the whole 2D tile directly because + # cute.copy() wants the 1st mode to be covered by the copy atom, + # and other modes as for loop. there is no fast way to + # "transpose" the tensor view. + q_tile = cute.local_tile( + q[token_id, None, None], + tiler=(self.coarsen, 16), + coord=(head_tile_id, sublane), + ) + cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) + for i in cutlass.range_constexpr(self.coarsen): + src = cute.recast_tensor(q_tile[i, None], Uint32) + cute.copy(cp_u32x8, src, q_bf16x2[i, None]) + + # RoPE applies only to the trailing rope_dim values. We keep the rounded + # BF16 result in q_bits so the later amax and quantization see BF16. + # cos_sin_cache layout: [max_pos, rope_dim] + if in_bounds and sublane * 16 >= self.nope_dim: + cos_vals = cute.make_rmem_tensor((8,), Float32) + sin_vals = cute.make_rmem_tensor((8,), Float32) + + pos = positions[token_id] + + # select 8 elems from cos and sin + cos_id = sublane - self.nope_dim // 16 + sin_id = cos_id + self.rope_dim // 16 + cos_src = cute.local_tile( + cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,) + ) + sin_src = cute.local_tile( + cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,) + ) + + cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + if const_expr(self.cos_sin_dtype is Float32): + cute.copy(cp_f32x8, cos_src, cos_vals) + cute.copy(cp_f32x8, sin_src, sin_vals) + else: + cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) + cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) + + for i in cutlass.range_constexpr(4): + cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) + sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) + cos_vals[i * 2] = cos0 + cos_vals[i * 2 + 1] = cos1 + sin_vals[i * 2] = sin0 + sin_vals[i * 2 + 1] = sin1 + + for i in cutlass.range_constexpr(self.coarsen): + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] + rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] + # convert back to BF16 to match numerics + q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) + + return ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, + ) + + +class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel): + """Eight-thread subwarps process one ``(token, head)`` row.""" + + @cute.jit + def __call__( + self, + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * self.threads_per_token + grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) + self.kernel( + positions, + q, + cos_sin_cache, + weights, + q_quant, + q_scale, + weights_out, + scale, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.kernel + def kernel( + self, + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, + ) = self._load_q_and_rope(positions, q, cos_sin_cache) + + cp_op = cute.nvgpu.CopyUniversalOp() + + # layout: [coarsen, 8] + q_fp4_tile = cute.local_tile( + q_quant[token_id, None, None], + tiler=(self.coarsen, 8), + coord=(head_tile_id, sublane), + ) + + for i in cutlass.range_constexpr(self.coarsen): + # compute amax in packed bf16x2 to save instructions + # Each thread holds 16 elems. Two adjacent threads form one 32-elem + # MXFP4 block, so a width-2 shuffle gives the block amax. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=MXFP4_BLOCK_SIZE // 16, + ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) + + if in_bounds: + # compute block scale with bit manipulation + # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask + # increments the exponent whenever fp4_scale is not exactly a power of 2 + eps = cutlass.const_expr(float.fromhex("0x6p-126")) + fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) + bits = recast_val(fp4_scale, Uint32) + ue8m0 = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) + + # Only one of the two threads in an MXFP4 block writes the shared scale. + if tid % 2 == 0: + mx_block = sublane // 2 + q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) + + # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent + # -A + 127 = 254 - ue8m0. + inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) + inv_fp4_scale = recast_val(inv_scale_bits, Float32) + + vals = cute.make_rmem_tensor(16, Float32) + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + vals[j * 2] = q0 * inv_fp4_scale + vals[j * 2 + 1] = q1 * inv_fp4_scale + + # pack to FP4 + packed = cute.make_rmem_tensor((2,), Uint32) + packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) + packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) + + dst = q_fp4_tile[i, None] + cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) + + # Weight scaling is independent of the Q subwarp work. The first + # num_tokens * num_heads logical threads cover one weight each. + if global_tid < num_token_heads: + weight_token_id = global_tid // self.num_heads + weight_head_id = global_tid % self.num_heads + weights_out[weight_token_id, weight_head_id] = ( + weights[weight_token_id, weight_head_id].to(Float32) * scale + ) + + @cache + @staticmethod + def compile( + head_dim: int = 128, + rope_dim: int = 64, + num_heads: int = 64, + cos_sin_dtype: type[cutlass.Numeric] = Float32, + coarsen: int = 4, + ): + num_tokens = cute.sym_int() + max_pos = cute.sym_int() + + q = make_fake_tensor( + BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + ) + positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) + cos_sin_cache = make_fake_tensor( + cos_sin_dtype, + (max_pos, rope_dim), + divisibility=8, + ) + weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) + q_fp4 = make_fake_tensor( + Uint8, + (num_tokens, num_heads, head_dim // 2), + divisibility=16, + ) + q_scale = make_fake_tensor( + Uint8, + (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE), + divisibility=4, + ) + weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) + + kernel = IndexerQMxFp4Kernel( + head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + positions, + q, + cos_sin_cache, + weights, + q_fp4, + q_scale, + weights_out, + Float32(0.0), + stream, + options="--enable-tvm-ffi", + ) + + +class IndexerQFp8Kernel(IndexerQRopeQuantKernel): + """Eight-thread subwarps process one ``(token, head)`` row and emit + float8 e4m3fn with a single per-(token, head) scalar scale folded + into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``). + """ + + def __init__( + self, + head_dim: int = 128, + rope_dim: int = 64, + num_heads: int = 64, + cos_sin_dtype: type[cutlass.Numeric] = Float32, + coarsen: int = 4, + ): + super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen) + # Each subwarp owns `coarsen` heads; we use the first `coarsen` + # threads of the subwarp to write the per-head weights using the + # fp8 scale computed in the matching loop iteration. + assert self.coarsen <= self.subwarp_size, ( + f"FP8 kernel requires coarsen ({self.coarsen}) <= " + f"subwarp_size ({self.subwarp_size}) for the weight-fold step" + ) + + @cute.jit + def __call__( + self, + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * self.threads_per_token + grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) + self.kernel( + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + scale, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.kernel + def kernel( + self, + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + _tid, + _global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + _num_token_heads, + ) = self._load_q_and_rope(positions, q, cos_sin_cache) + + cp_op = cute.nvgpu.CopyUniversalOp() + + # layout: [coarsen, 16] bytes (one e4m3fn per element). + q_fp8_tile = cute.local_tile( + q_fp8[token_id, None, None], + tiler=(self.coarsen, 16), + coord=(head_tile_id, sublane), + ) + + for i in cutlass.range_constexpr(self.coarsen): + # Reduce amax across the full head_dim: each thread already holds + # the max over its 16 lanes; a width=subwarp_size warp shuffle + # spreads the head-wide max to every lane in the subwarp. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=self.subwarp_size, + ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) + + # scale = max(amax, eps) / fp8_max, then rounded UP to the next + # power of two. Adding the mantissa mask before shifting out the + # mantissa bumps the exponent whenever s isn't a pure pow2. + fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) + bits = recast_val(fp32_scale, Uint32) + scale_exp = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) + + # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 + fp8_scale_bits = scale_exp << Uint32(23) + fp8_scale = recast_val(fp8_scale_bits, Float32) + # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 + inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) + inv_fp8_scale = recast_val(inv_scale_bits, Float32) + + # Weight fold: weights_out = weights * q_scale * scale_combined. + # All threads in the subwarp share the same fp8_scale after the + # warp_reduce above, so we let thread `sublane == i` write the + # weight for head `head_start + i`. + if in_bounds and sublane == i: + head_id = head_start + i + weights_out[token_id, head_id] = ( + weights[token_id, head_id].to(Float32) * scale * fp8_scale + ) + + if in_bounds: + # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s + # (one cp.async-shaped 128-bit store per row). + packed = cute.make_rmem_tensor((4,), Uint32) + for j in cutlass.range_constexpr(4): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) + q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) + packed[j] = cvt.fp32x4_to_fp8x4( + q0 * inv_fp8_scale, + q1 * inv_fp8_scale, + q2 * inv_fp8_scale, + q3 * inv_fp8_scale, + ) + + dst = q_fp8_tile[i, None] + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) + + @cache + @staticmethod + def compile( + head_dim: int = 128, + rope_dim: int = 64, + num_heads: int = 64, + cos_sin_dtype: type[cutlass.Numeric] = Float32, + coarsen: int = 4, + ): + num_tokens = cute.sym_int() + max_pos = cute.sym_int() + + q = make_fake_tensor( + BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + ) + positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) + cos_sin_cache = make_fake_tensor( + cos_sin_dtype, + (max_pos, rope_dim), + divisibility=8, + ) + weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) + q_fp8 = make_fake_tensor( + Uint8, + (num_tokens, num_heads, head_dim), + divisibility=16, + ) + weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) + + kernel = IndexerQFp8Kernel( + head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + Float32(0.0), + stream, + options="--enable-tvm-ffi", + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py new file mode 100644 index 00000000..7cdb39e9 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton input-staging kernel for DeepSeek V4 MegaMoE. + +Quantizes hidden states to fp8 with E8M0 group scales and repacks the +routing top-k tensors into the int64/float32 layout that the DeepGEMM +MegaMoE kernels consume. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _prepare_megamoe_inputs_kernel( + hidden_states, + x_fp8, + x_sf, + topk_ids, + topk_weights, + topk_idx_out, + topk_weights_out, + hidden_stride_m: tl.constexpr, + hidden_stride_k: tl.constexpr, + x_stride_m: tl.constexpr, + x_stride_k: tl.constexpr, + x_sf_stride_m: tl.constexpr, + x_sf_stride_k: tl.constexpr, + topk_ids_stride_m: tl.constexpr, + topk_ids_stride_k: tl.constexpr, + topk_weights_stride_m: tl.constexpr, + topk_weights_stride_k: tl.constexpr, + topk_idx_stride_m: tl.constexpr, + topk_idx_stride_k: tl.constexpr, + topk_weights_out_stride_m: tl.constexpr, + topk_weights_out_stride_k: tl.constexpr, + hidden_size: tl.constexpr, + top_k: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_K: tl.constexpr, + BLOCK_TOPK: tl.constexpr, +) -> None: + token_id = tl.program_id(0) + k_block_id = tl.program_id(1) + + k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K) + k_mask = k_offsets < hidden_size + hidden = tl.load( + hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k, + mask=k_mask, + other=0.0, + ).to(tl.float32) + + num_groups: tl.constexpr = BLOCK_K // GROUP_K + hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K]) + amax = tl.max(hidden_groups, axis=1) + amax = tl.maximum(amax, 1.0e-4) + + scale = amax / 448.0 + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + + hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K]) + scaled = hidden_groups * (1.0 / rounded_scale)[:, None] + scaled = tl.reshape(scaled, [BLOCK_K]) + fp8 = scaled.to(tl.float8e4nv) + tl.store( + x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k, + fp8, + mask=k_mask, + ) + + scale_offsets = tl.arange(0, num_groups) + packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + ) + + if k_block_id == 0: + topk_offsets = tl.arange(0, BLOCK_TOPK) + topk_mask = topk_offsets < top_k + + ids = tl.load( + topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k, + mask=topk_mask, + other=0, + ).to(tl.int64) + tl.store( + topk_idx_out + + token_id * topk_idx_stride_m + + topk_offsets * topk_idx_stride_k, + ids, + mask=topk_mask, + ) + + weights = tl.load( + topk_weights + + token_id * topk_weights_stride_m + + topk_offsets * topk_weights_stride_k, + mask=topk_mask, + other=0.0, + ) + tl.store( + topk_weights_out + + token_id * topk_weights_out_stride_m + + topk_offsets * topk_weights_out_stride_k, + weights, + mask=topk_mask, + ) + + +def prepare_megamoe_inputs( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + x_fp8: torch.Tensor, + x_sf: torch.Tensor, + topk_idx_out: torch.Tensor, + topk_weights_out: torch.Tensor, +) -> None: + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + if hidden_size % 128 != 0: + raise ValueError( + "DeepSeek V4 MegaMoE input staging requires hidden_size to be " + "a multiple of 128." + ) + top_k = topk_ids.shape[1] + if topk_weights.shape != topk_ids.shape: + raise ValueError( + "DeepSeek V4 MegaMoE input staging requires topk_weights and " + "topk_ids to have the same shape." + ) + + block_k = 128 + grid = (num_tokens, triton.cdiv(hidden_size, block_k)) + block_topk = triton.next_power_of_2(top_k) + _prepare_megamoe_inputs_kernel[grid]( + hidden_states, + x_fp8, + x_sf, + topk_ids, + topk_weights, + topk_idx_out, + topk_weights_out, + hidden_states.stride(0), + hidden_states.stride(1), + x_fp8.stride(0), + x_fp8.stride(1), + x_sf.stride(0), + x_sf.stride(1), + topk_ids.stride(0), + topk_ids.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + topk_idx_out.stride(0), + topk_idx_out.stride(1), + topk_weights_out.stride(0), + topk_weights_out.stride(1), + hidden_size, + top_k, + BLOCK_K=block_k, + GROUP_K=32, + BLOCK_TOPK=block_topk, + num_warps=4, + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py new file mode 100644 index 00000000..ed16ca6d --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -0,0 +1,1380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CuTe DSL sparse-attention compressor for DeepSeek V4. + +The public wrappers provide the C4 fused and C128 split kernels. +""" + +from __future__ import annotations + +from functools import cache +from typing import Any + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op +from quack.compile_utils import make_fake_tensor + +_TORCH_TO_CUTE = { + torch.bfloat16: BFloat16, + torch.float32: Float32, +} + + +@dsl_user_op +def _recast_val(x, dtype, *, loc=None, ip=None): + return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip))) + + +@dsl_user_op +def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32: + out = llvm.inline_asm( + T.i32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + ) + return Uint32(out) + + +@dsl_user_op +def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + out = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [data.ir_value(loc=loc, ip=ip)], + "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n", + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + ) + return ( + Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16: + out = llvm.inline_asm( + T.i16(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", + "=h,f,f", + has_side_effects=False, + is_align_stack=False, + ) + return Uint16(out) + + +class SparseAttnCompressNormRopeStoreC4Kernel: + min_scale = 1.0e-4 + rcp_ln2 = 1.4426950408889634 + + def __init__( + self, + head_size: int, + state_width: int, + rope_head_dim: int, + fp8_max: float, + quant_block: int, + token_stride: int, + scale_dim: int, + compress_ratio: int, + overlap: bool, + ): + self.head_dim = head_size + self.state_width = state_width + self.rope_dim = rope_head_dim + self.nope_dim = head_size - rope_head_dim + self.fp8_max = fp8_max + self.quant_block = quant_block + self.token_stride = token_stride + self.scale_dim = scale_dim + self.elems_per_lane = 8 + self.copy_elems = 4 + self.copy_chunks = self.elems_per_lane // self.copy_elems + self.lanes_per_group = quant_block // self.elems_per_lane + self.groups_per_warp = 32 // self.lanes_per_group + self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1 + self.scale_reduce_offset = self.lanes_per_group // 2 + self.num_warps = (head_size // quant_block) // self.groups_per_warp + self.nope_blocks = self.nope_dim // quant_block + self.tb_size = self.num_warps * 32 + self.compress_ratio = compress_ratio + self.overlap = overlap + self.window = (1 + int(overlap)) * compress_ratio + + @cute.jit + def __call__( + self, + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + self.kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.kernel + def kernel( + self, + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + group_lane = lane_id % self.lanes_per_group + group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group + elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane + + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + ) + has_req_idx = token_idx < token_to_req_indices.shape[0] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = ( + slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) + ) + + if active: + req_idx = token_to_req_indices[token_idx] + start = position - Int64(self.window - 1) + + smem = cutlass.utils.SmemAllocator() + s_block_numbers = smem.allocate_tensor( + Int32, cute.make_layout((self.window,)), byte_alignment=4 + ) + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) + + for row in cutlass.range_constexpr(self.window): + pos = start + Int64(row) + if tid == row: + block_number_i32 = Int32(0) + if pos >= Int64(0): + block_index = pos // block_size + block_number_i32 = block_table[req_idx, block_index] + s_block_numbers[row] = block_number_i32 + cute.arch.sync_threads() + + local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(self.elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (self.copy_chunks, self.copy_elems), + stride=(self.copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) + + for row in cutlass.range_constexpr(self.window): + pos = start + Int64(row) + if pos >= Int64(0): + block_index = pos // block_size + block_offset = pos - block_index * block_size + block_number = s_block_numbers[row].to(Int64) + head_offset = Int64((row // self.compress_ratio) * self.head_dim) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(self.copy_chunks): + copy_elem = const_expr(chunk * self.copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(self.copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(self.copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(self.copy_elems,), + coord=( + col_tile + Int64(self.state_width // self.copy_elems), + ), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) + + for e in cutlass.range_constexpr(self.elems_per_lane): + chunk = const_expr(e // self.copy_elems) + copy_elem = const_expr(e % self.copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = local_product[e] * old_scale + kv * new_scale + local_max[e] = new_max + + x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(self.elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] + + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(self.num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(self.head_dim) + rms_norm_eps, fastmath=True + ) + cute.arch.sync_threads() + + rrms = rrms_shared[0] + for e in cutlass.range_constexpr(self.elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) + + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) + scale_base = ( + page * k_cache.stride[0] + + kv_cache_block_size * Int64(self.token_stride) + + kv_offset * Int64(self.scale_dim) + ) + + if group_idx == self.nope_blocks: + compressed_pos = (position // Int64(self.compress_ratio)) * Int64( + self.compress_ratio + ) + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(self.rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = ( + value_base + + Int64(self.nope_dim) + + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64) + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_absmax = Float32(0.0) + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + q0, q1 = _bf16x2_to_fp32(q_packed) + q[elem] = q0 + q[elem + 1] = q1 + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), + ) + absmax = local_absmax + group_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - self.lanes_per_group) << 8 + | (cute.arch.WARP_SIZE - 1) + ) + for step in cutlass.range_constexpr(self.scale_reduce_steps): + offset = const_expr(self.scale_reduce_offset >> step) + absmax = cute.arch.fmax( + absmax, + cute.arch.shuffle_sync_bfly( + absmax, + offset=offset, + mask_and_clamp=group_mask_and_clamp, + ), + ) + scale_raw = cute.arch.fmax( + Float32(self.min_scale), + absmax / Float32(self.fp8_max), + ) + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + y0 = cutlass.min( + cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if group_lane == 0: + k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8) + if group_idx == 0: + k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( + 0 + ) + + @cache + @staticmethod + def compile( + head_size: int = 512, + state_width: int = 1024, + rope_head_dim: int = 64, + fp8_max: float = 448.0, + quant_block: int = 64, + token_stride: int = 576, + scale_dim: int = 8, + kv_block_stride: int = 74752, + compress_ratio: int = 4, + overlap: bool = True, + norm_weight_dtype: type[cutlass.Numeric] = Float32, + ): + if compress_ratio != 4 or not overlap: + raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") + if head_size != 512: + raise ValueError( + "CuTe DSL C4 fused sparse-attn currently requires head_size=512." + ) + if state_width != 2 * head_size: + raise ValueError( + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." + ) + if quant_block != 64: + raise ValueError( + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." + ) + if rope_head_dim != 64: + raise ValueError( + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." + ) + if token_stride < head_size + rope_head_dim: + raise ValueError("token_stride is too small for the packed FP8/BF16 row.") + expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 + if scale_dim < expected_scale_dim: + raise ValueError("scale_dim is too small for the UE8M0 scale row.") + + num_positions = cute.sym_int() + num_slots = cute.sym_int() + num_req_indices = cute.sym_int() + num_kv_slots = cute.sym_int() + num_state_blocks = cute.sym_int() + num_kv_blocks = cute.sym_int() + state_cache_block_size = cute.sym_int() + block_table_width = cute.sym_int() + max_pos = cute.sym_int() + state_cache_width = state_width * 2 + + state_cache = cute.runtime.make_fake_tensor( + Float32, + (num_state_blocks, state_cache_block_size, state_cache_width), + stride=( + cute.sym_int64(divisibility=16), + cute.sym_int64(divisibility=16), + 1, + ), + assumed_align=16, + ) + token_to_req_indices = make_fake_tensor( + Int32, (num_req_indices,), divisibility=4 + ) + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) + block_table = make_fake_tensor( + Int32, (cute.sym_int(), block_table_width), divisibility=1 + ) + rms_norm_weight = make_fake_tensor( + norm_weight_dtype, (head_size,), divisibility=4 + ) + cos_sin_cache = cute.runtime.make_fake_tensor( + Float32, + (max_pos, rope_head_dim), + stride=(cute.sym_int64(divisibility=4), 1), + assumed_align=4, + ) + k_cache = cute.runtime.make_fake_tensor( + Uint8, + (num_kv_blocks, cute.sym_int(), cute.sym_int()), + stride=( + cute.sym_int64(divisibility=16), + cute.sym_int64(divisibility=8), + 1, + ), + assumed_align=16, + ) + kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) + + kernel = SparseAttnCompressNormRopeStoreC4Kernel( + head_size, + state_width, + rope_head_dim, + fp8_max, + quant_block, + token_stride, + scale_dim, + compress_ratio, + overlap, + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + Int64(0), + rms_norm_weight, + Float32(0.0), + cos_sin_cache, + k_cache, + kv_slot_mapping, + Int64(0), + stream, + options="--enable-tvm-ffi", + ) + + +class SparseAttnCompressKernel: + head_tile = 64 + rows_per_warp = 16 + row_pairs_per_warp = rows_per_warp // 2 + elems_per_lane = 4 + lanes_per_row = head_tile // elems_per_lane + num_warps = 8 + stats_warp_stride = num_warps + 1 + tb_size = num_warps * 32 + rcp_ln2 = 1.4426950408889634 + + def __init__( + self, + head_size: int, + state_width: int, + compress_ratio: int, + overlap: bool, + ): + self.head_dim = head_size + self.num_splits = head_size // self.head_tile + self.state_width = state_width + self.compress_ratio = compress_ratio + self.overlap = overlap + self.window = (1 + int(overlap)) * compress_ratio + + @cute.jit + def __call__( + self, + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + compressed_kv: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0] * self.num_splits, 1, 1) + self.kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + compressed_kv, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.kernel + def kernel( + self, + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + compressed_kv: cute.Tensor, + ): + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + row_lane = lane_id // self.lanes_per_row + col_group = lane_id % self.lanes_per_row + + token_idx = block_id // self.num_splits + split_idx = block_id - token_idx * self.num_splits + col_base = split_idx * self.head_tile + col_group * self.elems_per_lane + + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + ) + has_req_idx = token_idx < token_to_req_indices.shape[0] + active = slot_id >= Int64(0) and has_req_idx and boundary + + if active: + smem = cutlass.utils.SmemAllocator() + s_max = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + self.lanes_per_row, + self.elems_per_lane, + self.stats_warp_stride, + ), + stride=( + self.elems_per_lane * self.stats_warp_stride, + self.stats_warp_stride, + 1, + ), + ), + byte_alignment=4, + ) + s_sum = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + self.lanes_per_row, + self.elems_per_lane, + self.stats_warp_stride, + ), + stride=( + self.elems_per_lane * self.stats_warp_stride, + self.stats_warp_stride, + 1, + ), + ), + byte_alignment=4, + ) + s_product = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + self.lanes_per_row, + self.elems_per_lane, + self.stats_warp_stride, + ), + stride=( + self.elems_per_lane * self.stats_warp_stride, + self.stats_warp_stride, + 1, + ), + ), + byte_alignment=4, + ) + + row_pair_layout = cute.make_layout( + (self.row_pairs_per_warp, self.elems_per_lane), + stride=(self.elems_per_lane, 1), + ) + kv_vals = cute.make_rmem_tensor(row_pair_layout, Float32) + score_vals = cute.make_rmem_tensor(row_pair_layout, Float32) + local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(self.elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + req_idx = token_to_req_indices[token_idx] + start = position - Int64(self.window - 1) + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + row_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - self.lanes_per_row) << 8 + | (cute.arch.WARP_SIZE - 1) + ) + + for i in cutlass.range_constexpr(self.row_pairs_per_warp): + row = warp_id * self.rows_per_warp + i * 2 + row_lane + pos = start + row.to(Int64) + valid = row < self.window and pos >= Int64(0) + head_offset = ((row // self.compress_ratio) * self.head_dim).to(Int64) + + for e in cutlass.range_constexpr(self.elems_per_lane): + kv = Float32(0.0) + score = -Float32.inf + kv_vals[i, e] = kv + score_vals[i, e] = score + + block_index = Int64(0) + block_offset = Int64(0) + block_number_i32 = Int32(0) + if valid: + block_index = pos // block_size + block_offset = pos - block_index * block_size + if col_group == 0: + block_number_i32 = block_table[req_idx, block_index] + block_number_i32 = cute.arch.shuffle_sync( + block_number_i32, + offset=0, + mask_and_clamp=row_mask_and_clamp, + ) + + if valid: + block_number = block_number_i32.to(Int64) + row_tensor = state_cache[block_number, block_offset, None] + col_tile = (head_offset + col_base.to(Int64)) // Int64( + self.elems_per_lane + ) + kv_src = cute.local_tile( + row_tensor, + tiler=(self.elems_per_lane,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(self.elems_per_lane,), + coord=( + col_tile + Int64(self.state_width // self.elems_per_lane), + ), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[i, None]) + cute.copy(cp_f32x4, score_src, score_vals[i, None]) + + for e in cutlass.range_constexpr(self.elems_per_lane): + local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) + + for e in cutlass.range_constexpr(self.elems_per_lane): + if local_max[e] > -Float32.inf: + for i in cutlass.range_constexpr(self.row_pairs_per_warp): + exp_score = cute.math.exp2( + (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2), + fastmath=True, + ) + local_sum[e] += exp_score + local_product[e] += kv_vals[i, e] * exp_score + + for e in cutlass.range_constexpr(self.elems_per_lane): + pair_max = cute.arch.shuffle_sync_bfly(local_max[e], offset=16) + pair_sum = cute.arch.shuffle_sync_bfly(local_sum[e], offset=16) + pair_product = cute.arch.shuffle_sync_bfly(local_product[e], offset=16) + warp_max = cute.arch.fmax(local_max[e], pair_max) + warp_sum = Float32(0.0) + warp_product = Float32(0.0) + if warp_max > -Float32.inf: + local_scale = cute.math.exp2( + (local_max[e] - warp_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + pair_scale = cute.math.exp2( + (pair_max - warp_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + warp_sum = local_sum[e] * local_scale + pair_sum * pair_scale + warp_product = ( + local_product[e] * local_scale + pair_product * pair_scale + ) + if lane_id < self.lanes_per_row: + s_max[col_group, e, warp_id] = warp_max + s_sum[col_group, e, warp_id] = warp_sum + s_product[col_group, e, warp_id] = warp_product + cute.arch.sync_threads() + + out_group = tid // self.num_warps + final_lane = tid % self.num_warps + final_groups_per_pass = const_expr(self.tb_size // self.num_warps) + for pass_idx in cutlass.range_constexpr( + self.head_tile // final_groups_per_pass + ): + out_idx = pass_idx * final_groups_per_pass + out_group + out_lane = out_idx // self.elems_per_lane + out_elem = out_idx % self.elems_per_lane + + local_warp_max = s_max[out_lane, out_elem, final_lane] + global_max = local_warp_max + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) + global_max = cute.arch.fmax( + global_max, + cute.arch.shuffle_sync_bfly( + global_max, + offset=offset, + mask_and_clamp=row_mask_and_clamp, + ), + ) + + scale = cute.math.exp2( + (local_warp_max - global_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + global_sum = s_sum[out_lane, out_elem, final_lane] * scale + global_product = s_product[out_lane, out_elem, final_lane] * scale + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) + global_sum += cute.arch.shuffle_sync_bfly( + global_sum, + offset=offset, + mask_and_clamp=row_mask_and_clamp, + ) + global_product += cute.arch.shuffle_sync_bfly( + global_product, + offset=offset, + mask_and_clamp=row_mask_and_clamp, + ) + + if final_lane == 0: + compressed_kv.iterator[ + token_idx.to(Int64) * compressed_kv.stride[0] + + (split_idx * self.head_tile + out_idx).to(Int64) + ] = global_product / global_sum + + @cache + @staticmethod + def compile( + head_size: int = 512, + state_width: int = 512, + compress_ratio: int = 128, + overlap: bool = False, + ): + if head_size % SparseAttnCompressKernel.head_tile != 0: + raise ValueError("head_size must be divisible by the 64-wide head tile.") + num_positions = cute.sym_int() + num_slots = cute.sym_int() + num_req_indices = cute.sym_int() + num_blocks = cute.sym_int() + state_cache_block_size = cute.sym_int() + block_table_width = cute.sym_int() + state_cache_width = state_width * 2 + + state_cache = cute.runtime.make_fake_tensor( + Float32, + (num_blocks, state_cache_block_size, state_cache_width), + stride=( + cute.sym_int64(divisibility=16), + cute.sym_int64(divisibility=16), + 1, + ), + assumed_align=16, + ) + token_to_req_indices = make_fake_tensor( + Int32, (num_req_indices,), divisibility=4 + ) + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) + block_table = make_fake_tensor( + Int32, (cute.sym_int(), block_table_width), divisibility=1 + ) + compressed_kv = cute.runtime.make_fake_tensor( + Float32, + (num_slots, head_size), + stride=(cute.sym_int64(divisibility=4), 1), + assumed_align=4, + ) + + kernel = SparseAttnCompressKernel( + head_size, + state_width, + compress_ratio, + overlap, + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + Int64(0), + compressed_kv, + stream, + options="--enable-tvm-ffi", + ) + + +class SparseAttnNormRopeStoreKernel: + min_scale = 1.0e-4 + + def __init__( + self, + head_size: int, + rope_head_dim: int, + fp8_max: float, + quant_block: int, + token_stride: int, + scale_dim: int, + compress_ratio: int, + ): + self.head_dim = head_size + self.rope_dim = rope_head_dim + self.nope_dim = head_size - rope_head_dim + self.fp8_max = fp8_max + self.quant_block = quant_block + self.token_stride = token_stride + self.scale_dim = scale_dim + self.num_warps = head_size // quant_block + self.nope_blocks = self.nope_dim // quant_block + self.tb_size = head_size // 2 + self.compress_ratio = compress_ratio + + @cute.jit + def __call__( + self, + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + self.kernel( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + + @cute.kernel + def kernel( + self, + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + elem0 = tid * 2 + + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + ) + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + + if active: + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) + x0 = compressed_kv.iterator[base] + x1 = compressed_kv.iterator[base + Int64(1)] + + local_sumsq = x0 * x0 + x1 * x1 + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + + smem = cutlass.utils.SmemAllocator() + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(self.num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(self.head_dim) + rms_norm_eps, fastmath=True + ) + cute.arch.sync_threads() + + rrms = rrms_shared[0] + x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) + x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) + + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) + scale_base = ( + page * k_cache.stride[0] + + kv_cache_block_size * Int64(self.token_stride) + + kv_offset * Int64(self.scale_dim) + ) + + if warp_id == self.nope_blocks: + pair_idx = lane_id + compressed_pos = (position // Int64(self.compress_ratio)) * Int64( + self.compress_ratio + ) + cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] + real = x0 * cos_v - x1 * sin_v + imag = x0 * sin_v + x1 * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64) + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q_packed = _fp32x2_to_bf16x2(x0, x1) + q0, q1 = _bf16x2_to_fp32(q_packed) + abs0 = cute.math.absf(q0) + abs1 = cute.math.absf(q1) + local_absmax = cute.arch.fmax(abs0, abs1) + absmax = local_absmax + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + absmax = cute.arch.fmax( + absmax, cute.arch.shuffle_sync_bfly(absmax, offset) + ) + scale_raw = cute.arch.fmax( + Float32(self.min_scale), + absmax / Float32(self.fp8_max), + ) + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) + y0 = cutlass.min( + cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to( + Int64 + ) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if lane_id == 0: + k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8) + if warp_id == 0: + k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( + 0 + ) + + @cache + @staticmethod + def compile( + head_size: int = 512, + rope_head_dim: int = 64, + fp8_max: float = 448.0, + quant_block: int = 64, + token_stride: int = 576, + scale_dim: int = 8, + kv_block_stride: int = 74752, + compress_ratio: int = 128, + norm_weight_dtype: type[cutlass.Numeric] = Float32, + ): + if quant_block != 64: + raise ValueError( + "CuTe DSL sparse-attn store currently requires quant_block=64." + ) + if rope_head_dim != 64: + raise ValueError( + "CuTe DSL sparse-attn store currently requires rope_head_dim=64." + ) + if head_size % quant_block != 0: + raise ValueError("head_size must be divisible by quant_block.") + if token_stride < head_size + rope_head_dim: + raise ValueError("token_stride is too small for the packed FP8/BF16 row.") + expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 + if scale_dim < expected_scale_dim: + raise ValueError("scale_dim is too small for the UE8M0 scale row.") + num_positions = cute.sym_int() + num_slots = cute.sym_int() + num_kv_slots = cute.sym_int() + max_pos = cute.sym_int() + num_blocks = cute.sym_int() + + compressed_kv = cute.runtime.make_fake_tensor( + Float32, + (num_slots, head_size), + stride=(cute.sym_int64(divisibility=4), 1), + assumed_align=4, + ) + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) + rms_norm_weight = make_fake_tensor( + norm_weight_dtype, (head_size,), divisibility=4 + ) + cos_sin_cache = cute.runtime.make_fake_tensor( + Float32, + (max_pos, rope_head_dim), + stride=(cute.sym_int64(divisibility=4), 1), + assumed_align=4, + ) + k_cache = cute.runtime.make_fake_tensor( + Uint8, + (num_blocks, cute.sym_int(), cute.sym_int()), + stride=( + cute.sym_int64(divisibility=16), + cute.sym_int64(divisibility=8), + 1, + ), + assumed_align=16, + ) + kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) + + kernel = SparseAttnNormRopeStoreKernel( + head_size, + rope_head_dim, + fp8_max, + quant_block, + token_stride, + scale_dim, + compress_ratio, + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + Float32(0.0), + cos_sin_cache, + k_cache, + kv_slot_mapping, + Int64(0), + stream, + options="--enable-tvm-ffi", + ) + + +def compress_kv_sparse_attn_cutedsl( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + compressed_kv: torch.Tensor, + head_size: int = 512, + state_width: int = 512, + compress_ratio: int = 128, + overlap: bool = False, +) -> None: + if positions.numel() == 0: + return + compiled = SparseAttnCompressKernel.compile( + head_size=head_size, + state_width=state_width, + compress_ratio=compress_ratio, + overlap=overlap, + ) + compiled( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + compressed_kv, + ) + + +def norm_rope_insert_sparse_attn_cutedsl( + compressed_kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + k_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + kv_block_stride: int, + head_size: int = 512, + rope_head_dim: int = 64, + fp8_max: float = 448.0, + quant_block: int = 64, + token_stride: int = 576, + scale_dim: int = 8, + compress_ratio: int = 128, +) -> None: + if positions.numel() == 0: + return + norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) + if norm_weight_dtype is None: + raise ValueError( + "CuTe DSL sparse-attn store supports rms_norm_weight dtype " + f"bf16/fp32, got {rms_norm_weight.dtype}." + ) + if k_cache.ndim != 3: + raise ValueError( + "CuTe DSL sparse-attn store expects the real DeepSeek V4 " + f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." + ) + compiled = SparseAttnNormRopeStoreKernel.compile( + head_size=head_size, + rope_head_dim=rope_head_dim, + fp8_max=fp8_max, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + kv_block_stride=kv_block_stride, + compress_ratio=compress_ratio, + norm_weight_dtype=norm_weight_dtype, + ) + compiled( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ) + + +def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + k_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + kv_block_stride: int, + head_size: int = 512, + state_width: int = 1024, + rope_head_dim: int = 64, + fp8_max: float = 448.0, + quant_block: int = 64, + token_stride: int = 576, + scale_dim: int = 8, + compress_ratio: int = 4, + overlap: bool = True, +) -> None: + if positions.numel() == 0: + return + norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) + if norm_weight_dtype is None: + raise ValueError( + "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype " + f"bf16/fp32, got {rms_norm_weight.dtype}." + ) + if k_cache.ndim != 3: + raise ValueError( + "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " + f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." + ) + compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile( + head_size=head_size, + state_width=state_width, + rope_head_dim=rope_head_dim, + fp8_max=fp8_max, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + kv_block_stride=kv_block_stride, + compress_ratio=compress_ratio, + overlap=overlap, + norm_weight_dtype=norm_weight_dtype, + ) + compiled( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ) + + +def compress_norm_rope_store_cutedsl( + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, +) -> None: + if compress_ratio == 4: + # For C4A, the single fused kernel is faster than the two-kernel version. + fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache, + k_cache_metadata.slot_mapping, + kv_cache.shape[1], # paged KV cache block size + kv_cache.stride(0), + head_size=head_dim, + state_width=state_width, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + compress_ratio=compress_ratio, + overlap=overlap, + ) + else: + # For C128, the two-kernel version is faster than the single fused kernel. + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) + compress_kv_sparse_attn_cutedsl( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + compressed_kv, + head_size=head_dim, + state_width=state_width, + compress_ratio=compress_ratio, + overlap=overlap, + ) + norm_rope_insert_sparse_attn_cutedsl( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache, + k_cache_metadata.slot_mapping, + kv_cache.shape[1], # paged KV cache block size + kv_cache.stride(0), + head_size=head_dim, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + compress_ratio=compress_ratio, + ) diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py b/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py new file mode 100644 index 00000000..85a78883 --- /dev/null +++ b/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantization config for DeepSeek V4.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) + +_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ) + + +class DeepseekV4FP8Config(Fp8Config): + """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch. + + DeepSeek V4 checkpoints always use FP8 block quantization for + linear/attention layers. The MoE expert weights vary by checkpoint: + - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts + with ue8m0 (e8m0fnu) FP8 linear scales. + - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block + experts with float32 FP8 linear scales. + + The dispatch and the linear scale dtype are both keyed off + ``expert_dtype`` from the model's hf_config; missing values default + to ``"fp4"`` so existing FP4 checkpoints stay unchanged. + + NOTE: ``expert_dtype`` is resolved lazily because this config is + constructed during VllmConfig setup, before ``set_current_vllm_config`` + is active. Reading hf_config eagerly in ``__init__`` would always see + the default ``"fp4"`` and silently misroute Flash-Base checkpoints. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._resolved_expert_dtype: str | None = None + self._resolved_moe_quant_algo: str | None = None + self._nvfp4_config: ModelOptNvFp4Config | None = None + # ``is_scale_e8m0`` is a property that resolves on first read, + # by which time the current vllm_config has been set. + + @property + def expert_dtype(self) -> str: + if self._resolved_expert_dtype is None: + try: + hf_config = get_current_vllm_config().model_config.hf_config + except Exception: + # vllm_config not yet set; defer the decision until a + # later call lands inside set_current_vllm_config. + return "fp4" + expert_dtype = getattr(hf_config, "expert_dtype", "fp4") + if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES: + raise ValueError( + f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; " + f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}." + ) + self._resolved_expert_dtype = expert_dtype + from vllm.logger import init_logger + + init_logger(__name__).info_once( + "DeepSeek V4 expert_dtype resolved to %r", expert_dtype + ) + return self._resolved_expert_dtype + + @property + def is_scale_e8m0(self) -> bool: + # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert + # checkpoints (Flash-Base) store them as float32. + return self.expert_dtype == "fp4" + + def _resolve_moe_overrides(self) -> None: + if self._resolved_moe_quant_algo is not None: + return + try: + hf_config = get_current_vllm_config().model_config.hf_config + except Exception: + return + quant_cfg = getattr(hf_config, "quantization_config", None) or {} + algo = (quant_cfg.get("moe_quant_algo") or "").upper() or None + self._resolved_moe_quant_algo = algo or "" + + @property + def moe_quant_algo(self) -> str: + self._resolve_moe_overrides() + return self._resolved_moe_quant_algo or "" + + def _get_nvfp4_config(self) -> ModelOptNvFp4Config: + if self._nvfp4_config is None: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ) + + self._nvfp4_config = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + group_size=16, + ) + return self._nvfp4_config + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "deepseek_v4_fp8" + + @classmethod + def override_quantization_method( + cls, hf_quant_cfg, user_quant, hf_config=None + ) -> QuantizationMethods | None: + if not ( + isinstance(hf_quant_cfg, dict) + and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") + ): + return None + model_type = getattr(hf_config, "model_type", None) + if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8": + return "deepseek_v4_fp8" + return None + + def get_quant_method(self, layer, prefix): + if isinstance(layer, FusedMoE): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedFusedMoEMethod(layer.moe_config) + if self.expert_dtype == "fp4": + if self.moe_quant_algo == "NVFP4": + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4FusedMoE, + ) + + return ModelOptNvFp4FusedMoE( + quant_config=self._get_nvfp4_config(), + moe_config=layer.moe_config, + ) + return Mxfp4MoEMethod(layer.moe_config) + # expert_dtype == "fp8": fall through to Fp8Config which + # returns Fp8MoEMethod with block-wise float32 scales. + return super().get_quant_method(layer, prefix) + + def is_mxfp4_quant(self, prefix, layer): + if not isinstance(layer, FusedMoE) or self.expert_dtype != "fp4": + return False + return self.moe_quant_algo != "NVFP4" diff --git a/single_shot_inference.py b/single_shot_inference.py index f416d99f..29d841a0 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -365,8 +365,10 @@ class Compressor: n_comp = compressed.shape[0] # Vectorized position computation — no Python loop, no .item() + # Block-aligned: use FIRST position of each block (vLLM cross-check confirmed) + # Wrong: ((bi+1)*r - 1) uses LAST position → off by r-1 (3 for CSA, 127 for HCA) bi = torch.arange(n_comp, device=dev) - pos_idx = ((bi + 1) * r - 1).clamp(max=positions.numel() - 1) + pos_idx = (bi * r).clamp(max=positions.numel() - 1) comp_pos = positions[pos_idx] # Return FP32 compressed output — caller handles RoPE + NVFP4 quantize @@ -990,6 +992,7 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, moe_runner=None, se_runner=None, router=None, prod_lin=None, _profile_detail=False, _profile_times=None, _use_fused_rmsnorm_quantize=True, + comp_rope_cos=None, comp_rope_sin=None, ): """Forward one transformer layer. """ @@ -1406,6 +1409,14 @@ def main(): rtheta = cfg.get("rope_theta", 10000.); romax = rp.get("original_max_position_embeddings", 65536) rbfast, rbslow = rp.get("beta_fast", 32), rp.get("beta_slow", 1) rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)} + # Compressed-entry RoPE uses separate theta (vLLM cross-check: compress_rope_theta) + # If compress_rope_theta differs from rope_theta, compressed KV entries need their own cache + comp_rtheta = cfg.get("compress_rope_theta", rtheta) + if comp_rtheta != rtheta: + comp_rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", comp_rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)} + print(f" Compressed RoPE theta: {comp_rtheta} (different from normal: {rtheta})") + else: + comp_rope_caches = rope_caches # Same theta, reuse normal cache # KV caches, compressors, indexers kv_caches, compressors, indexers = {}, {}, {}