diff --git a/TEMP/CROSS_REFERENCE.md b/TEMP/CROSS_REFERENCE.md deleted file mode 100644 index 7affd75f..00000000 --- a/TEMP/CROSS_REFERENCE.md +++ /dev/null @@ -1,100 +0,0 @@ -# vLLM Independent-Reference Cross-Checks (DSV4) - -**Purpose:** catch the architectural bugs your cosine backlog is blind to. Every check in `/home/openclaw/dev/nvfp4-megamoe-kernel/DEGENERATION_TESTS.md` and `/home/openclaw/dev/nvfp4-megamoe-kernel/CORRECTNESS_BACKLOG.md` compares your kernel to your *own* PyTorch reference — if both share a wrong assumption (RoPE position, RoPE style, norm ordering, compression overlap), they agree at cos 0.9999 while both being wrong. vLLM's DSV4 is an **independent** implementation; it can't share your blind spots. Use it as the tiebreaker. - -**How to use this:** these are mostly *static code diffs*, not GPU runs — fast. For each check, find the line in `single_shot_inference.py` (or `dsv4/`), find the cited line in the vLLM tree, and record match / mismatch + the fix. - -**Ground rules** -- The attached vLLM tree "works but is unstable" and strayed from the paper in places. So trust it as an oracle for the **math** of the compressor / mHC / RoPE / norm-ordering — it produces coherent output against the real weights — but NOT as a stability or paper-fidelity reference everywhere. -- **DO NOT mirror vLLM's attention.** `attention.py` / `nvidia/flashmla.py` shoehorn CSA/HCA onto the V3.2 MLA-sparse backend + a bolted-on SWA cache. That's the unstable stopgap. Your from-scratch shared-KV MQA + grouped-output is more paper-faithful — keep it. -- Where vLLM and your code disagree on a detail your cosines can't see (below), **vLLM is the tiebreaker worth investigating** — match it and re-test. - -Reference files: `deepseek_v4/nvidia/model.py`, `deepseek_v4/compressor.py`, `deepseek_v4/common/rope.py`, `deepseek_v4/common/ops/*`, `deepseek_v4/nvidia/ops/*cutedsl.py`. - ---- - -## Cross-check 1 — Final tail ordering: collapse → final RMSNorm → lm_head -*(Resolves the Test-2 / mHC question. If this is wrong, it explains the |X|=860 confusion.)* - -**vLLM reference (`nvidia/model.py`):** the expanded mHC residual stays `n_hc × d` "until `hc_head()` collapses it" (comment ~line 1130). Tail sequence: -- `mhc_post_tilelang(...)` collapse (`hc_head`, `n_hc·d → d`) — ~line 1172 -- `hidden_states = self.norm(hidden_states)` where `self.norm = RMSNorm(hidden_size)` — ~line 1191 -- `compute_logits` → `lm_head` — ~lines 1444–1448 - -So the canonical tail is **collapse → final RMSNorm → lm_head, in that order.** - -**Check:** trace the last layer's residual in `single_shot`. Does it (a) collapse the expanded mHC residual via `hc_head`/`mhc_post`, then (b) apply a model-final `RMSNorm`, then (c) `lm_head` — in that order? - -**Decision** -- Missing collapse, missing final norm, or reordered → **real bug, fix the ordering.** This is the legitimate "the final norm is broken" branch from the decode runbook — a small fix, NOT residual clipping. -- All present and correct → RMSNorm is scale-invariant, so |X|=860 is normalized away → **mHC growth is not the cause**, corroborating the Test-2 falsification. - ---- - -## Cross-check 2 — Compressed-entry RoPE position (block-aligned) — HIGHEST PRIORITY -*(Off-by-(r−1) here is invisible to same-input cosines and wrecks long-range attention.)* - -**vLLM reference (`compressor.py`):** the position used to RoPE a compressed entry is -``` -(positions // compress_ratio) * compress_ratio # block-aligned = FIRST position of the block -``` - -**Your earlier code used:** `((bi+1) * r - 1)` — the **LAST** position of the block. Difference = `r − 1`: **3 for CSA, 127 for HCA.** A constant 127-position offset on every HCA entry degrades long-range attention while every per-layer cosine stays high (your reference makes the same choice). - -**Check:** what position does `single_shot`'s compressor assign when RoPE-ing compressed KV entries **and** the indexer keys? Compare against the block-aligned floor formula. - -**Decision:** anything other than block-aligned `(pos // ratio) * ratio` → change to it and re-test. This is the single most likely concrete bug in this document. - ---- - -## Cross-check 3 — RoPE style: GPT-J interleaved, not NeoX split-half - -**vLLM reference (`compressor.py` compress kernel comment + `common/rope.py`):** -- `is_neox_style=False` → **interleaved pairs**, NOT split-half -- applied to the **last `rope_head_dim`** elements of `head_dim` -- `cos_sin_cache` layout `[max_pos, rope_head_dim]`: first half cos, second half sin (per-pair, each length `rope_head_dim // 2`) - -**Check:** is `single_shot`'s RoPE interleaved (GPT-J) or split-half (NeoX)? Verify **both** the forward RoPE on Q/KV **and** the inverse RoPE on the attention output, and that the cos/sin cache layout matches. - -**Decision:** split-half where the model expects interleaved (or a mismatched cos/sin layout) scrambles every roped vector — invisible to same-input cosines, wrong against the real weights. Match vLLM's interleaved layout. - ---- - -## Cross-check 4 — CSA overlap (2m) vs HCA no-overlap - -**vLLM reference (`compressor.py`):** `overlap = (compress_ratio == 4)` and `coff = 1 + (compress_ratio == 4)`. So CSA (ratio 4) uses **overlapped** compression — each compressed entry derived from `2m` KV entries via the `C^a` + `C^b` overlap (paper eqs. 11–12) — while HCA (ratio 128) does **not** overlap. - -**Check:** does `single_shot`'s CSA compressor produce overlapped blocks (each entry from `2m` entries with the `C^b` overlap), and HCA non-overlapped? - -**Decision:** if CSA isn't overlapping, the compressed representation differs from what the model was trained on → degraded attention. Match the overlap behavior. - ---- - -## Reporting - -Same harness/workflow as the decode runbook: edit locally → commit → push → pull on B200 → test; `TEST_LAYERS` as an **env var**, never a CLI arg. Most checks here are static diffs (no GPU). For each, paste: `single_shot` line, vLLM reference line, MATCH/MISMATCH, and the fix if mismatched. - ---- - -## Appendix A — vLLM reference → your kernel map (for the kernel work, not the bug) - -Independent CuteDSL/op references for the kernels you hand-wrote. BF16/FP8, not NVFP4 — so a *correctness/layout* reference, not liftable code (and useful evidence for the NVIDIA gaps writeup): - -| Your kernel | vLLM independent reference | -|---|---| -| compress → RMSNorm → RoPE → FP8 quant → cache-write | `nvidia/ops/sparse_attn_compress_cutedsl.py` (`compress_norm_rope_store_cutedsl`) | -| dequant + gather K (your `gather_mixed_*`) | `nvidia/ops/dequant_gather_k_cutedsl.py` | -| indexer Q | `nvidia/ops/fused_indexer_q_cutedsl.py`, `common/ops/fused_indexer_q.py` | -| inverse RoPE + FP8 quant | `common/ops/fused_inv_rope_fp8_quant.py` | -| QK RMSNorm | `common/ops/fused_qk_rmsnorm.py` | -| uncompressed tail / partial-state buffer (your decode buffering) | `common/ops/save_partial_states.py` | -| MoE dispatch prep | `nvidia/ops/prepare_megamoe.py` | -| mHC pre/post | `vllm…kernels.mhc.tilelang` (`mhc_pre`, `mhc_post`, `mhc_fused_post_pre`) — used in `nvidia/model.py` | - -## Appendix B — determinism gotcha worth stealing - -`compressor.py` disables PDL (`launch_pdl=False`) because the compress kernels consume a preceding GEMM's output without emitting/waiting on PDL grid-dependency primitives, and `launch_pdl=True` caused a read-after-write race → non-deterministic output. You target bitwise determinism, so this is a hazard already mapped: any kernel that depends on a prior kernel's output without an explicit grid-dependency primitive must not use PDL. - -## Appendix C — what NOT to take - -`attention.py`, `nvidia/flashmla.py`, and the `DeepseekV4SparseMLA` / `DeepseekV4SWACache` machinery are the MLA shoehorn — the unstable stopgap that motivated this whole project. Use vLLM's compressor / mHC / RoPE / norm-ordering as oracles; leave its attention impl alone. \ No newline at end of file diff --git a/TEMP/deepseek_v4.zip b/TEMP/deepseek_v4.zip deleted file mode 100644 index 77ae49cd..00000000 Binary files a/TEMP/deepseek_v4.zip and /dev/null differ diff --git a/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py b/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py deleted file mode 100644 index abaa794f..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 deleted file mode 100644 index 208f01a7..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# 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 deleted file mode 100644 index fb724fbe..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py +++ /dev/null @@ -1,972 +0,0 @@ -# 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 deleted file mode 100644 index 5938cde6..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py +++ /dev/null @@ -1,509 +0,0 @@ -# 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 deleted file mode 100644 index 2af93fba..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py +++ /dev/null @@ -1,856 +0,0 @@ -# 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 deleted file mode 100644 index 55cb3d94..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/attention.py +++ /dev/null @@ -1,806 +0,0 @@ -# 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 deleted file mode 100644 index 208f01a7..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# 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 deleted file mode 100644 index dc6f3c60..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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 deleted file mode 100644 index ac66751e..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py +++ /dev/null @@ -1,594 +0,0 @@ -# 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 deleted file mode 100644 index 9a5e478e..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ /dev/null @@ -1,666 +0,0 @@ -# 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 deleted file mode 100644 index d5aaf10f..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py +++ /dev/null @@ -1,438 +0,0 @@ -# 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 deleted file mode 100644 index 97fc0962..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ /dev/null @@ -1,318 +0,0 @@ -# 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 deleted file mode 100644 index 874b4f60..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py +++ /dev/null @@ -1,203 +0,0 @@ -# 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 deleted file mode 100644 index 0dd348a4..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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 deleted file mode 100644 index e3d7d38f..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py +++ /dev/null @@ -1,101 +0,0 @@ -# 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 deleted file mode 100644 index 44ae3286..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 deleted file mode 100644 index f36dc8f1..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/compressor.py +++ /dev/null @@ -1,380 +0,0 @@ -# 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 deleted file mode 100644 index 208f01a7..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# 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 deleted file mode 100644 index 5c8b08d4..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py +++ /dev/null @@ -1,424 +0,0 @@ -# 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 deleted file mode 100644 index 547048ab..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py +++ /dev/null @@ -1,1476 +0,0 @@ -# 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 deleted file mode 100644 index 133a96e3..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py +++ /dev/null @@ -1,516 +0,0 @@ -# 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 deleted file mode 100644 index 20752bc1..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# 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 deleted file mode 100644 index 6c443246..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ /dev/null @@ -1,331 +0,0 @@ -# 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 deleted file mode 100644 index 3bd5e1b5..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ /dev/null @@ -1,610 +0,0 @@ -# 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 deleted file mode 100644 index 7cdb39e9..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ /dev/null @@ -1,173 +0,0 @@ -# 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 deleted file mode 100644 index ed16ca6d..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ /dev/null @@ -1,1380 +0,0 @@ -# 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 deleted file mode 100644 index 85a78883..00000000 --- a/TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py +++ /dev/null @@ -1,158 +0,0 @@ -# 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/TEMP/dsv4thing.zip b/TEMP/dsv4thing.zip new file mode 100644 index 00000000..2b1685dd Binary files /dev/null and b/TEMP/dsv4thing.zip differ diff --git a/TEMP/dsv4thing/README.md b/TEMP/dsv4thing/README.md new file mode 100644 index 00000000..c62e58cb --- /dev/null +++ b/TEMP/dsv4thing/README.md @@ -0,0 +1,26 @@ +# Inference code for DeepSeek models + +First convert huggingface model weight files to the format of this project. +```bash +export EXPERTS=384 +export MP=8 +export CONFIG=config.json +python convert.py --hf-ckpt-path ${HF_CKPT_PATH} --save-path ${SAVE_PATH} --n-experts ${EXPERTS} --model-parallel ${MP} +``` + +Then chat with DeepSeek model at will! +```bash +torchrun --nproc-per-node ${MP} generate.py --ckpt-path ${SAVE_PATH} --config ${CONFIG} --interactive +``` + +Or batch inference from file. +```bash +torchrun --nproc-per-node ${MP} generate.py --ckpt-path ${SAVE_PATH} --config ${CONFIG} --input-file ${FILE} +``` + +Or multi nodes inference. +```bash +torchrun --nnodes ${NODES} --nproc-per-node $((MP / NODES)) --node-rank $RANK --master-addr $ADDR generate.py --ckpt-path ${SAVE_PATH} --config ${CONFIG} --input-file ${FILE} +``` + +If you want to use fp8, just remove `"expert_dtype": "fp4"` in `config.json` and specify `--expert-dtype fp8` in `convert.py`. diff --git a/TEMP/dsv4thing/config.json b/TEMP/dsv4thing/config.json new file mode 100644 index 00000000..e2196678 --- /dev/null +++ b/TEMP/dsv4thing/config.json @@ -0,0 +1,35 @@ +{ + "vocab_size": 129280, + "dim": 7168, + "moe_inter_dim": 3072, + "n_layers": 61, + "n_hash_layers": 3, + "n_heads": 128, + "n_routed_experts": 384, + "n_shared_experts": 1, + "n_activated_experts": 6, + "score_func": "sqrtsoftplus", + "route_scale": 2.5, + "swiglu_limit": 10.0, + "q_lora_rank": 1536, + "head_dim": 512, + "rope_head_dim": 64, + "o_groups": 16, + "o_lora_rank": 1024, + "window_size": 128, + "original_seq_len": 65536, + "rope_theta": 10000, + "rope_factor": 16, + "beta_fast": 32, + "beta_slow": 1, + "index_n_heads": 64, + "index_head_dim": 128, + "index_topk": 1024, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "dtype": "fp8", + "scale_fmt": "ue8m0", + "expert_dtype": "fp4", + "compress_rope_theta": 160000, + "compress_ratios": [128, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0] +} \ No newline at end of file diff --git a/TEMP/dsv4thing/convert.py b/TEMP/dsv4thing/convert.py new file mode 100644 index 00000000..77748646 --- /dev/null +++ b/TEMP/dsv4thing/convert.py @@ -0,0 +1,168 @@ +import os +import shutil +from argparse import ArgumentParser +from glob import glob +from tqdm import tqdm, trange + +import torch +from safetensors.torch import safe_open, save_file + + +FP4_TABLE = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0 +], dtype=torch.float32) + + +def cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Casts a tensor from e2m1fn to e4m3fn losslessly. + """ + assert x.dtype == torch.int8 + assert x.ndim == 2 + out_dim, in_dim = x.size() + in_dim *= 2 + fp8_block_size = 128 + fp4_block_size = 32 + assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0 + assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size + + x = x.view(torch.uint8) + low = x & 0x0F + high = (x >> 4) & 0x0F + x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2) + + # max_fp4 (6.0) * MAX_OFFSET must fit in e4m3fn (max 448) + # 6.0 * 2^6 = 384 < 448; 6.0 * 2^7 = 768 > 448; so MAX_OFFSET_BITS = 6 + MAX_OFFSET_BITS = 6 + + bOut = out_dim // fp8_block_size + bIn = in_dim // fp8_block_size + # bOut, bIn, 128, 128 + x = x.view(bOut, fp8_block_size, bIn, fp8_block_size).transpose(1, 2) + # bOut, bIn, 128*4 + scale = scale.float().view(bOut, fp8_block_size, bIn, -1).transpose(1, 2).flatten(2) + ## bOut, bIn, 1 + scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS) + # bOut, bIn, 128*4 + offset = scale / scale_max_offset_bits + # bOut, bIn, 128, 128 + offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1) + x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim) + return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu) + + +mapping = { + "embed_tokens": ("embed", 0), + "input_layernorm": ("attn_norm", None), + "post_attention_layernorm": ("ffn_norm", None), + "q_proj": ("wq", 0), + "q_a_proj": ("wq_a", None), + "q_a_layernorm": ("q_norm", None), + "q_b_proj": ("wq_b", 0), + "kv_a_proj_with_mqa": ("wkv_a", None), + "kv_a_layernorm": ("kv_norm", None), + "kv_b_proj": ("wkv_b", 0), + "o_proj": ("wo", 1), + "gate_proj": ("w1", 0), + "down_proj": ("w2", 1), + "up_proj": ("w3", 0), + "lm_head": ("head", 0), + + "embed": ("embed", 0), + "wq_b": ("wq_b", 0), + "wo_a": ("wo_a", 0), + "wo_b": ("wo_b", 1), + "head": ("head", 0), + "attn_sink": ("attn_sink", 0), + "weights_proj": ("weights_proj", 0), +} + + +def main(hf_ckpt_path, save_path, n_experts, mp, expert_dtype): + """ + Converts and saves model checkpoint files into a specified format. + + Args: + hf_ckpt_path (str): Path to the directory containing the input checkpoint files. + save_path (str): Path to the directory where the converted checkpoint files will be saved. + n_experts (int): Total number of experts in the model. + mp (int): Model parallelism factor. + + Returns: + None + """ + torch.set_num_threads(8) + n_local_experts = n_experts // mp + state_dicts = [{} for _ in range(mp)] + + for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))): + with safe_open(file_path, framework="pt", device="cpu") as f: + for name in f.keys(): + param: torch.Tensor = f.get_tensor(name) + if name.startswith("model."): + name = name[len("model."):] + if name.startswith("mtp.") and ("emb" in name or name.endswith("head.weight")): + continue + name = name.replace("self_attn", "attn") + name = name.replace("mlp", "ffn") + name = name.replace("weight_scale_inv", "scale") + name = name.replace("e_score_correction_bias", "bias") + if any(x in name for x in ["hc", "attn_sink", "tie2eid", "ape"]): # without .weight + key = name.split(".")[-1] + else: + key = name.split(".")[-2] + if key in mapping: + new_key, dim = mapping[key] + else: + new_key, dim = key, None + name = name.replace(key, new_key) + for i in range(mp): + new_param = param + if "experts" in name and "shared_experts" not in name: + idx = int(name.split(".")[-3]) + if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts: + continue + elif dim is not None: + assert param.size(dim) % mp == 0, f"Dimension {dim} must be divisible by {mp}" + shard_size = param.size(dim) // mp + new_param = param.narrow(dim, i * shard_size, shard_size).contiguous() + state_dicts[i][name] = new_param + + os.makedirs(save_path, exist_ok=True) + + for i in trange(mp): + names = list(state_dicts[i].keys()) + for name in names: + if name.endswith("wo_a.weight"): + weight = state_dicts[i][name] + scale = state_dicts[i].pop(name.replace("weight", "scale")) + weight = weight.unflatten(0, (-1, 128)).unflatten(-1, (-1, 128)).float() * scale[:, None, :, None].float() + state_dicts[i][name] = weight.flatten(2, 3).flatten(0, 1).bfloat16() + elif "experts" in name and state_dicts[i][name].dtype == torch.int8: + if expert_dtype == "fp8": + scale_name = name.replace("weight", "scale") + weight = state_dicts[i].pop(name) + scale = state_dicts[i].pop(scale_name) + state_dicts[i][name], state_dicts[i][scale_name] = cast_e2m1fn_to_e4m3fn(weight, scale) + else: + state_dicts[i][name] = state_dicts[i][name].view(torch.float4_e2m1fn_x2) + save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors")) + + for file in ["tokenizer.json", "tokenizer_config.json"]: + old_file_path = os.path.join(hf_ckpt_path, file) + new_file_path = os.path.join(save_path, file) + if os.path.exists(old_file_path): + shutil.copyfile(old_file_path, new_file_path) + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--hf-ckpt-path", type=str, required=True) + parser.add_argument("--save-path", type=str, required=True) + parser.add_argument("--n-experts", type=int, required=True) + parser.add_argument("--model-parallel", type=int, required=True) + parser.add_argument("--expert-dtype", type=str, choices=["fp8", "fp4"], required=False, default=None) + args = parser.parse_args() + assert args.n_experts % args.model_parallel == 0, "Number of experts must be divisible by model parallelism" + main(args.hf_ckpt_path, args.save_path, args.n_experts, args.model_parallel, args.expert_dtype) diff --git a/TEMP/dsv4thing/generate.py b/TEMP/dsv4thing/generate.py new file mode 100644 index 00000000..c35c8030 --- /dev/null +++ b/TEMP/dsv4thing/generate.py @@ -0,0 +1,155 @@ +import os +import json +import sys +from argparse import ArgumentParser +from typing import List + +import torch +import torch.distributed as dist +from transformers import AutoTokenizer +from safetensors.torch import load_model + +from model import Transformer, ModelArgs +current_dir = os.path.dirname(os.path.abspath(__file__)) +encoding_dir = os.path.join(current_dir, '../encoding') +sys.path.insert(0, os.path.abspath(encoding_dir)) +from encoding_dsv4 import encode_messages, parse_message_from_completion_text + + +def sample(logits, temperature: float = 1.0): + """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU, + since it avoids the GPU-to-CPU sync in torch.multinomial.""" + logits = logits / max(temperature, 1e-5) + probs = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1) + + +@torch.inference_mode() +def generate( + model: Transformer, + prompt_tokens: List[List[int]], + max_new_tokens: int, + eos_id: int, + temperature: float = 1.0 +) -> List[List[int]]: + """Batch generation with left-padded prompts. + + The first forward pass processes [min_prompt_len:] tokens (prefill phase). + Subsequent passes generate one token at a time (decode phase). For positions + still within a prompt, the ground-truth token overrides the model's prediction. + """ + prompt_lens = [len(t) for t in prompt_tokens] + assert max(prompt_lens) <= model.max_seq_len, f"Prompt length exceeds model maximum sequence length (max_seq_len={model.max_seq_len})" + total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens)) + tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long) + for i, t in enumerate(prompt_tokens): + tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long) + prev_pos = 0 + finished = torch.tensor([False] * len(prompt_tokens)) + prompt_mask = tokens != -1 + for cur_pos in range(min(prompt_lens), total_len): + logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos) + if temperature > 0: + next_token = sample(logits, temperature) + else: + next_token = logits.argmax(dim=-1) + next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token) + tokens[:, cur_pos] = next_token + finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id) + prev_pos = cur_pos + if finished.all(): + break + completion_tokens = [] + for i, toks in enumerate(tokens.tolist()): + toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens] + if eos_id in toks: + toks = toks[:toks.index(eos_id)] + toks.append(eos_id) + completion_tokens.append(toks) + return completion_tokens + + +def main( + ckpt_path: str, + config: str, + input_file: str = "", + interactive: bool = True, + max_new_tokens: int = 100, + temperature: float = 1.0, +) -> None: + world_size = int(os.getenv("WORLD_SIZE", "1")) + rank = int(os.getenv("RANK", "0")) + local_rank = int(os.getenv("LOCAL_RANK", "0")) + if world_size > 1: + dist.init_process_group("nccl") + global print + if rank != 0: + print = lambda *_, **__: None + torch.cuda.set_device(local_rank) + torch.cuda.memory._set_allocator_settings("expandable_segments:True") + torch.set_default_dtype(torch.bfloat16) + torch.set_num_threads(8) + torch.manual_seed(33377335) + with open(config) as f: + args = ModelArgs(**json.load(f)) + if interactive: + args.max_batch_size = 1 + print(args) + with torch.device("cuda"): + model = Transformer(args) + tokenizer = AutoTokenizer.from_pretrained(ckpt_path) + print("load model") + load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"), strict=False) + torch.set_default_device("cuda") + print("I'm DeepSeek 👋") + + if interactive: + messages = [] + while True: + if world_size == 1: + prompt = input(">>> ") + elif rank == 0: + prompt = input(">>> ") + objects = [prompt] + dist.broadcast_object_list(objects, 0) + else: + objects = [None] + dist.broadcast_object_list(objects, 0) + prompt = objects[0] + if prompt == "/exit": + break + elif prompt == "/clear": + messages.clear() + continue + messages.append({"role": "user", "content": prompt}) + prompt_tokens = tokenizer.encode(encode_messages(messages, thinking_mode="chat")) + completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature) + completion = tokenizer.decode(completion_tokens[0]) + print(completion) + messages.append(parse_message_from_completion_text(completion, thinking_mode="chat")) + else: + with open(input_file) as f: + prompts = f.read().split("\n\n") + prompt_tokens = [tokenizer.encode(encode_messages([{"role": "user", "content": prompt}], thinking_mode="chat")) for prompt in prompts] + completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature) + completions = tokenizer.batch_decode(completion_tokens) + for prompt, completion in zip(prompts, completions): + print("Prompt:", prompt) + print("Completion:", completion) + print() + + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--ckpt-path", type=str, required=True) + parser.add_argument("--config", type=str, required=True) + parser.add_argument("--input-file", type=str, default="") + parser.add_argument("--interactive", action="store_true") + parser.add_argument("--max-new-tokens", type=int, default=300) + parser.add_argument("--temperature", type=float, default=0.6) + args = parser.parse_args() + assert args.input_file or args.interactive, "Either input-file or interactive mode must be specified" + main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature) diff --git a/TEMP/dsv4thing/kernel.py b/TEMP/dsv4thing/kernel.py new file mode 100644 index 00000000..ea7976fa --- /dev/null +++ b/TEMP/dsv4thing/kernel.py @@ -0,0 +1,536 @@ +import torch +import tilelang +import tilelang.language as T +from typing import Tuple, Optional + + +tilelang.set_log_level("WARNING") + +pass_configs = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} + +FP8 = "float8_e4m3" +FP4 = "float4_e2m1fn" +FE8M0 = "float8_e8m0fnu" +BF16 = "bfloat16" +FP32 = "float32" +INT32 = "int32" + + +def fast_log2_ceil(x): + """Compute ceil(log2(x)) via IEEE 754 bit manipulation. Avoids slow log/ceil intrinsics.""" + bits_x = T.reinterpret("uint32", x) + exp_x = (bits_x >> 23) & 0xFF + man_bits = bits_x & ((1 << 23) - 1) + return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0)) + + +def fast_pow2(x): + """Compute 2^x for integer x via IEEE 754 bit manipulation.""" + bits_x = (x + 127) << 23 + return T.reinterpret("float32", bits_x) + + +def fast_round_scale(amax, fp8_max_inv): + return fast_pow2(fast_log2_ceil(amax * fp8_max_inv)) + + +@tilelang.jit(pass_configs=pass_configs) +def act_quant_kernel( + N, block_size=128, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, + round_scale=False, inplace=False +): + """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16.""" + M = T.symbolic("M") + fp8_min = -448.0 + fp8_max = 448.0 + fp8_max_inv = 1 / fp8_max + num_stages = 0 if round_scale or inplace else 2 + blk_m = 32 + group_size = block_size + # Internal computation in FP32; scale_dtype controls output storage format. + compute_dtype = FP32 + out_dtype = in_dtype if inplace else out_dtype + + @T.prim_func + def act_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), compute_dtype) + s_local = T.alloc_fragment((blk_m,), compute_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=num_stages): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 1e-4) + if round_scale: + s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv) + else: + s_local[i] = amax_local[i] * fp8_max_inv + if inplace: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.Cast( + out_dtype, + T.Cast(compute_dtype, T.Cast(out_dtype, T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ))) * s_local[i], + ) + else: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i]) + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return act_quant_kernel_ + + +def act_quant( + x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None, + scale_dtype: torch.dtype = torch.float32, inplace: bool = False, +) -> torch.Tensor: + """Block-wise FP8 quantization. inplace=True does fused quant+dequant back to BF16. + When scale_fmt is set, scales are rounded to power-of-2 (MXFP).""" + N = x.size(-1) + assert N % block_size == 0 + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + z = x.contiguous() + y = torch.empty_like(z) if inplace else torch.empty_like(z, dtype=torch.float8_e4m3fn) + s = z.new_empty(*z.size()[:-1], N // block_size, dtype=scale_dtype) + kernel = act_quant_kernel( + N, block_size, scale_dtype=tl_dtype, + round_scale=scale_fmt is not None, inplace=inplace, + ) + kernel(z.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) + if inplace: + x.copy_(y) + return x + return y, s + + +@tilelang.jit(pass_configs=pass_configs) +def fp4_quant_kernel( + N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False +): + """Block-wise FP4 quantization. Power-of-2 scale via bit ops. inplace=True does fused quant+dequant.""" + M = T.symbolic("M") + fp4_max = 6.0 + fp4_max_inv = 1.0 / fp4_max + blk_m = 32 + group_size = block_size + compute_dtype = FP32 + out_dtype = in_dtype if inplace else FP4 + + @T.prim_func + def fp4_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), compute_dtype) + s_local = T.alloc_fragment((blk_m,), compute_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=2): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 6 * (2**-126)) + s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv) + if inplace: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.Cast( + out_dtype, + T.Cast(compute_dtype, T.Cast(FP4, T.clamp( + x_local[i, j] / s_local[i], -fp4_max, fp4_max + ))) * s_local[i], + ) + else: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], -fp4_max, fp4_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i]) + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return fp4_quant_kernel_ + + +def fp4_act_quant( + x: torch.Tensor, block_size: int = 32, inplace: bool = False, +) -> torch.Tensor: + """Block-wise FP4 quantization. inplace=True does fused quant+dequant back to BF16.""" + N = x.size(-1) + assert N % block_size == 0 + z = x.contiguous() + y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2) + s = z.new_empty(*z.size()[:-1], N // block_size, dtype=torch.float8_e8m0fnu) + kernel = fp4_quant_kernel(N, block_size, inplace=inplace) + kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size)) + if inplace: + x.copy_(y) + return x + return y, s + + +@tilelang.jit(pass_configs=pass_configs) +def fp8_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32): + assert out_dtype in [BF16, FP32] + + M = T.symbolic("M") + group_size = 128 + block_M = 32 + block_N = 128 + block_K = 128 + + @T.prim_func + def fp8_gemm_kernel_( + A: T.Tensor[(M, K), FP8], + B: T.Tensor[(N, K), FP8], + C: T.Tensor[(M, N), out_dtype], + scales_a: T.Tensor[(M, T.ceildiv(K, group_size)), scale_dtype], + scales_b: T.Tensor[(T.ceildiv(N, group_size), T.ceildiv(K, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + A_shared = T.alloc_shared((block_M, block_K), FP8) + B_shared = T.alloc_shared((block_N, block_K), FP8) + C_shared = T.alloc_shared((block_M, block_N), out_dtype) + Scale_C_shared = T.alloc_shared((block_M), FP32) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Improve L2 Cache + T.use_swizzle(panel_size=10) + T.clear(C_local) + T.clear(C_local_accum) + + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=4): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) + # Cast scales to FP32 for computation; scales_b has one value per block_N group + Scale_B = T.Cast(FP32, scales_b[bx * block_N // group_size, k]) + for i in T.Parallel(block_M): + Scale_C_shared[i] = T.Cast(FP32, scales_a[by * block_M + i, k]) * Scale_B + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + # Separate accumulator for scale-corrected results (2x accumulation precision) + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i] + T.clear(C_local) + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return fp8_gemm_kernel_ + + +def fp8_gemm( + a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """C[M,N] = A[M,K] @ B[N,K]^T with per-128 block FP8 scaling on both A and B.""" + assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous" + assert a_s.is_contiguous() and b_s.is_contiguous(), ( + "Scaling factor tensors must be contiguous" + ) + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype()) + kernel = fp8_gemm_kernel(N, K, scale_dtype=tl_dtype) + kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s) + return c + + +@tilelang.jit(pass_configs=pass_configs) +def sparse_attn_kernel(h: int, d: int, scale=None): + """Sparse multi-head attention via index gathering + online softmax (FlashAttention-style). + For each (batch, seq_pos), gathers top-k KV positions by index, computes attention + with numerically stable running max/sum, and includes a learnable attn_sink bias.""" + b = T.symbolic("b") + m = T.symbolic("m") + n = T.symbolic("n") + topk = T.symbolic("topk") + if scale is None: + scale = (1.0 / d) ** 0.5 + + num_stages = 2 + threads = 256 + block = 64 + num_blocks = tilelang.cdiv(topk, block) + + @T.prim_func + def sparse_attn_kernel_( + q: T.Tensor[(b, m, h, d), BF16], + kv: T.Tensor[(b, n, d), BF16], + o: T.Tensor[(b, m, h, d), BF16], + attn_sink: T.Tensor[(h,), FP32], + topk_idxs: T.Tensor[(b, m, topk), INT32], + ): + with T.Kernel(m, b, threads=threads) as (bx, by): + q_shared = T.alloc_shared((h, d), BF16) + kv_shared = T.alloc_shared((block, d), BF16) + o_shared = T.alloc_shared((h, d), BF16) + acc_s_cast = T.alloc_shared((h, block), BF16) + + idxs = T.alloc_fragment(block, INT32) + acc_s = T.alloc_fragment((h, block), FP32) + acc_o = T.alloc_fragment((h, d), FP32) + scores_max = T.alloc_fragment(h, FP32) + scores_max_prev = T.alloc_fragment(h, FP32) + scores_scale = T.alloc_fragment(h, FP32) + scores_sum = T.alloc_fragment(h, FP32) + sum_exp = T.alloc_fragment(h, FP32) + + T.clear(acc_o) + T.clear(sum_exp) + T.fill(scores_max, -T.infinity(FP32)) + T.copy(q[by, bx, :, :], q_shared) + + for t in T.Pipelined(num_blocks, num_stages=num_stages): + for i in T.Parallel(block): + idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1) + for i, j in T.Parallel(block, d): + kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0) + for i, j in T.Parallel(h, block): + acc_s[i, j] = T.if_then_else(idxs[j] != -1, 0, -T.infinity(FP32)) + T.gemm(q_shared, kv_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + for i, j in T.Parallel(h, block): + acc_s[i, j] *= scale + T.copy(scores_max, scores_max_prev) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(h): + scores_scale[i] = T.exp(scores_max_prev[i] - scores_max[i]) + for i, j in T.Parallel(h, block): + acc_s[i, j] = T.exp(acc_s[i, j] - scores_max[i]) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(h): + sum_exp[i] = sum_exp[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + for i, j in T.Parallel(h, d): + acc_o[i, j] *= scores_scale[i] + T.gemm(acc_s_cast, kv_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for i in T.Parallel(h): + sum_exp[i] += T.exp(attn_sink[i] - scores_max[i]) + for i, j in T.Parallel(h, d): + acc_o[i, j] /= sum_exp[i] + T.copy(acc_o, o_shared) + T.copy(o_shared, o[by, bx, :, :]) + + return sparse_attn_kernel_ + + +def sparse_attn( + q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float +) -> torch.Tensor: + b, s, h, d = q.size() + # Pad heads to 16 for kernel efficiency (stripped after) + if h < 16: + q = torch.cat([q, q.new_zeros(b, s, 16 - h, d)], dim=2) + attn_sink = torch.cat([attn_sink, attn_sink.new_zeros(16 - h)]) + o = torch.empty_like(q) + kernel = sparse_attn_kernel(q.size(2), d, softmax_scale) + kernel(q, kv, o, attn_sink, topk_idxs) + if h < 16: + o = o.narrow(2, 0, h).contiguous() + return o + + +@tilelang.jit(pass_configs=pass_configs) +def hc_split_sinkhorn_kernel(hc: int, sinkhorn_iters: int, eps: float): + n = T.symbolic("n") + mix_hc = (2 + hc) * hc + threads = 64 + + @T.prim_func + def hc_split_sinkhorn_kernel_( + mixes: T.Tensor[(n, mix_hc), FP32], + hc_scale: T.Tensor[(3,), FP32], + hc_base: T.Tensor[(mix_hc,), FP32], + pre: T.Tensor[(n, hc), FP32], + post: T.Tensor[(n, hc), FP32], + comb: T.Tensor[(n, hc, hc), FP32], + ): + with T.Kernel(n, threads=threads) as i: + mixes_shared = T.alloc_shared(mix_hc, FP32) + comb_frag = T.alloc_fragment((hc, hc), FP32) + T.copy(mixes[i, :], mixes_shared) + + for j in T.Parallel(hc): + pre[i, j] = T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + eps + for j in T.Parallel(hc): + post[i, j] = 2 * T.sigmoid(mixes_shared[j + hc] * hc_scale[1] + hc_base[j + hc]) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = mixes_shared[j * hc + k + hc * 2] * hc_scale[2] + hc_base[j * hc + k + hc * 2] + + row_sum = T.alloc_fragment(hc, FP32) + col_sum = T.alloc_fragment(hc, FP32) + + # comb = comb.softmax(-1) + eps + row_max = T.alloc_fragment(hc, FP32) + T.reduce_max(comb_frag, row_max, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = T.exp(comb_frag[j, k] - row_max[j]) + T.reduce_sum(comb_frag, row_sum, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / row_sum[j] + eps + + # comb = comb / (comb.sum(-2) + eps) + T.reduce_sum(comb_frag, col_sum, dim=0) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps) + + for _ in T.serial(sinkhorn_iters - 1): + # comb = comb / (comb.sum(-1) + eps) + T.reduce_sum(comb_frag, row_sum, dim=1) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (row_sum[j] + eps) + # comb = comb / (comb.sum(-2) + eps) + T.reduce_sum(comb_frag, col_sum, dim=0) + for j, k in T.Parallel(hc, hc): + comb_frag[j, k] = comb_frag[j, k] / (col_sum[k] + eps) + + T.copy(comb_frag, comb[i, :, :]) + + return hc_split_sinkhorn_kernel_ + + +def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6): + b, s, _ = mixes.size() + pre = mixes.new_empty(b, s, hc_mult) + post = mixes.new_empty(b, s, hc_mult) + comb = mixes.new_empty(b, s, hc_mult, hc_mult) + kernel = hc_split_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps) + kernel(mixes.view(-1, (2 + hc_mult) * hc_mult), hc_scale, hc_base, + pre.view(-1, hc_mult), post.view(-1, hc_mult), comb.view(-1, hc_mult, hc_mult)) + return pre, post, comb + + +@tilelang.jit(pass_configs=pass_configs) +def fp4_gemm_kernel(N, K, out_dtype=BF16, accum_dtype=FP32, scale_dtype=FP32): + """FP8 act x FP4 weight GEMM kernel. + + C[M, N] = A_fp8[M, K] @ B_fp4[N, K]^T + + Act: 1x128 quant on K (reduce dim), FP8 with configurable scale dtype + Weight: 1x32 quant on K (reduce dim), FP4 with E8M0 scale + + B is stored as [N, K//2] in float4_e2m1fn_x2, logical [N, K] in fp4. + The FP4 values are packed along the K (last) dimension. + + Strategy: load FP4 sub-blocks of size [block_N, sub_K] (sub_K=32), + cast FP4 to FP8 via float, then do FP8xFP8 GEMM. + Apply act scale (per 128 on K) and weight scale (per 32 on K) to the accumulator. + """ + M = T.symbolic("M") + act_group_size = 128 + weight_group_size = 32 + block_M = 32 + block_N = 128 + block_K = 32 # matches weight_group_size for simple scale handling + n_sub = act_group_size // block_K # 4 sub-blocks per act scale group + + @T.prim_func + def fp4_gemm_kernel_( + A: T.Tensor[(M, K), FP8], + B: T.Tensor[(N, K), FP4], + C: T.Tensor[(M, N), out_dtype], + scales_a: T.Tensor[(M, T.ceildiv(K, act_group_size)), scale_dtype], + scales_b: T.Tensor[(N, T.ceildiv(K, weight_group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + A_shared = T.alloc_shared((block_M, block_K), FP8) + B_fp4_shared = T.alloc_shared((block_N, block_K), FP4) + B_shared = T.alloc_shared((block_N, block_K), FP8) + C_shared = T.alloc_shared((block_M, block_N), out_dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_local_accum = T.alloc_fragment((block_M, block_N), accum_dtype) + scale_a_frag = T.alloc_fragment((block_M,), FP32) + scale_b_frag = T.alloc_fragment((block_N,), FP32) + + T.use_swizzle(panel_size=10) + T.clear(C_local) + T.clear(C_local_accum) + + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=2): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_fp4_shared) + # FP4->FP8 cast must go through FP32 to avoid ambiguous C++ overload + for i, j in T.Parallel(block_N, block_K): + B_shared[i, j] = T.Cast(FP8, T.Cast(FP32, B_fp4_shared[i, j])) + + # Weight scale: per 32 on K, indexed by k (each k is one block_K=32) + for i in T.Parallel(block_N): + scale_b_frag[i] = T.Cast(FP32, scales_b[bx * block_N + i, k]) + + # Act scale: per 128 on K, indexed by k // 4 + for i in T.Parallel(block_M): + scale_a_frag[i] = T.Cast(FP32, scales_a[by * block_M + i, k // n_sub]) + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * scale_a_frag[i] * scale_b_frag[j] + T.clear(C_local) + + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return fp4_gemm_kernel_ + + +def fp4_gemm( + a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor, + scale_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """C[M,N] = A_fp8[M,K] @ B_fp4[N,K]^T. + A has per-128 act scale; B has per-32 E8M0 weight scale. + B is stored as [N, K//2] in float4_e2m1fn_x2 (2 FP4 values per byte, packed along K).""" + assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous" + assert a_s.is_contiguous() and b_s.is_contiguous(), ( + "Scaling factor tensors must be contiguous" + ) + tl_dtype = FE8M0 if scale_dtype == torch.float8_e8m0fnu else FP32 + K = a.size(-1) + M = a.numel() // K + N = b.size(0) + c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype()) + kernel = fp4_gemm_kernel(N, K, scale_dtype=tl_dtype) + kernel(a.view(M, K), b, c.view(M, N), a_s.view(M, -1), b_s) + return c diff --git a/TEMP/dsv4thing/model.py b/TEMP/dsv4thing/model.py new file mode 100644 index 00000000..167ade8f --- /dev/null +++ b/TEMP/dsv4thing/model.py @@ -0,0 +1,827 @@ +import math +from dataclasses import dataclass +from typing import Tuple, Optional, Literal +from functools import lru_cache +from contextlib import contextmanager + +import torch +from torch import nn +import torch.nn.functional as F +import torch.distributed as dist + +from kernel import act_quant, fp4_act_quant, fp8_gemm, fp4_gemm, sparse_attn, hc_split_sinkhorn + + +world_size = 1 +rank = 0 +block_size = 128 +fp4_block_size = 32 +default_dtype = torch.bfloat16 +scale_fmt = None +scale_dtype = torch.float32 + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch default dtype, restoring it on exit (even if an exception occurs).""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + +@dataclass +class ModelArgs: + """Model hyperparameters. Field names match the config JSON keys.""" + max_batch_size: int = 4 + max_seq_len: int = 4096 + dtype: Literal["bf16", "fp8"] = "fp8" + scale_fmt: Literal[None, "ue8m0"] = "ue8m0" + expert_dtype: Literal[None, "fp4"] = None + scale_dtype: Literal["fp32", "fp8"] = "fp8" + vocab_size: int = 129280 + dim: int = 4096 + moe_inter_dim: int = 4096 + n_layers: int = 7 + n_hash_layers: int = 0 + n_mtp_layers: int = 1 + n_heads: int = 64 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + route_scale: float = 1. + swiglu_limit: float = 0. + # mqa + q_lora_rank: int = 1024 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + o_groups: int = 8 + o_lora_rank: int = 1024 + window_size: int = 128 + compress_ratios: Tuple[int] = (0, 0, 4, 128, 4, 128, 4, 0) + # yarn + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # index + index_n_heads: int = 64 + index_head_dim: int = 128 + index_topk: int = 512 + # hc + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + + +class ParallelEmbedding(nn.Module): + """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows. + Out-of-range indices are zero-masked before all_reduce to combine partial embeddings.""" + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})" + self.part_vocab_size = (vocab_size // world_size) + self.vocab_start_idx = rank * self.part_vocab_size + self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if world_size > 1: + mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) + x = x - self.vocab_start_idx + x[mask] = 0 + y = F.embedding(x, self.weight) + if world_size > 1: + y[mask] = 0 + dist.all_reduce(y) + return y + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor: + """Dispatches to fp4_gemm / fp8_gemm / F.linear based on weight dtype. + For quantized weights, x is first quantized to FP8 via act_quant.""" + assert bias is None + + if weight.dtype == torch.float4_e2m1fn_x2: + x, s = act_quant(x, block_size, scale_fmt, scale_dtype) + return fp4_gemm(x, s, weight, weight.scale, scale_dtype) + elif weight.dtype == torch.float8_e4m3fn: + x, s = act_quant(x, block_size, scale_fmt, scale_dtype) + return fp8_gemm(x, s, weight, weight.scale, scale_dtype) + else: + return F.linear(x, weight) + + +class Linear(nn.Module): + """Linear layer supporting BF16, FP8, and FP4 weight formats with per-block scaling.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or default_dtype + if dtype == torch.float4_e2m1fn_x2: + # FP4: weight is [out, in//2] in float4_e2m1fn_x2, logically [out, in] in fp4 + # Scale is [out, in//32] in float8_e8m0fnu (1 scale per 32 fp4 elements along K) + self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2)) + scale_out_features = out_features + scale_in_features = in_features // fp4_block_size + self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu)) + elif dtype == torch.float8_e4m3fn: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + scale_out_features = (out_features + block_size - 1) // block_size + scale_in_features = (in_features + block_size - 1) // block_size + self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float8_e8m0fnu)) + else: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class ColumnParallelLinear(Linear): + """Shards output dim across TP ranks. No all-reduce needed on output.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Shards input dim across TP ranks. All-reduce on output to sum partial results.""" + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): + assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})" + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if world_size > 1: + y = y.float() + dist.all_reduce(y) + if self.bias is not None: + y += self.bias + return y.type_as(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + # rmsnorm in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient. + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +@lru_cache(2) +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor: + """Precomputes complex exponentials for rotary embeddings with YaRN scaling. + When original_seq_len > 0, applies frequency interpolation with a smooth + linear ramp between beta_fast and beta_slow correction ranges.""" + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim-1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Applies rotary positional embeddings in-place. Uses conjugate for inverse (de-rotation).""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Applies randomized Hadamard rotation to spread information across dims before FP8 quant.""" + assert x.dtype == torch.bfloat16 + from fast_hadamard_transform import hadamard_transform + return hadamard_transform(x, scale=x.size(-1) ** -0.5) + + +@lru_cache(1) +def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + if start_pos >= window_size - 1: + start_pos %= window_size + matrix = torch.cat([torch.arange(start_pos + 1, window_size), torch.arange(0, start_pos + 1)], dim=0) + elif start_pos > 0: + matrix = F.pad(torch.arange(start_pos + 1), (0, window_size - start_pos - 1), value=-1) + else: + base = torch.arange(seqlen).unsqueeze(1) + matrix = (base - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + matrix = torch.where(matrix > base, -1, matrix) + return matrix.unsqueeze(0).expand(bsz, -1, -1) + + +@lru_cache(2) +def get_compress_topk_idxs(ratio: int, bsz: int, seqlen: int, start_pos: int, offset: int): + if start_pos > 0: + matrix = torch.arange(0, (start_pos + 1) // ratio) + offset + else: + matrix = torch.arange(seqlen // ratio).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix.unsqueeze(0).expand(bsz, -1, -1) + + +class Compressor(nn.Module): + """Compresses KV cache via learned gated pooling over `compress_ratio` consecutive tokens. + When overlap=True (ratio==4), uses overlapping windows for smoother compression boundaries.""" + + def __init__(self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False): + super().__init__() + self.dim = args.dim + self.head_dim = head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = head_dim - args.rope_head_dim + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.rotate = rotate + coff = 1 + self.overlap + + self.ape = nn.Parameter(torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32)) + # wkv and wgate in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for convenient. + # When overlap, the first half of dims is for overlapping compression, second half for normal. + self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.norm = RMSNorm(self.head_dim, args.norm_eps) + self.kv_cache: torch.Tensor = None # assigned lazily from Attention.kv_cache + # State buffers for decode-phase incremental compression. + # With overlap: state[:, :ratio] = overlapping window, state[:, ratio:] = current window. + self.register_buffer("kv_state", torch.zeros(args.max_batch_size, coff * compress_ratio, coff * self.head_dim, dtype=torch.float32), persistent=False) + self.register_buffer("score_state", torch.full((args.max_batch_size, coff * compress_ratio, coff * self.head_dim), float("-inf"), dtype=torch.float32), persistent=False) + self.freqs_cis: torch.Tensor = None + + def overlap_transform(self, tensor: torch.Tensor, value=0): + # tensor: [b,s,r,2d] + b, s, _, _ = tensor.size() + ratio, d = self.compress_ratio, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int): + assert self.kv_cache is not None + bsz, seqlen, _ = x.size() + ratio, overlap, d, rd = self.compress_ratio, self.overlap, self.head_dim, self.rope_head_dim + dtype = x.dtype + # compression need fp32 + x = x.float() + kv = self.wkv(x) + score = self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + offset = ratio if overlap else 0 + if overlap and cutoff >= ratio: + self.kv_state[:bsz, :ratio] = kv[:, cutoff-ratio : cutoff] + self.score_state[:bsz, :ratio] = score[:, cutoff-ratio : cutoff] + self.ape + if remainder > 0: + kv, self.kv_state[:bsz, offset : offset+remainder] = kv.split([cutoff, remainder], dim=1) + self.score_state[:bsz, offset : offset+remainder] = score[:, cutoff:] + self.ape[:remainder] + score = score[:, :cutoff] + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + if overlap: + kv = self.overlap_transform(kv, 0) + score = self.overlap_transform(score, float("-inf")) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: + should_compress = (start_pos + 1) % self.compress_ratio == 0 + score += self.ape[start_pos % ratio] + if overlap: + self.kv_state[:bsz, ratio + start_pos % ratio] = kv.squeeze(1) + self.score_state[:bsz, ratio + start_pos % ratio] = score.squeeze(1) + if should_compress: + kv_state = torch.cat([self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], dim=1) + score_state = torch.cat([self.score_state[:bsz, :ratio, :d], self.score_state[:bsz, ratio:, d:]], dim=1) + kv = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True) + self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:] + self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:] + else: + self.kv_state[:bsz, start_pos % ratio] = kv.squeeze(1) + self.score_state[:bsz, start_pos % ratio] = score.squeeze(1) + if should_compress: + kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + if not should_compress: + return + kv = self.norm(kv.to(dtype)) + if start_pos == 0: + freqs_cis = self.freqs_cis[:cutoff:ratio] + else: + freqs_cis = self.freqs_cis[start_pos + 1 - self.compress_ratio].unsqueeze(0) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + if self.rotate: + kv = rotate_activation(kv) + fp4_act_quant(kv, fp4_block_size, True) + else: + act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True) + if start_pos == 0: + self.kv_cache[:bsz, :seqlen // ratio] = kv + else: + self.kv_cache[:bsz, start_pos // ratio] = kv.squeeze(1) + return kv + + +class Indexer(torch.nn.Module): + """Selects top-k compressed KV positions for sparse attention via learned scoring. + Has its own Compressor (with Hadamard rotation) to build compressed KV for scoring.""" + + def __init__(self, args: ModelArgs, compress_ratio: int = 4): + super().__init__() + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.softmax_scale = self.head_dim ** -0.5 + self.compress_ratio = compress_ratio + + self.compressor = Compressor(args, compress_ratio, self.head_dim, True) + self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len // compress_ratio, self.head_dim), persistent=False) + self.freqs_cis = None + + def forward(self, x: torch.Tensor, qr: torch.Tensor, start_pos: int, offset: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] + ratio = self.compress_ratio + rd = self.rope_head_dim + end_pos = start_pos + seqlen + if self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache + self.compressor.freqs_cis = self.freqs_cis + q = self.wq_b(qr) + q = q.unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + q = rotate_activation(q) + # use fp4 simulation for q and kv in indexer + fp4_act_quant(q, fp4_block_size, True) + self.compressor(x, start_pos) + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16 + index_score = torch.einsum("bshd,btd->bsht", q, self.kv_cache[:bsz, :end_pos // ratio]) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if world_size > 1: + dist.all_reduce(index_score) + if start_pos == 0: + mask = torch.arange(seqlen // ratio).repeat(seqlen, 1) >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + index_score += torch.where(mask, float("-inf"), 0) + topk_idxs = index_score.topk(min(self.index_topk, end_pos // ratio), dim=-1)[1] + if start_pos == 0: + mask = topk_idxs >= torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + topk_idxs = torch.where(mask, -1, topk_idxs + offset) + else: + topk_idxs += offset + return topk_idxs + + +class Attention(nn.Module): + """Multi-head Latent Attention (MLA) with sliding window + optional KV compression. + Uses low-rank Q projection (wq_a -> q_norm -> wq_b) and grouped low-rank O projection.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + self.n_heads = args.n_heads + self.n_local_heads = args.n_heads // world_size + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.head_dim = args.head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = args.head_dim - args.rope_head_dim + self.n_groups = args.o_groups + self.n_local_groups = self.n_groups // world_size + self.window_size = args.window_size + self.compress_ratio = args.compress_ratios[layer_id] + self.eps = args.norm_eps + + self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32)) + self.wq_a = Linear(self.dim, self.q_lora_rank) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.wkv = Linear(self.dim, self.head_dim) + 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 * args.o_lora_rank, dtype=torch.bfloat16) + self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim) + self.softmax_scale = self.head_dim ** -0.5 + + if self.compress_ratio: + self.compressor = Compressor(args, self.compress_ratio, self.head_dim) + if self.compress_ratio == 4: + self.indexer = Indexer(args, self.compress_ratio) + else: + self.indexer = None + + kv_cache_size = args.window_size + (args.max_seq_len // self.compress_ratio if self.compress_ratio else 0) + self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim), persistent=False) + if self.compress_ratio: + original_seq_len, rope_theta = args.original_seq_len, args.compress_rope_theta + else: + # disable YaRN and use base rope_theta in pure sliding-window attention + original_seq_len, rope_theta = 0, args.rope_theta + freqs_cis = precompute_freqs_cis(self.rope_head_dim, args.max_seq_len, original_seq_len, + rope_theta, args.rope_factor, args.beta_fast, args.beta_slow) + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + def forward(self, x: torch.Tensor, start_pos: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] + win = self.window_size + ratio = self.compress_ratio + rd = self.rope_head_dim + if self.compress_ratio and self.compressor.kv_cache is None: + self.compressor.kv_cache = self.kv_cache[:, win:] + self.compressor.freqs_cis = self.freqs_cis + if self.indexer is not None: + self.indexer.freqs_cis = self.freqs_cis + # q + qr = q = self.q_norm(self.wq_a(x)) + q = self.wq_b(q).unflatten(-1, (self.n_local_heads, self.head_dim)) + q *= torch.rsqrt(q.square().mean(-1, keepdim=True) + self.eps) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + # win kv & topk_idxs + kv = self.wkv(x) + kv = self.kv_norm(kv) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + # FP8-simulate non-rope dims to match QAT; rope dims stay bf16 for positional precision + act_quant(kv[..., :-rd], 64, scale_fmt, scale_dtype, True) + topk_idxs = get_window_topk_idxs(win, bsz, seqlen, start_pos) + if self.compress_ratio: + offset = kv.size(1) if start_pos == 0 else win + if self.indexer is not None: + compress_topk_idxs = self.indexer(x, qr, start_pos, offset) + else: + compress_topk_idxs = get_compress_topk_idxs(ratio, bsz, seqlen, start_pos, offset) + topk_idxs = torch.cat([topk_idxs, compress_topk_idxs], dim=-1) + topk_idxs = topk_idxs.int() + + # compress kv & attn + if start_pos == 0: + if seqlen <= win: + self.kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.kv_cache[:bsz, cutoff: win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split([win - cutoff, cutoff], dim=1) + if self.compress_ratio: + if (kv_compress := self.compressor(x, start_pos)) is not None: + kv = torch.cat([kv, kv_compress], dim=1) + # We performed QAT here, kv could also use fp8 format, though current implementation uses bf16 + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + else: + self.kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + if self.compress_ratio: + self.compressor(x, start_pos) + o = sparse_attn(q, self.kv_cache[:bsz], self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + # o + o = o.view(bsz, seqlen, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + # NOTE: wo_a is FP8 in checkpoint; could do FP8 einsum here for better perf, + # but using BF16 for simplicity. + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class Gate(nn.Module): + """MoE gating: computes expert routing scores and selects top-k experts. + Supports hash-based routing (first n_hash_layers) where expert indices are + predetermined per token ID, and score-based routing (remaining layers).""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.dim = args.dim + self.topk = args.n_activated_experts + self.score_func = args.score_func + self.route_scale = args.route_scale + self.hash = layer_id < args.n_hash_layers + self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim)) + if self.hash: + self.tid2eid = nn.Parameter(torch.empty(args.vocab_size, args.n_activated_experts, dtype=torch.int32), requires_grad=False) + self.bias = None + else: + self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.float32)) + + def forward(self, x: torch.Tensor, input_ids: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: + scores = linear(x.float(), self.weight.float()) + if self.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + original_scores = scores + # Bias shifts scores for expert selection (topk) but does not affect routing weights. + if self.bias is not None: + scores = scores + self.bias + if self.hash: + indices = self.tid2eid[input_ids] + else: + indices = scores.topk(self.topk, dim=-1)[1] + weights = original_scores.gather(1, indices) + if self.score_func != "softmax": + weights /= weights.sum(dim=-1, keepdim=True) + weights *= self.route_scale + return weights, indices + + +class Expert(nn.Module): + """Single MoE expert: SwiGLU FFN (w1, w2, w3). Computation in float32 for stability.""" + def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0): + super().__init__() + self.w1 = Linear(dim, inter_dim, dtype=dtype) + self.w2 = Linear(inter_dim, dim, dtype=dtype) + self.w3 = Linear(dim, inter_dim, dtype=dtype) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor] = None) -> torch.Tensor: + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + +class MoE(nn.Module): + """Mixture-of-Experts: gate routes each token to top-k routed experts + 1 shared expert. + Experts are sharded across TP ranks; each rank handles n_routed_experts // world_size experts.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + assert args.n_routed_experts % world_size == 0, f"Number of experts must be divisible by world size (world_size={world_size})" + self.n_routed_experts = args.n_routed_experts + self.n_local_experts = args.n_routed_experts // world_size + self.n_activated_experts = args.n_activated_experts + self.experts_start_idx = rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.gate = Gate(layer_id, args) + expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None + self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim, dtype=expert_dtype, swiglu_limit=args.swiglu_limit) if self.experts_start_idx <= i < self.experts_end_idx else None + for i in range(self.n_routed_experts)]) + assert args.n_shared_experts == 1 + self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit) + + def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, input_ids.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + y[idx] += expert(x[idx], weights[idx, top, None]) + if world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + +class Block(nn.Module): + """Transformer block with Hyper-Connections (HC) mixing. + Instead of a simple residual, HC maintains `hc_mult` copies of the hidden state. + hc_pre: reduces hc copies -> 1 via learned weighted sum (pre-weights from Sinkhorn). + hc_post: expands 1 -> hc copies via learned post-weights + combination matrix.""" + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.norm_eps = args.norm_eps + self.attn = Attention(layer_id, args) + self.ffn = MoE(layer_id, args) + self.attn_norm = RMSNorm(args.dim, self.norm_eps) + self.ffn_norm = RMSNorm(args.dim, self.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + self.hc_sinkhorn_iters = args.hc_sinkhorn_iters + self.hc_eps = args.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_attn_scale = nn.Parameter(torch.empty(3)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3)) + + def hc_pre(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + # x: [b,s,hc,d], hc_fn: [mix_hc,hc*d], hc_scale: [3], hc_base: [mix_hc], y: [b,s,hc,d] + shape, dtype = x.size(), x.dtype + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + pre, post, comb = hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps) + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype), post, comb + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + # x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc], y: [b,s,hc,d] + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.type_as(x) + + def forward(self, x: torch.Tensor, start_pos: int, input_ids: Optional[torch.Tensor]) -> torch.Tensor: + 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(x, start_pos) + 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 + + +class ParallelHead(nn.Module): + + def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.part_vocab_size = (vocab_size // world_size) + # lm_head in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for easier computation of logits later. + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32)) + + def get_logits(self, x): + return F.linear(x[:, -1].float(), self.weight) + + def forward(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, norm: RMSNorm): + # x: [b,s,hc,d] + x = self.hc_head(x, hc_fn, hc_scale, hc_base) + logits = self.get_logits(norm(x)) + if world_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(world_size)] + dist.all_gather(all_logits, logits) + logits = torch.cat(all_logits, dim=-1) + return logits + + def hc_head(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + shape, dtype = x.size(), x.dtype + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.hc_eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype) + + +class MTPBlock(Block): + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__(layer_id, args) + self.e_proj = Linear(args.dim, args.dim) + self.h_proj = Linear(args.dim, args.dim) + self.enorm = RMSNorm(args.dim, args.norm_eps) + self.hnorm = RMSNorm(args.dim, args.norm_eps) + self.norm = RMSNorm(args.dim, args.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim)) + self.hc_head_base = nn.Parameter(torch.empty(hc_mult)) + self.hc_head_scale = nn.Parameter(torch.empty(1)) + self.embed: ParallelEmbedding = None + self.head: ParallelHead = None + + @torch.inference_mode() + def forward(self, x: torch.Tensor, start_pos: int, input_ids: torch.Tensor) -> torch.Tensor: + # x: [b,s,hc,d] + assert self.embed is not None and self.head is not None + e = self.embed(input_ids) + e = self.enorm(e) + x = self.hnorm(x) + x = self.e_proj(e).unsqueeze(2) + self.h_proj(x) + x = super().forward(x, start_pos, input_ids) + logits = self.head(x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm) + return logits + + +class Transformer(nn.Module): + """Full DeepSeek-V4 model: embed -> HC-expand -> N blocks -> HC-head -> logits. + Sets global state (world_size, rank, default_dtype, scale_fmt, scale_dtype) in __init__.""" + def __init__(self, args: ModelArgs): + global world_size, rank, default_dtype, scale_fmt, scale_dtype + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + scale_fmt = "ue8m0" if args.scale_dtype == "fp8" else args.scale_fmt + scale_dtype = torch.float8_e8m0fnu if args.scale_dtype == "fp8" else torch.float32 + super().__init__() + self.max_seq_len = args.max_seq_len + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.embed = ParallelEmbedding(args.vocab_size, args.dim) + self.layers = torch.nn.ModuleList() + for layer_id in range(args.n_layers): + self.layers.append(Block(layer_id, args)) + self.norm = RMSNorm(args.dim, self.norm_eps) + self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) + self.mtp = torch.nn.ModuleList() + for layer_id in range(args.n_mtp_layers): + self.mtp.append(MTPBlock(args.n_layers + layer_id, args)) + self.mtp[-1].embed = self.embed + self.mtp[-1].head = self.head + self.hc_mult = hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim)) + self.hc_head_base = nn.Parameter(torch.empty(hc_mult)) + self.hc_head_scale = nn.Parameter(torch.empty(1)) + + @torch.inference_mode() + def forward(self, input_ids: torch.Tensor, start_pos: int = 0): + h = self.embed(input_ids) + # Expand to hc_mult copies for Hyper-Connections + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + for layer in self.layers: + h = layer(h, start_pos, input_ids) + logits = self.head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm) + return logits + + +if __name__ == "__main__": + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(0) + args = ModelArgs(n_hash_layers=0) + x = torch.randint(0, args.vocab_size, (2, 128)) + model = Transformer(args) + + print(model(x).size()) + for i in range(128, 150): + print(i, model(x[:, 0:1], i).size()) + + h = torch.randn(2, 128, args.hc_mult, args.dim) + mtp = model.mtp[0] + print(mtp(h, 0, x).size()) + print(mtp(h[:, 0:1], 1, x[:, 0:1]).size()) diff --git a/TEMP/dsv4thing/requirements.txt b/TEMP/dsv4thing/requirements.txt new file mode 100644 index 00000000..7e1cc78c --- /dev/null +++ b/TEMP/dsv4thing/requirements.txt @@ -0,0 +1,5 @@ +torch>=2.10.0 +transformers>=5.0.0 +safetensors>=0.7.0 +fast_hadamard_transform +tilelang==0.1.8 \ No newline at end of file diff --git a/single_shot_inference.py b/single_shot_inference.py index 7f445d19..369b1d50 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -47,8 +47,16 @@ PROMPT = _args.prompt or "The capital of France is" NUM_GPUS = _args.num_gpus SEED = _args.seed VERBOSE = _args.verbose -THINK_START, THINK_END = 128821, 128822 -USER_TOKEN, ASSISTANT_TOKEN = 128803, 128804 +# Special token IDs — derived from official encoding module strings + tokenizer. +# Do NOT hardcode these; the encoding module defines the canonical token strings. +from encoding.deepseek_v4_encoding import ( + thinking_start_token as _THINK_START_STR, + thinking_end_token as _THINK_END_STR, + USER_SP_TOKEN as _USER_STR, + ASSISTANT_SP_TOKEN as _ASSISTANT_STR, + eos_token as _EOS_STR, + bos_token as _BOS_STR, +) FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) # ===================================================================== @@ -1461,6 +1469,14 @@ def main(): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) + # Derive special token IDs from official encoding strings + tokenizer. + # This is the ONLY source of truth — never hardcode these IDs. + THINK_START = tokenizer.convert_tokens_to_ids(_THINK_START_STR) + THINK_END = tokenizer.convert_tokens_to_ids(_THINK_END_STR) + USER_TOKEN = tokenizer.convert_tokens_to_ids(_USER_STR) + ASSISTANT_TOKEN = tokenizer.convert_tokens_to_ids(_ASSISTANT_STR) + bos = tokenizer.bos_token_id or 0 + # A1: Build explicit stop set — DSV4 uses special turn-end tokens beyond eos STOP_IDS = set() eos_id = tokenizer.eos_token_id @@ -1476,7 +1492,6 @@ def main(): print(f" Special tokens: {tokenizer.special_tokens_map}") print(f" THINK_START={THINK_START} THINK_END={THINK_END} USER={USER_TOKEN} ASST={ASSISTANT_TOKEN}") - bos = tokenizer.bos_token_id or 0 if _args.prefill_tokens: generated = [int(x) for x in _args.prefill_tokens.split(',')] else: @@ -1749,12 +1764,35 @@ def main(): print(f" L{li} {tag}: {dt_ms:.2f}ms") prev_t = t - out = tokenizer.decode(all_tokens, skip_special_tokens=True) - print(f"\n{'='*70}") - print(f"Input: '{PROMPT}'") - print(f"Output: '{out}'") - print(f"Total: {time.time()-t0:.1f}s") - print(f"{'='*70}") + out_raw = tokenizer.decode(all_tokens, skip_special_tokens=False) + # Use official DSV4 parser for structured output + try: + from encoding.deepseek_v4_encoding import parse_message_from_completion_text + # Find the assistant portion — after the last ASSISTANT token + assistant_start = out_raw.find(_ASSISTANT_STR) + if assistant_start >= 0: + assistant_text = out_raw[assistant_start + len(_ASSISTANT_STR):] + else: + assistant_text = out_raw + parsed = parse_message_from_completion_text(assistant_text, thinking_mode=_args.thinking_mode) + reasoning = parsed.get('reasoning', '') + content = parsed.get('content', '') + print(f"\n{'='*70}") + print(f"Input: '{PROMPT}'") + if reasoning: + print(f"Reasoning: {reasoning[:500]}{'...' if len(reasoning) > 500 else ''}") + print(f"Content: {content}") + print(f"Total: {time.time()-t0:.1f}s") + print(f"{'='*70}") + except Exception as e: + # Fallback: raw decode (shouldn't happen with correct output) + out = tokenizer.decode(all_tokens, skip_special_tokens=True) + print(f"\n{'='*70}") + print(f"Input: '{PROMPT}'") + print(f"Output (raw): '{out}'") + print(f"Parse error: {e}") + print(f"Total: {time.time()-t0:.1f}s") + print(f"{'='*70}") if __name__ == "__main__": main()