forked from Karylab-cklius/vllm
Woosuk Kwon
@@ -715,6 +715,7 @@ class CompilationConfig:
|
||||
"vllm::kda_attention",
|
||||
"vllm::sparse_attn_indexer",
|
||||
"vllm::rocm_aiter_sparse_attn_indexer",
|
||||
"vllm::monolithic_attn",
|
||||
]
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
|
||||
@@ -2,8 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Monolithic decoder layer for DeepSeek V3.2 on SM100 (Blackwell).
|
||||
Direct kernel calls, no module wrappers for norms.
|
||||
Gate weight inlined, FusedMoE kept for quantized expert kernels.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,24 +9,161 @@ from __future__ import annotations
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config import 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.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
|
||||
|
||||
from .allreduce_rms import AllReduceRMSParams, allreduce_add_rms_norm
|
||||
from .attention import MonolithicMLAAttention
|
||||
from .ops import fused_norm_rope, fused_q, rms_norm
|
||||
from .ops import fused_norm_rope, fused_q
|
||||
from .sparse_indexer import sparse_attn_indexer
|
||||
|
||||
_side_stream: torch.cuda.Stream | None = None
|
||||
|
||||
def monolithic_attn(
|
||||
positions: torch.Tensor,
|
||||
q_c: torch.Tensor,
|
||||
kv_c: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
index_k: torch.Tensor,
|
||||
index_weights: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> torch.Tensor:
|
||||
layer = get_forward_context().no_compile_layers[layer_name]
|
||||
attn = layer.attn
|
||||
mla = attn.mla_attn
|
||||
|
||||
attn_metadata = get_forward_context().attn_metadata
|
||||
if not isinstance(attn_metadata, dict):
|
||||
output.zero_()
|
||||
return output
|
||||
|
||||
mla_attn_metadata = attn_metadata.get(mla.layer_name)
|
||||
if mla_attn_metadata is None:
|
||||
output.zero_()
|
||||
return output
|
||||
|
||||
num_actual_toks = mla_attn_metadata.num_actual_tokens
|
||||
if num_actual_toks == 0:
|
||||
output.zero_()
|
||||
return output
|
||||
|
||||
# Step 2. fused norm + rope + cache writes
|
||||
slot_mapping = None
|
||||
indexer_k_cache = None
|
||||
mla_kv_cache = None
|
||||
mla_k_scale = None
|
||||
idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix)
|
||||
if idx_meta is not None:
|
||||
slot_mapping = idx_meta.slot_mapping
|
||||
indexer_k_cache = attn.indexer_k_cache.kv_cache
|
||||
mla_kv_cache = attn.mla_attn.kv_cache
|
||||
mla_k_scale = attn.mla_attn._k_scale
|
||||
|
||||
q_c = fused_norm_rope(
|
||||
positions,
|
||||
q_c,
|
||||
attn.q_a_layernorm_weight,
|
||||
layer.rms_norm_eps,
|
||||
kv_c,
|
||||
attn.kv_a_layernorm_weight,
|
||||
attn.rms_norm_eps,
|
||||
k_pe,
|
||||
attn.rotary_emb.cos_sin_cache,
|
||||
index_k,
|
||||
attn.indexer_k_norm.weight,
|
||||
attn.indexer_k_norm.bias,
|
||||
attn.rms_norm_eps,
|
||||
attn.indexer_rope_emb.cos_sin_cache,
|
||||
attn.topk_indices_buffer,
|
||||
slot_mapping=slot_mapping,
|
||||
indexer_k_cache=indexer_k_cache,
|
||||
mla_kv_cache=mla_kv_cache,
|
||||
mla_kv_cache_dtype=attn.mla_attn.kv_cache_dtype,
|
||||
mla_k_scale=mla_k_scale,
|
||||
)
|
||||
|
||||
# Step 3. q_c -> index_q, q
|
||||
step3_out = torch.mm(q_c, layer._fused_step3_q_w.T)
|
||||
index_q, q = step3_out.split(
|
||||
[layer._step3_index_q_dim, step3_out.shape[-1] - layer._step3_index_q_dim],
|
||||
dim=-1,
|
||||
)
|
||||
index_q = index_q.view(-1, attn.index_n_heads, attn.index_head_dim)
|
||||
q = q.view(-1, attn.num_local_heads, attn.qk_head_dim)
|
||||
|
||||
# Step 4. Q RoPE + W_UK_T absorption + FP8 packing
|
||||
q_nope, q_pe = q.split(
|
||||
[mla.qk_nope_head_dim, mla.qk_rope_head_dim],
|
||||
dim=-1,
|
||||
)
|
||||
q_nope = q_nope.transpose(0, 1)
|
||||
ql_nope = torch.bmm(q_nope, mla.W_UK_T)
|
||||
ql_nope = ql_nope.transpose(0, 1)
|
||||
|
||||
index_q_fp8, index_weights, mqa_q = fused_q(
|
||||
positions,
|
||||
q_pe,
|
||||
attn.rotary_emb.cos_sin_cache,
|
||||
index_q,
|
||||
attn.indexer_rope_emb.cos_sin_cache,
|
||||
ql_nope,
|
||||
mla._q_scale,
|
||||
index_weights,
|
||||
attn.indexer_softmax_scale,
|
||||
attn.index_n_heads**-0.5,
|
||||
)
|
||||
|
||||
# Steps 5-6. Sparse indexer + MLA sparse decode attention
|
||||
sparse_attn_indexer(
|
||||
attn.indexer_k_cache.prefix,
|
||||
attn.indexer_k_cache.kv_cache,
|
||||
index_q_fp8,
|
||||
index_weights,
|
||||
attn.topk_tokens,
|
||||
attn.index_head_dim,
|
||||
layer.max_model_len,
|
||||
layer.indexer_workspace_size,
|
||||
attn.topk_indices_buffer,
|
||||
)
|
||||
|
||||
mqa_q = mqa_q[:num_actual_toks]
|
||||
kv_cache = mla.kv_cache
|
||||
if mla.kv_cache_dtype.startswith("fp8") and mla.kv_cache_dtype != "fp8_ds_mla":
|
||||
kv_cache = kv_cache.view(torch.float8_e4m3fn)
|
||||
attn_out, _ = mla.impl.forward_mqa(mqa_q, kv_cache, mla_attn_metadata, mla)
|
||||
x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1)
|
||||
|
||||
out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim)
|
||||
out = out.transpose(0, 1)
|
||||
torch.bmm(x, mla.W_UV, out=out)
|
||||
return output
|
||||
|
||||
|
||||
def _get_side_stream() -> torch.cuda.Stream:
|
||||
"""Lazily created CUDA stream shared by all decoder layers."""
|
||||
global _side_stream
|
||||
if _side_stream is None:
|
||||
_side_stream = torch.cuda.Stream()
|
||||
return _side_stream
|
||||
def monolithic_attn_fake(
|
||||
positions: torch.Tensor,
|
||||
q_c: torch.Tensor,
|
||||
kv_c: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
index_k: torch.Tensor,
|
||||
index_weights: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> torch.Tensor:
|
||||
del positions, q_c, kv_c, k_pe, index_k, index_weights, layer_name
|
||||
return output
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="monolithic_attn",
|
||||
op_func=monolithic_attn,
|
||||
fake_impl=monolithic_attn_fake,
|
||||
mutates_args=["output"],
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
|
||||
class MonolithicDecoderLayer(nn.Module):
|
||||
@@ -45,9 +180,14 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
layer_idx: int,
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
prefix: str = "",
|
||||
fi_params: AllReduceRMSParams | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
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
|
||||
|
||||
self.layer_name = prefix
|
||||
self.layer_idx = layer_idx
|
||||
self.hidden_size = config.hidden_size
|
||||
self.rms_norm_eps = config.rms_norm_eps
|
||||
@@ -55,7 +195,6 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
self.kv_lora_rank = config.kv_lora_rank
|
||||
self.qk_rope_head_dim = config.qk_rope_head_dim
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self._fi_params = fi_params
|
||||
|
||||
cache_config = vllm_config.cache_config
|
||||
quant_config = vllm_config.quant_config
|
||||
@@ -63,13 +202,18 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
self.indexer_workspace_size = get_max_prefill_buffer_size(vllm_config)
|
||||
self.max_model_len = vllm_config.model_config.max_model_len
|
||||
|
||||
# LayerNorm weights (raw)
|
||||
# Use the regular vLLM RMSNorm modules so the compiler sees the
|
||||
# canonical residual-add + RMSNorm pattern.
|
||||
dtype = torch.get_default_dtype()
|
||||
self.input_layernorm_weight = nn.Parameter(
|
||||
torch.ones(config.hidden_size, dtype=dtype)
|
||||
self.input_layernorm = RMSNorm(
|
||||
hidden_size=config.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
dtype=dtype,
|
||||
)
|
||||
self.post_attention_layernorm_weight = nn.Parameter(
|
||||
torch.ones(config.hidden_size, dtype=dtype)
|
||||
self.post_attention_layernorm = RMSNorm(
|
||||
hidden_size=config.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Fused QKV A-projection lives inside self_attn namespace
|
||||
@@ -103,7 +247,6 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
topk_indices_buffer=topk_indices_buffer,
|
||||
prefix=f"{prefix}.self_attn",
|
||||
)
|
||||
self.attn.o_proj.reduce_results = False
|
||||
|
||||
# MoE or Dense MLP
|
||||
moe_layer_freq = getattr(config, "moe_layer_freq", 1)
|
||||
@@ -126,15 +269,12 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
self.mlp.skip_final_allreduce = True
|
||||
self.mlp.skip_scale_and_add = True
|
||||
else:
|
||||
self.mlp = DeepseekV2MLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
reduce_results=False,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
|
||||
@@ -183,56 +323,11 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
hidden_states: torch.Tensor,
|
||||
residual: torch.Tensor | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from vllm.forward_context import get_forward_context
|
||||
|
||||
fwd_ctx = get_forward_context()
|
||||
attn_metadata = fwd_ctx.attn_metadata
|
||||
mla = self.attn.mla_attn
|
||||
mla_attn_metadata = None
|
||||
slot_mapping = None
|
||||
if isinstance(attn_metadata, dict):
|
||||
idx_meta = attn_metadata[self.attn.indexer_k_cache.prefix]
|
||||
mla_attn_metadata = attn_metadata[mla.layer_name]
|
||||
# Indexer and MLA caches share the same block_size and track
|
||||
# the same requests, so their slot_mappings are identical.
|
||||
slot_mapping = idx_meta.slot_mapping
|
||||
|
||||
if slot_mapping is not None:
|
||||
indexer_k_cache = self.attn.indexer_k_cache.kv_cache
|
||||
mla_kv_cache = mla.kv_cache
|
||||
mla_k_scale = mla._k_scale
|
||||
else:
|
||||
indexer_k_cache = None
|
||||
mla_kv_cache = None
|
||||
mla_k_scale = None
|
||||
|
||||
# Input norm + residual
|
||||
# When fused_allreduce_rms is enabled, hidden_states arriving from
|
||||
# the previous layer is the *unreduced* MLP/MoE output. We fuse
|
||||
# AllReduce + residual-add + RMSNorm into a single kernel.
|
||||
if residual is None:
|
||||
# First layer: hidden_states is from embed_tokens (already
|
||||
# fully materialised), no allreduce needed.
|
||||
residual = hidden_states
|
||||
hidden_states = rms_norm(
|
||||
hidden_states, self.input_layernorm_weight, self.rms_norm_eps
|
||||
)
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states, residual = allreduce_add_rms_norm(
|
||||
hidden_states,
|
||||
residual,
|
||||
self.input_layernorm_weight,
|
||||
self.rms_norm_eps,
|
||||
self._fi_params,
|
||||
)
|
||||
|
||||
if not hasattr(self, "_fused_step1_hidden_w") or not hasattr(
|
||||
self, "_fused_step3_q_w"
|
||||
):
|
||||
raise RuntimeError(
|
||||
"Monolithic decoder fused weights are not initialized. "
|
||||
"Call fuse_indexer_weights() after weight loading."
|
||||
)
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
|
||||
# Step 1. hidden_states -> q_c, kv_c, k_pe, index_k, index_weights
|
||||
step1_out = torch.mm(hidden_states, self._fused_step1_hidden_w.T)
|
||||
@@ -241,202 +336,27 @@ class MonolithicDecoderLayer(nn.Module):
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
# Step 2. Q RMS norm
|
||||
# + KV RMS norm + KV RoPE + MLA cache write
|
||||
# + Index K layer norm + RoPE + FP8 quant + cache write
|
||||
# + Init topk indices
|
||||
q_c = fused_norm_rope(
|
||||
positions,
|
||||
# Q RMS norm
|
||||
q_c,
|
||||
self.attn.q_a_layernorm_weight,
|
||||
self.rms_norm_eps,
|
||||
# KV RMS norm
|
||||
kv_c,
|
||||
self.attn.kv_a_layernorm_weight,
|
||||
self.attn.rms_norm_eps,
|
||||
# KV RoPE
|
||||
k_pe,
|
||||
self.attn.rotary_emb.cos_sin_cache,
|
||||
# Index K layer norm + RoPE
|
||||
index_k,
|
||||
self.attn.indexer_k_norm.weight,
|
||||
self.attn.indexer_k_norm.bias,
|
||||
self.attn.rms_norm_eps,
|
||||
self.attn.indexer_rope_emb.cos_sin_cache,
|
||||
# Top k indices
|
||||
self.attn.topk_indices_buffer,
|
||||
# Fused cache writes (single slot_mapping for both caches)
|
||||
slot_mapping=slot_mapping,
|
||||
indexer_k_cache=indexer_k_cache,
|
||||
mla_kv_cache=mla_kv_cache,
|
||||
mla_kv_cache_dtype=self.attn.mla_attn.kv_cache_dtype,
|
||||
mla_k_scale=mla_k_scale,
|
||||
)
|
||||
|
||||
# Step 3. q_c -> index_q, q
|
||||
step3_out = torch.mm(q_c, self._fused_step3_q_w.T)
|
||||
index_q, q = step3_out.split(
|
||||
[self._step3_index_q_dim, step3_out.shape[-1] - self._step3_index_q_dim],
|
||||
dim=-1,
|
||||
)
|
||||
index_q = index_q.view(-1, self.attn.index_n_heads, self.attn.index_head_dim)
|
||||
q = q.view(-1, self.attn.num_local_heads, self.attn.qk_head_dim)
|
||||
|
||||
# Step 4. Second fused stage:
|
||||
# Q RoPE + Index Q RoPE + Index Q FP8 + index-weight scaling
|
||||
# + W_UK_T absorption + MQA FP8 query packing.
|
||||
q_nope, q_pe = q.split([mla.qk_nope_head_dim, mla.qk_rope_head_dim], dim=-1)
|
||||
q_nope = q_nope.transpose(0, 1)
|
||||
ql_nope = torch.bmm(q_nope, mla.W_UK_T)
|
||||
ql_nope = ql_nope.transpose(0, 1)
|
||||
|
||||
assert mla.kv_cache_dtype.startswith("fp8")
|
||||
assert mla.impl.supports_quant_query_input
|
||||
|
||||
index_q_fp8, index_weights, mqa_q = fused_q(
|
||||
positions,
|
||||
q_pe,
|
||||
self.attn.rotary_emb.cos_sin_cache,
|
||||
index_q,
|
||||
self.attn.indexer_rope_emb.cos_sin_cache,
|
||||
ql_nope,
|
||||
mla._q_scale,
|
||||
index_weights,
|
||||
self.attn.indexer_softmax_scale,
|
||||
self.attn.index_n_heads**-0.5,
|
||||
)
|
||||
|
||||
# Step 5. Sparse indexer.
|
||||
# The FP8 quant + cache write for index_k is already done in
|
||||
# fused_norm_rope (step 2) when slot_mapping is available.
|
||||
sparse_attn_indexer(
|
||||
self.attn.indexer_k_cache.prefix,
|
||||
self.attn.indexer_k_cache.kv_cache,
|
||||
index_q_fp8,
|
||||
index_weights,
|
||||
self.attn.topk_tokens,
|
||||
self.attn.index_head_dim,
|
||||
self.max_model_len,
|
||||
self.indexer_workspace_size,
|
||||
self.attn.topk_indices_buffer,
|
||||
)
|
||||
|
||||
# Step 6. MLA sparse decode attention (inlined).
|
||||
# The KV cache update was already done in fused_norm_rope (step 2).
|
||||
# Steps 2-6. Combined: fused norm/rope + Q projections + sparse MLA.
|
||||
mla = self.attn.mla_attn
|
||||
output_shape = (hidden_states.shape[0], mla.num_heads * mla.v_head_dim)
|
||||
output_dtype = mla.W_UV.dtype
|
||||
if mla_attn_metadata is None or slot_mapping is None:
|
||||
attn_out = torch.zeros(
|
||||
output_shape,
|
||||
dtype=output_dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
else:
|
||||
num_actual_toks = mla_attn_metadata.num_actual_tokens
|
||||
mqa_q = mqa_q[:num_actual_toks]
|
||||
kv_cache = mla.kv_cache
|
||||
if (
|
||||
mla.kv_cache_dtype.startswith("fp8")
|
||||
and mla.kv_cache_dtype != "fp8_ds_mla"
|
||||
):
|
||||
kv_cache = kv_cache.view(torch.float8_e4m3fn)
|
||||
attn_out = torch.empty(
|
||||
output_shape,
|
||||
dtype=output_dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
attn_out = torch.ops.vllm.monolithic_attn(
|
||||
positions,
|
||||
q_c,
|
||||
kv_c,
|
||||
k_pe,
|
||||
index_k,
|
||||
index_weights,
|
||||
attn_out,
|
||||
self.layer_name,
|
||||
)
|
||||
|
||||
attn_out, _ = mla.impl.forward_mqa(
|
||||
mqa_q,
|
||||
kv_cache,
|
||||
mla_attn_metadata,
|
||||
mla,
|
||||
)
|
||||
|
||||
output = torch.empty(
|
||||
output_shape,
|
||||
dtype=output_dtype,
|
||||
device=kv_cache.device,
|
||||
)
|
||||
x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1)
|
||||
out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim)
|
||||
out = out.transpose(0, 1)
|
||||
torch.bmm(x, mla.W_UV, out=out)
|
||||
attn_out = output
|
||||
|
||||
# Step 7. Output projection (AllReduce disabled when fused).
|
||||
hidden_states, _ = self.attn.o_proj(attn_out)
|
||||
|
||||
# Post-attn norm + residual
|
||||
# Fuse the o_proj AllReduce with post-attention RMSNorm.
|
||||
hidden_states, residual = allreduce_add_rms_norm(
|
||||
hidden_states,
|
||||
residual,
|
||||
self.post_attention_layernorm_weight,
|
||||
self.rms_norm_eps,
|
||||
self._fi_params,
|
||||
)
|
||||
|
||||
# MLP / MoE
|
||||
# When fused_allreduce_rms is enabled, the MLP/MoE AllReduce is
|
||||
# deferred — it will be fused with the next layer's input norm.
|
||||
if self.is_moe:
|
||||
# MoE returns raw (shared_output, routed_output) without
|
||||
# applying scale + add. torch.compile fuses the elementwise ops.
|
||||
shared_output, routed_output = self.mlp(hidden_states)
|
||||
hidden_states = scale_and_add(
|
||||
routed_output, self.routed_scaling_factor, shared_output
|
||||
)
|
||||
else:
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
|
||||
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
return hidden_states, residual
|
||||
|
||||
def fuse_shared_expert_act_quant(self) -> None:
|
||||
"""Fuse SiLU-and-Mul + NVFP4 quantize in the shared expert MLP.
|
||||
|
||||
Replaces the shared expert forward so that the activation and the
|
||||
FP4 input quantisation happen in a single kernel. The monolithic
|
||||
path always runs on Blackwell with FLASHINFER_CUTLASS for linear
|
||||
NVFP4, so we hard-code that backend.
|
||||
"""
|
||||
if not self.is_moe:
|
||||
return
|
||||
shared_experts = self.mlp.shared_experts
|
||||
if shared_experts is None:
|
||||
return
|
||||
|
||||
from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptNvFp4LinearMethod,
|
||||
)
|
||||
|
||||
if not isinstance(
|
||||
shared_experts.down_proj.quant_method, ModelOptNvFp4LinearMethod
|
||||
):
|
||||
return
|
||||
|
||||
from vllm._custom_ops import silu_and_mul_nvfp4_quant
|
||||
from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm
|
||||
|
||||
dp = shared_experts.down_proj
|
||||
|
||||
def _fused_forward(x: torch.Tensor) -> torch.Tensor:
|
||||
gate_up, _ = shared_experts.gate_up_proj(x)
|
||||
x_fp4, x_bs = silu_and_mul_nvfp4_quant(gate_up, dp.input_global_scale_inv)
|
||||
out = flashinfer_scaled_fp4_mm(
|
||||
x_fp4,
|
||||
dp.weight,
|
||||
x_bs,
|
||||
dp.weight_scale,
|
||||
dp.alpha,
|
||||
gate_up.dtype,
|
||||
backend="cutlass",
|
||||
)
|
||||
return out
|
||||
|
||||
shared_experts.forward = _fused_forward
|
||||
|
||||
|
||||
@torch.compile
|
||||
def scale_and_add(x: torch.Tensor, scale: float, y: torch.Tensor) -> torch.Tensor:
|
||||
orig_dtype = x.dtype
|
||||
x = x.to(torch.float32)
|
||||
x = x * scale
|
||||
z = x + y.to(torch.float32)
|
||||
return z.to(orig_dtype)
|
||||
|
||||
@@ -10,28 +10,24 @@ from collections.abc import Iterable
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_gather,
|
||||
)
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .allreduce_rms import (
|
||||
AllReduceRMSParams,
|
||||
allreduce_add_rms_norm,
|
||||
should_use_allreduce_rms,
|
||||
)
|
||||
from .decoder_layer import MonolithicDecoderLayer
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@support_torch_compile
|
||||
class MonolithicDeepseekV32Model(nn.Module):
|
||||
"""Transformer backbone."""
|
||||
|
||||
@@ -43,13 +39,6 @@ class MonolithicDeepseekV32Model(nn.Module):
|
||||
quant_config = vllm_config.quant_config
|
||||
self.config = config
|
||||
self.device = current_platform.device_type
|
||||
self.fi_params: AllReduceRMSParams | None = None
|
||||
if should_use_allreduce_rms():
|
||||
self.fi_params = AllReduceRMSParams(vllm_config, config.hidden_size)
|
||||
logger.info(
|
||||
"Enabling fused AllReduce + RMSNorm in monolithic "
|
||||
"DeepSeek V3.2 decoder layers."
|
||||
)
|
||||
|
||||
topk_tokens = config.index_topk
|
||||
self.topk_indices_buffer = torch.empty(
|
||||
@@ -74,16 +63,16 @@ class MonolithicDeepseekV32Model(nn.Module):
|
||||
layer_idx=i,
|
||||
topk_indices_buffer=self.topk_indices_buffer,
|
||||
prefix=f"{prefix}.layers.{i}",
|
||||
fi_params=self.fi_params,
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm_weight = nn.Parameter(
|
||||
torch.ones(config.hidden_size, dtype=torch.get_default_dtype())
|
||||
self.norm = RMSNorm(
|
||||
hidden_size=config.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
dtype=torch.get_default_dtype(),
|
||||
)
|
||||
self.rms_norm_eps = config.rms_norm_eps
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -94,15 +83,7 @@ class MonolithicDeepseekV32Model(nn.Module):
|
||||
residual = None
|
||||
for layer in self.layers:
|
||||
hidden_states, residual = layer(positions, hidden_states, residual)
|
||||
# After the last layer, hidden_states is unreduced when fused.
|
||||
# AllReduce before the final norm.
|
||||
hidden_states, _ = allreduce_add_rms_norm(
|
||||
hidden_states,
|
||||
residual,
|
||||
self.norm_weight,
|
||||
self.rms_norm_eps,
|
||||
self.fi_params,
|
||||
)
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
return hidden_states
|
||||
|
||||
|
||||
@@ -134,6 +115,7 @@ class DeepseekV32MonolithicForCausalLM(nn.Module):
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.lm_head" if prefix else "lm_head",
|
||||
)
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
self.num_redundant_experts = 0
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
@@ -149,10 +131,7 @@ class DeepseekV32MonolithicForCausalLM(nn.Module):
|
||||
return self.model(input_ids, positions)
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
|
||||
logits = self.lm_head.quant_method.apply(self.lm_head, hidden_states)
|
||||
logits = tensor_model_parallel_all_gather(logits)
|
||||
logits = logits[..., : self.config.vocab_size]
|
||||
return logits
|
||||
return self.logits_processor(self.lm_head, hidden_states)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Delegate to the original DeepSeek V2 weight loader.
|
||||
@@ -176,21 +155,15 @@ class DeepseekV32MonolithicForCausalLM(nn.Module):
|
||||
# Fuse indexer linear weights after loading.
|
||||
for layer in self.model.layers:
|
||||
layer.fuse_indexer_weights()
|
||||
layer.fuse_shared_expert_act_quant()
|
||||
|
||||
return loaded
|
||||
|
||||
def _remap(self, name: str) -> str:
|
||||
"""Remap only names that differ from original model structure."""
|
||||
# Only remap layernorms (raw params) and indexer (underscore prefix).
|
||||
# Only remap layernorms and indexer (underscore prefix).
|
||||
# Everything else (fused_qkv_a_proj, experts, gate, etc.) uses the
|
||||
# same module paths as the original model.
|
||||
replacements = [
|
||||
("input_layernorm.weight", "input_layernorm_weight"),
|
||||
(
|
||||
"post_attention_layernorm.weight",
|
||||
"post_attention_layernorm_weight",
|
||||
),
|
||||
(
|
||||
"self_attn.q_a_layernorm.weight",
|
||||
"attn.q_a_layernorm_weight",
|
||||
@@ -203,7 +176,6 @@ class DeepseekV32MonolithicForCausalLM(nn.Module):
|
||||
("self_attn.kv_b_proj", "attn.kv_b_proj"),
|
||||
("self_attn.o_proj", "attn.o_proj"),
|
||||
("self_attn.indexer.", "attn.indexer_"),
|
||||
("model.norm.weight", "model.norm_weight"),
|
||||
]
|
||||
for old, new in replacements:
|
||||
if old in name:
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Triton fused SiLU-and-Mul + NVFP4 (e2m1) quantization.
|
||||
|
||||
Uses PTX inline assembly to match the CUDA kernel's fast-math intrinsics
|
||||
and e2m1 conversion bitwise-exactly:
|
||||
- rcp.approx.ftz.f32 (reciprocal_approximate_ftz)
|
||||
- ex2.approx.ftz.f32 (__expf via base-2 fast exp)
|
||||
- cvt.rn.satfinite.e2m1x2.f32 (float32 → packed e2m1)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from vllm._custom_ops import create_fp4_output_tensors
|
||||
|
||||
# ── PTX helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _rcp_approx_ftz(x):
|
||||
"""rcp.approx.ftz.f32 — fast reciprocal (~1 mantissa-bit precision)."""
|
||||
return tl.inline_asm_elementwise(
|
||||
asm="rcp.approx.ftz.f32 $0, $1;",
|
||||
constraints="=f,f",
|
||||
args=[x],
|
||||
dtype=tl.float32,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _ex2_approx_ftz(x):
|
||||
"""ex2.approx.ftz.f32 — fast 2^x (used to implement __expf)."""
|
||||
return tl.inline_asm_elementwise(
|
||||
asm="ex2.approx.ftz.f32 $0, $1;",
|
||||
constraints="=f,f",
|
||||
args=[x],
|
||||
dtype=tl.float32,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _cvt_e2m1x2(even, odd):
|
||||
"""cvt.rn.satfinite.e2m1x2.f32 — pack two f32 into one byte of e2m1."""
|
||||
return tl.inline_asm_elementwise(
|
||||
asm=(
|
||||
"{ .reg .b8 tmp;"
|
||||
" cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1;"
|
||||
" cvt.u32.u8 $0, tmp; }"
|
||||
),
|
||||
constraints="=r,f,f",
|
||||
args=[even, odd],
|
||||
dtype=tl.int32,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
|
||||
|
||||
# ── kernel ───────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _silu_mul_nvfp4_quant_kernel(
|
||||
input_ptr, # [M, 2*N] bfloat16
|
||||
output_ptr, # [M, N//2] uint8 (packed e2m1 pairs)
|
||||
sf_out_ptr, # swizzled scale factors (viewed as uint8*)
|
||||
sf_scale_ptr, # global scale (float32 scalar)
|
||||
M,
|
||||
N, # output dim (= input last dim / 2)
|
||||
stride_in_m,
|
||||
stride_out_m,
|
||||
num_k_tiles, # ceil(N / 64)
|
||||
PAIRS: tl.constexpr, # = 8 (16 values → 8 packed bytes)
|
||||
):
|
||||
"""One program = one (row, quantisation-group-of-16) pair."""
|
||||
LOG2E: tl.constexpr = 1.4426950408889634
|
||||
|
||||
pid_m = tl.program_id(0)
|
||||
pid_g = tl.program_id(1)
|
||||
|
||||
col_base = pid_g * (PAIRS * 2) # = pid_g * 16
|
||||
pair_off = tl.arange(0, PAIRS) # [0..7]
|
||||
even_cols = col_base + 2 * pair_off # [0,2,4,...,14] + col_base
|
||||
odd_cols = even_cols + 1 # [1,3,5,...,15] + col_base
|
||||
|
||||
row_off = pid_m * stride_in_m
|
||||
|
||||
# ── load gate / up (BF16 → float32) ──
|
||||
gate_e = tl.load(input_ptr + row_off + even_cols).to(tl.float32)
|
||||
up_e = tl.load(input_ptr + row_off + N + even_cols).to(tl.float32)
|
||||
gate_o = tl.load(input_ptr + row_off + odd_cols).to(tl.float32)
|
||||
up_o = tl.load(input_ptr + row_off + N + odd_cols).to(tl.float32)
|
||||
|
||||
# ── fast-math SiLU: __fdividef(x, 1 + __expf(-x)) ──
|
||||
# __expf(y) = ex2.approx.ftz(y * log2(e))
|
||||
# __fdividef(a, b) = a * rcp.approx.ftz(b)
|
||||
exp_e = _ex2_approx_ftz((-gate_e) * LOG2E)
|
||||
exp_o = _ex2_approx_ftz((-gate_o) * LOG2E)
|
||||
silu_e = gate_e * _rcp_approx_ftz(1.0 + exp_e)
|
||||
silu_o = gate_o * _rcp_approx_ftz(1.0 + exp_o)
|
||||
|
||||
res_e = silu_e * up_e
|
||||
res_o = silu_o * up_o
|
||||
|
||||
# ── BF16 round-trip (matches compute_silu_mul returning PackedVec<BF16>) ──
|
||||
res_e = res_e.to(tl.bfloat16).to(tl.float32)
|
||||
res_o = res_o.to(tl.bfloat16).to(tl.float32)
|
||||
|
||||
# ── per-group abs-max (over all 16 values) ──
|
||||
amax = tl.maximum(tl.max(tl.abs(res_e)), tl.max(tl.abs(res_o)))
|
||||
|
||||
# ── scale factor: sf_scale * (amax * rcp.approx.ftz(6.0)) ──
|
||||
sf_scale_val = tl.load(sf_scale_ptr).to(tl.float32)
|
||||
sf_raw = sf_scale_val * (amax * _rcp_approx_ftz(6.0))
|
||||
|
||||
# quantise sf → float8_e4m3fn → float32 (exact round-trip)
|
||||
sf_fp8 = sf_raw.to(tl.float8e4nv)
|
||||
sf_rounded = sf_fp8.to(tl.float32)
|
||||
|
||||
# ── write scale byte into swizzled layout ──
|
||||
sf_col = pid_g
|
||||
m_tile = pid_m // 128
|
||||
outer_m = pid_m % 32
|
||||
inner_m = (pid_m // 32) % 4
|
||||
k_tile = sf_col // 4
|
||||
inner_k = sf_col % 4
|
||||
sf_offset = (
|
||||
(m_tile * num_k_tiles + k_tile) * 512 + outer_m * 16 + inner_m * 4 + inner_k
|
||||
)
|
||||
sf_byte = sf_fp8.to(tl.uint8, bitcast=True)
|
||||
tl.store(sf_out_ptr + sf_offset, sf_byte)
|
||||
|
||||
# ── output scale: rcp(sf_rounded * rcp(sf_scale)) ──
|
||||
rcp_sf_scale = _rcp_approx_ftz(sf_scale_val)
|
||||
out_scale = tl.where(
|
||||
sf_rounded != 0.0,
|
||||
_rcp_approx_ftz(sf_rounded * rcp_sf_scale),
|
||||
0.0,
|
||||
)
|
||||
|
||||
# ── scale values ──
|
||||
scaled_e = res_e * out_scale
|
||||
scaled_o = res_o * out_scale
|
||||
|
||||
# ── PTX e2m1 conversion + pack ──
|
||||
packed = _cvt_e2m1x2(scaled_e, scaled_o).to(tl.uint8)
|
||||
|
||||
# ── store 8 packed bytes ──
|
||||
out_off = pid_m * stride_out_m + col_base // 2 + pair_off
|
||||
tl.store(output_ptr + out_off, packed)
|
||||
|
||||
|
||||
# ── Python entry point ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def silu_and_mul_nvfp4_quant(
|
||||
input: torch.Tensor,
|
||||
input_global_scale: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Drop-in replacement for vllm._custom_ops.silu_and_mul_nvfp4_quant."""
|
||||
if input.ndim == 1:
|
||||
input = input.unsqueeze(0)
|
||||
else:
|
||||
input = input.reshape(-1, input.shape[-1])
|
||||
|
||||
M, two_N = input.shape
|
||||
N = two_N // 2
|
||||
assert N % 16 == 0, f"N must be a multiple of 16, got {N}"
|
||||
|
||||
output, output_scale = create_fp4_output_tensors(
|
||||
M,
|
||||
N,
|
||||
input.device,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
|
||||
num_groups = N // 16
|
||||
num_k_tiles = (num_groups + 3) // 4 # ceil(N / 64)
|
||||
|
||||
grid = (M, num_groups)
|
||||
_silu_mul_nvfp4_quant_kernel[grid](
|
||||
input,
|
||||
output,
|
||||
output_scale.view(torch.uint8),
|
||||
input_global_scale,
|
||||
M,
|
||||
N,
|
||||
input.stride(0),
|
||||
output.stride(0),
|
||||
num_k_tiles,
|
||||
PAIRS=8,
|
||||
)
|
||||
|
||||
output_scale = output_scale.view(torch.float8_e4m3fn)
|
||||
return output, output_scale
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Triton implementation of concat_and_cache_mla_kernel.
|
||||
|
||||
Concatenates compressed KV (kv_c) and positional-encoding keys (k_pe)
|
||||
into a packed MLA KV cache, with optional FP8 quantization.
|
||||
|
||||
Cache layout per slot: [kv_lora_rank | pe_dim] elements, contiguous.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _concat_and_cache_mla_kernel(
|
||||
# Pointers
|
||||
kv_c_ptr, # [num_tokens, kv_lora_rank]
|
||||
k_pe_ptr, # [num_tokens, pe_dim]
|
||||
kv_cache_ptr, # [num_blocks, block_size, entry_stride]
|
||||
slot_mapping_ptr, # [num_tokens], int64
|
||||
scale_ptr, # [1], float32
|
||||
# Strides (in elements, not bytes)
|
||||
kv_c_stride,
|
||||
k_pe_stride,
|
||||
block_stride,
|
||||
entry_stride,
|
||||
# Dimensions
|
||||
kv_lora_rank: tl.constexpr,
|
||||
pe_dim: tl.constexpr,
|
||||
block_size,
|
||||
# Mode
|
||||
QUANTIZE_FP8: tl.constexpr,
|
||||
KV_C_BLOCK: tl.constexpr, # next_power_of_2(kv_lora_rank)
|
||||
PE_BLOCK: tl.constexpr, # next_power_of_2(pe_dim)
|
||||
):
|
||||
"""Copy one token's kv_c and k_pe into the KV cache.
|
||||
|
||||
Grid: (num_tokens,)
|
||||
"""
|
||||
token_idx = tl.program_id(0)
|
||||
|
||||
slot_idx = tl.load(slot_mapping_ptr + token_idx)
|
||||
# slot_idx == -1 for padded tokens
|
||||
if slot_idx < 0:
|
||||
return
|
||||
|
||||
block_idx = slot_idx // block_size
|
||||
block_offset = slot_idx % block_size
|
||||
dst_base = block_idx * block_stride + block_offset * entry_stride
|
||||
|
||||
# -- Copy kv_c (kv_lora_rank elements) --
|
||||
kv_c_offsets = tl.arange(0, KV_C_BLOCK)
|
||||
kv_c_mask = kv_c_offsets < kv_lora_rank
|
||||
kv_c_vals = tl.load(
|
||||
kv_c_ptr + token_idx * kv_c_stride + kv_c_offsets,
|
||||
mask=kv_c_mask,
|
||||
other=0.0,
|
||||
)
|
||||
if QUANTIZE_FP8:
|
||||
scale = tl.load(scale_ptr)
|
||||
kv_c_fp8 = (kv_c_vals.to(tl.float32) / scale).to(tl.float8e4nv)
|
||||
tl.store(kv_cache_ptr + dst_base + kv_c_offsets, kv_c_fp8, mask=kv_c_mask)
|
||||
else:
|
||||
tl.store(kv_cache_ptr + dst_base + kv_c_offsets, kv_c_vals, mask=kv_c_mask)
|
||||
|
||||
# -- Copy k_pe (pe_dim elements) at offset kv_lora_rank --
|
||||
pe_offsets = tl.arange(0, PE_BLOCK)
|
||||
pe_mask = pe_offsets < pe_dim
|
||||
k_pe_vals = tl.load(
|
||||
k_pe_ptr + token_idx * k_pe_stride + pe_offsets,
|
||||
mask=pe_mask,
|
||||
other=0.0,
|
||||
)
|
||||
if QUANTIZE_FP8:
|
||||
k_pe_fp8 = (k_pe_vals.to(tl.float32) / scale).to(tl.float8e4nv)
|
||||
tl.store(
|
||||
kv_cache_ptr + dst_base + kv_lora_rank + pe_offsets,
|
||||
k_pe_fp8,
|
||||
mask=pe_mask,
|
||||
)
|
||||
else:
|
||||
tl.store(
|
||||
kv_cache_ptr + dst_base + kv_lora_rank + pe_offsets,
|
||||
k_pe_vals,
|
||||
mask=pe_mask,
|
||||
)
|
||||
|
||||
|
||||
def concat_and_cache_mla(
|
||||
kv_c: torch.Tensor, # [num_tokens, kv_lora_rank]
|
||||
k_pe: torch.Tensor, # [num_tokens, pe_dim]
|
||||
kv_cache: torch.Tensor, # [num_blocks, block_size, entry_stride]
|
||||
slot_mapping: torch.Tensor, # [num_tokens]
|
||||
kv_cache_dtype: str,
|
||||
scale: torch.Tensor, # [1]
|
||||
) -> None:
|
||||
"""Concatenate kv_c and k_pe into an MLA KV cache.
|
||||
|
||||
Drop-in replacement for the CUDA ``concat_and_cache_mla`` op.
|
||||
|
||||
Args:
|
||||
kv_c: Compressed KV of shape ``[num_tokens, kv_lora_rank]``.
|
||||
k_pe: Positional-encoding keys of shape ``[num_tokens, pe_dim]``.
|
||||
kv_cache: Cache of shape ``[num_blocks, block_size, entry_stride]``.
|
||||
slot_mapping: Slot index per token (``-1`` = padding).
|
||||
kv_cache_dtype: ``"auto"`` for direct copy, ``"fp8_e4m3"`` for
|
||||
FP8 quantization with the provided ``scale``.
|
||||
scale: Scalar FP8 scale factor (used only when
|
||||
``kv_cache_dtype != "auto"``).
|
||||
"""
|
||||
num_tokens = slot_mapping.shape[0]
|
||||
kv_lora_rank = kv_c.shape[1]
|
||||
pe_dim = k_pe.shape[1]
|
||||
block_size = kv_cache.shape[1]
|
||||
|
||||
kv_c_stride = kv_c.stride(0)
|
||||
k_pe_stride = k_pe.stride(0)
|
||||
block_stride = kv_cache.stride(0)
|
||||
entry_stride = kv_cache.stride(1)
|
||||
|
||||
quantize_fp8 = kv_cache_dtype != "auto"
|
||||
|
||||
# When quantizing to FP8, view the (uint8) cache as float8_e4m3fn so
|
||||
# that Triton's store sees an fp8-typed pointer.
|
||||
if quantize_fp8 and kv_cache.dtype == torch.uint8:
|
||||
kv_cache_view = kv_cache.view(torch.float8_e4m3fn)
|
||||
else:
|
||||
kv_cache_view = kv_cache
|
||||
|
||||
grid = (num_tokens,)
|
||||
_concat_and_cache_mla_kernel[grid](
|
||||
kv_c,
|
||||
k_pe,
|
||||
kv_cache_view,
|
||||
slot_mapping,
|
||||
scale,
|
||||
kv_c_stride,
|
||||
k_pe_stride,
|
||||
block_stride,
|
||||
entry_stride,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
pe_dim=pe_dim,
|
||||
block_size=block_size,
|
||||
QUANTIZE_FP8=quantize_fp8,
|
||||
KV_C_BLOCK=triton.next_power_of_2(kv_lora_rank),
|
||||
PE_BLOCK=triton.next_power_of_2(pe_dim),
|
||||
)
|
||||
Reference in New Issue
Block a user