biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Canonical encoding: derive special token IDs from official encoding module + tokenizer

Browse Source 
- Remove hardcoded THINK_START/THINK_END/USER_TOKEN/ASSISTANT_TOKEN IDs
- Import token strings from encoding.deepseek_v4_encoding (official source)
- Resolve IDs via tokenizer.convert_tokens_to_ids() at runtime
- Use parse_message_from_completion_text() for structured output parsing
- No more hand-rolled prompt construction or hardcoded token IDs
- Clean up TEMP: replace old deepseek_v4_ref with dsv4thing.zip reference
This commit is contained in:
biondizzle
2026-06-03 10:23:02 +00:00
parent a86d6d90a5
commit 8cfc1cae58
38 changed files with 1799 additions and 11229 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
TEMP/CROSS_REFERENCE.md
View File

@@ -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.

BIN
TEMP/deepseek_v4.zip
View File

Binary file not shown.

30
TEMP/deepseek_v4_ref/deepseek_v4/__init__.py
View File

@@ -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",
]

2
TEMP/deepseek_v4_ref/deepseek_v4/amd/__init__.py
View File

@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

972
TEMP/deepseek_v4_ref/deepseek_v4/amd/model.py
View File

@@ -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()

509
TEMP/deepseek_v4_ref/deepseek_v4/amd/mtp.py
View File

@@ -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

856
TEMP/deepseek_v4_ref/deepseek_v4/amd/rocm.py
View File

@@ -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],
)

806
TEMP/deepseek_v4_ref/deepseek_v4/attention.py
View File

@@ -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)

2
TEMP/deepseek_v4_ref/deepseek_v4/common/__init__.py
View File

@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

28
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/__init__.py
View File

@@ -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",
]

594
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/cache_utils.py
View File

@@ -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)

666
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_compress_quant_cache.py
View File

@@ -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)

438
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_indexer_q.py
View File

@@ -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

318
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
View File

@@ -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,
)

203
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
View File

@@ -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

96
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/fused_qk_rmsnorm.py
View File

@@ -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

101
TEMP/deepseek_v4_ref/deepseek_v4/common/ops/save_partial_states.py
View File

@@ -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,
)

36
TEMP/deepseek_v4_ref/deepseek_v4/common/rope.py
View File

@@ -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,
)

380
TEMP/deepseek_v4_ref/deepseek_v4/compressor.py
View File

@@ -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,
)

2
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/__init__.py
View File

@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

424
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/flashmla.py
View File

@@ -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],
)

1476
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/model.py
View File

File diff suppressed because it is too large Load Diff

516
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/mtp.py
View File

@@ -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

13
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/__init__.py
View File

@@ -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.
"""

331
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
View File

@@ -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",
)

610
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
View File

@@ -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",
)

173
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/prepare_megamoe.py
View File

@@ -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,
)

1380
TEMP/deepseek_v4_ref/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
View File

File diff suppressed because it is too large Load Diff

158
TEMP/deepseek_v4_ref/deepseek_v4/quant_config.py
View File

@@ -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"

BIN
TEMP/dsv4thing.zip Normal file
View File

Binary file not shown.

26
TEMP/dsv4thing/README.md Normal file
View File

@@ -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`.

35
TEMP/dsv4thing/config.json Normal file
View File

@@ -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]
}

168
TEMP/dsv4thing/convert.py Normal file
View File

@@ -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)

155
TEMP/dsv4thing/generate.py Normal file
View File

@@ -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)

536
TEMP/dsv4thing/kernel.py Normal file
View File

@@ -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

827
TEMP/dsv4thing/model.py Normal file
View File

@@ -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())

5
TEMP/dsv4thing/requirements.txt Normal file
View File

@@ -0,0 +1,5 @@
torch>=2.10.0
transformers>=5.0.0
safetensors>=0.7.0
fast_hadamard_transform
tilelang==0.1.8

56
single_shot_inference.py
View File

@@ -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()