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