first commit
This commit is contained in:
871
upstream/vllm/model_executor/layers/quantization/fp8.py
Normal file
871
upstream/vllm/model_executor/layers/quantization/fp8.py
Normal file
@@ -0,0 +1,871 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
import vllm.envs as envs
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
init_fp8_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.scaled_mm import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
MarlinFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
FusedMoEMethodBase,
|
||||
FusedMoeWeightScaleSupported,
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
UnquantizedFusedMoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
Fp8MoeBackend,
|
||||
convert_to_fp8_moe_kernel_format,
|
||||
make_fp8_moe_kernel,
|
||||
make_fp8_moe_quant_config,
|
||||
select_fp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
LinearBase,
|
||||
LinearMethodBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_input_scale,
|
||||
create_fp8_scale_parameter,
|
||||
create_fp8_weight_parameter,
|
||||
process_fp8_input_tensor_strategy_moe,
|
||||
process_fp8_weight_tensor_strategy,
|
||||
process_fp8_weight_tensor_strategy_moe,
|
||||
validate_fp8_block_shape,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
get_marlin_input_dtype,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
create_fp8_quant_key,
|
||||
is_layer_skipped,
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
cutlass_block_fp8_supported,
|
||||
cutlass_fp8_supported,
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
BlockQuantScaleParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_ppu():
|
||||
from vllm.utils.ppu_deep_gemm import (
|
||||
is_deep_gemm_supported,
|
||||
)
|
||||
else:
|
||||
from vllm.utils.deep_gemm import (
|
||||
is_deep_gemm_supported,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
ACTIVATION_SCHEMES = ["static", "dynamic"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Fp8Config(QuantizationConfig):
|
||||
"""Config class for FP8."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
is_checkpoint_fp8_serialized: bool = False,
|
||||
activation_scheme: str = "dynamic",
|
||||
ignored_layers: list[str] | None = None,
|
||||
weight_block_size: list[int] | None = None,
|
||||
store_dtype: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
|
||||
|
||||
if activation_scheme not in ACTIVATION_SCHEMES:
|
||||
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
|
||||
self.activation_scheme = activation_scheme
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.store_dtype = store_dtype
|
||||
if weight_block_size is not None:
|
||||
if not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"The block-wise quantization only supports fp8-serialized "
|
||||
"checkpoint for now."
|
||||
)
|
||||
if len(weight_block_size) != 2:
|
||||
raise ValueError(
|
||||
"The quantization block size of weight must have 2 "
|
||||
f"dimensions, but got {len(weight_block_size)} dimensions"
|
||||
)
|
||||
if activation_scheme != "dynamic":
|
||||
raise ValueError(
|
||||
"The block-wise quantization only supports "
|
||||
"dynamic activation scheme for now, but got "
|
||||
f"{activation_scheme} activation scheme."
|
||||
)
|
||||
self.weight_block_size = weight_block_size
|
||||
self.use_deep_gemm: bool | None = None
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> QuantizationMethods:
|
||||
return "fp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 75
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
||||
if self.ignored_layers is not None:
|
||||
self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "Fp8Config":
|
||||
quant_method = cls.get_from_keys(config, ["quant_method"])
|
||||
is_checkpoint_fp8_serialized = "fp8" in quant_method
|
||||
activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
|
||||
ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
|
||||
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
||||
store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None)
|
||||
if not ignored_layers:
|
||||
ignored_layers = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(
|
||||
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
|
||||
activation_scheme=activation_scheme,
|
||||
ignored_layers=ignored_layers,
|
||||
weight_block_size=weight_block_size,
|
||||
store_dtype=store_dtype,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> "QuantizeMethodBase | None":
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
if not self.is_checkpoint_fp8_serialized:
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineLinearMethod,
|
||||
)
|
||||
|
||||
online_method = Fp8PerTensorOnlineLinearMethod()
|
||||
online_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
|
||||
return online_method
|
||||
else:
|
||||
offline_method = Fp8LinearMethod(self)
|
||||
offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
|
||||
return offline_method
|
||||
elif isinstance(layer, RoutedExperts):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedFusedMoEMethod(layer.moe_config)
|
||||
if self.store_dtype == "mxfp4":
|
||||
from vllm.model_executor.layers.quantization.mxfp4 import (
|
||||
Mxfp4MoEMethod,
|
||||
)
|
||||
|
||||
return Mxfp4MoEMethod(layer.moe_config)
|
||||
if self.is_checkpoint_fp8_serialized:
|
||||
return Fp8MoEMethod(self, layer)
|
||||
else:
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import (
|
||||
Fp8PerTensorOnlineMoEMethod,
|
||||
)
|
||||
|
||||
return Fp8PerTensorOnlineMoEMethod(layer=layer)
|
||||
elif isinstance(layer, Attention):
|
||||
return Fp8KVCacheMethod(self)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_cache_scale_mapper() -> "WeightsMapper":
|
||||
"""Map compressed-tensors KV-cache scale names to vLLM names."""
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
orig_to_new_suffix = {
|
||||
".k_proj.output_scale": ".attn.k_scale",
|
||||
".v_proj.output_scale": ".attn.v_scale",
|
||||
".q_proj.output_scale": ".attn.q_scale",
|
||||
".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",
|
||||
}
|
||||
cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix)
|
||||
return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper()
|
||||
|
||||
|
||||
class CopyNumelCounter(TorchDispatchMode):
|
||||
"""
|
||||
Tracks total number of elements modified with `copy_`. Useful for keeping
|
||||
track of weight loading where underlying weights can be arbitrarily
|
||||
transformed (such as with `narrow`) before calling copy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.copied_numel = 0
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
out = func(*args, **kwargs)
|
||||
if func == torch.ops.aten.copy_.default:
|
||||
self.copied_numel += args[0].numel()
|
||||
return out
|
||||
|
||||
|
||||
def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None:
|
||||
"""Copies any attrs present in `old` but not in `new` to `new`"""
|
||||
new_attrs = set(dir(new))
|
||||
attrs_to_set = {}
|
||||
for attr in dir(old):
|
||||
if attr not in new_attrs:
|
||||
attrs_to_set[attr] = getattr(old, attr)
|
||||
set_weight_attrs(new, attrs_to_set)
|
||||
|
||||
|
||||
class Fp8LinearMethod(LinearMethodBase):
|
||||
"""Linear method for FP8.
|
||||
Supports loading FP8 checkpoints with static weight scale and
|
||||
dynamic/static activation scale.
|
||||
|
||||
Limitations:
|
||||
1. Only support float8_e4m3fn data type due to the limitation of
|
||||
torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
|
||||
|
||||
Args:
|
||||
quant_config: The quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config):
|
||||
self.quant_config = quant_config
|
||||
self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False)
|
||||
self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.input_dtype = get_current_vllm_config().model_config.dtype
|
||||
|
||||
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
|
||||
# kernel for fast weight-only FP8 quantization
|
||||
self.marlin_input_dtype = None
|
||||
self.use_marlin = False
|
||||
|
||||
if self.quant_config.use_deep_gemm is not None:
|
||||
self.use_deep_gemm = self.quant_config.use_deep_gemm
|
||||
else:
|
||||
self.use_deep_gemm = is_deep_gemm_supported()
|
||||
|
||||
self.weight_block_size = self.quant_config.weight_block_size
|
||||
self.block_quant = self.weight_block_size is not None
|
||||
self.act_q_static = self.quant_config.activation_scheme == "static"
|
||||
|
||||
if self.block_quant:
|
||||
assert not self.act_q_static
|
||||
assert self.weight_block_size is not None
|
||||
|
||||
self.activation_quant_key = create_fp8_quant_key(
|
||||
static=self.act_q_static,
|
||||
group_shape=GroupShape(1, self.weight_block_size[0]),
|
||||
)
|
||||
self.weight_quant_key = create_fp8_quant_key(
|
||||
static=True, group_shape=GroupShape(*self.weight_block_size)
|
||||
)
|
||||
else:
|
||||
self.weight_quant_key = kFp8StaticTensorSym
|
||||
# Use per-token quantization for better perf if dynamic and cutlass
|
||||
if self.act_q_static:
|
||||
self.activation_quant_key = kFp8StaticTensorSym
|
||||
elif cutlass_fp8_supported():
|
||||
self.activation_quant_key = kFp8DynamicTokenSym
|
||||
else:
|
||||
self.activation_quant_key = kFp8DynamicTensorSym
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
validate_fp8_block_shape(
|
||||
layer,
|
||||
input_size,
|
||||
output_size,
|
||||
input_size_per_partition,
|
||||
output_partition_sizes,
|
||||
self.weight_block_size,
|
||||
)
|
||||
|
||||
weight = create_fp8_weight_parameter(
|
||||
output_size_per_partition, input_size_per_partition, weight_loader
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
if not self.block_quant:
|
||||
scale = create_fp8_scale_parameter(
|
||||
PerTensorScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
None,
|
||||
weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", scale)
|
||||
else:
|
||||
assert not self.act_q_static
|
||||
assert self.weight_block_size is not None
|
||||
scale = create_fp8_scale_parameter(
|
||||
BlockQuantScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
self.weight_block_size,
|
||||
weight_loader,
|
||||
scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None),
|
||||
)
|
||||
# The weight_scale_inv name is intentional for deepseekv3
|
||||
layer.register_parameter("weight_scale_inv", scale)
|
||||
|
||||
# INPUT ACTIVATION SCALE
|
||||
if self.act_q_static:
|
||||
scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
|
||||
set_weight_attrs(scale, {"scale_type": "input_scale"})
|
||||
layer.register_parameter("input_scale", scale)
|
||||
|
||||
self.fp8_linear = init_fp8_linear_kernel(
|
||||
activation_quant_key=self.activation_quant_key,
|
||||
weight_quant_key=self.weight_quant_key,
|
||||
weight_shape=layer.weight.shape,
|
||||
input_dtype=self.input_dtype,
|
||||
out_dtype=self.out_dtype,
|
||||
module_name=self.__class__.__name__,
|
||||
)
|
||||
|
||||
self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if self.use_marlin:
|
||||
if not self.block_quant:
|
||||
# Canonicalize to (K, N) for the kernel.
|
||||
replace_parameter(layer, "weight", layer.weight.t())
|
||||
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
|
||||
# AttributeError if backend selection changes.
|
||||
if hasattr(self.fp8_linear, "marlin_input_dtype"):
|
||||
self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
return
|
||||
|
||||
input_scale = None
|
||||
# TODO(rob): refactor block quant into separate class.
|
||||
if self.block_quant:
|
||||
assert not self.act_q_static
|
||||
|
||||
# If checkpoint not serialized fp8, quantize the weights.
|
||||
else:
|
||||
# If checkpoint is fp8 per-tensor, handle that there are N scales for N
|
||||
# shards in a fused module
|
||||
weight = layer.weight
|
||||
weight_scale = layer.weight_scale
|
||||
|
||||
# If using w8a8, torch._scaled_mm needs per tensor, so
|
||||
# requantize the logical shards as a single weight.
|
||||
weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
|
||||
weight,
|
||||
weight_scale,
|
||||
layer.logical_widths,
|
||||
getattr(layer, "input_scale", None),
|
||||
)
|
||||
if self.act_q_static:
|
||||
assert input_scale is not None
|
||||
input_scale = input_scale.max()
|
||||
weight = weight.t()
|
||||
|
||||
# Update layer with new values.
|
||||
replace_parameter(layer, "weight", weight.data)
|
||||
replace_parameter(layer, "weight_scale", weight_scale.data)
|
||||
|
||||
if input_scale is not None:
|
||||
replace_parameter(layer, "input_scale", input_scale)
|
||||
else:
|
||||
layer.input_scale = None
|
||||
|
||||
self.fp8_linear.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
# if batch invariant mode is enabled, prefer direct FP8 path
|
||||
# we will use BF16 dequant when direct FP8 is not supported.
|
||||
if envs.VLLM_BATCH_INVARIANT:
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
return self.fp8_linear.apply_weights(
|
||||
layer,
|
||||
x,
|
||||
bias,
|
||||
)
|
||||
else:
|
||||
if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel):
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
# per-tensor/channel: dequant to BF16 and run GEMM
|
||||
weight_fp8 = layer.weight.to(torch.bfloat16)
|
||||
weight_scale = layer.weight_scale.to(torch.bfloat16)
|
||||
if weight_scale.numel() == 1:
|
||||
# Per-tensor: simple scalar multiplication
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
else:
|
||||
# Multiple scales (fused modules like QKV)
|
||||
# Try to infer correct broadcasting
|
||||
# weight is [K, N], scale could be [num_logical_weights]
|
||||
# Need to figure out how to broadcast - for now just try
|
||||
# direct multiplication
|
||||
if (
|
||||
weight_scale.dim() == 1
|
||||
and weight_scale.shape[0] == weight_fp8.shape[0]
|
||||
):
|
||||
# Per-row scaling
|
||||
weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
|
||||
else:
|
||||
# Fallback
|
||||
weight_bf16 = weight_fp8 * weight_scale
|
||||
return torch.nn.functional.linear(x, weight_bf16.t(), bias)
|
||||
|
||||
return self.fp8_linear.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
"""MoE method for FP8.
|
||||
Supports loading FP8 checkpoints with static weight scale and
|
||||
dynamic/static activation scale.
|
||||
|
||||
Also supports loading quantized FP16/BF16 model checkpoints with dynamic
|
||||
activation scaling. The weight scaling factor will be initialized after
|
||||
the model weights are loaded.
|
||||
|
||||
Args:
|
||||
quant_config: The quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
|
||||
super().__init__(layer.moe_config)
|
||||
self.quant_config = quant_config
|
||||
self.weight_block_size = self.quant_config.weight_block_size
|
||||
self.block_quant: bool = self.weight_block_size is not None
|
||||
self.weight_scale_name = (
|
||||
"weight_scale_inv" if self.block_quant else "weight_scale"
|
||||
)
|
||||
|
||||
# Set weight key and activation key for kernel compatibility
|
||||
if self.block_quant:
|
||||
weight_key = kFp8Static128BlockSym
|
||||
activation_key = kFp8Dynamic128Sym
|
||||
else:
|
||||
weight_key = kFp8StaticTensorSym
|
||||
activation_key = (
|
||||
kFp8StaticTensorSym
|
||||
if self.quant_config.activation_scheme == "static"
|
||||
else kFp8DynamicTensorSym
|
||||
)
|
||||
|
||||
# Select Fp8 MoE backend
|
||||
self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
|
||||
config=self.moe,
|
||||
weight_key=weight_key,
|
||||
activation_key=activation_key,
|
||||
allow_vllm_cutlass=False,
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
layer.num_experts = num_experts
|
||||
layer.orig_dtype = params_dtype
|
||||
layer.weight_block_size = None
|
||||
|
||||
assert self.quant_config.is_checkpoint_fp8_serialized
|
||||
params_dtype = torch.float8_e4m3fn
|
||||
|
||||
if self.block_quant:
|
||||
assert self.weight_block_size is not None
|
||||
layer.weight_block_size = self.weight_block_size
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
block_n, block_k = (
|
||||
self.weight_block_size[0],
|
||||
self.weight_block_size[1],
|
||||
)
|
||||
# NOTE: To ensure proper alignment of the block-wise quantization
|
||||
# scales, the output_size of the weights for both the gate and up
|
||||
# layers must be divisible by block_n.
|
||||
# Required by column parallel or enabling merged weights
|
||||
if intermediate_size_per_partition % block_n != 0:
|
||||
raise ValueError(
|
||||
f"The output_size of gate's and up's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_n = {block_n}."
|
||||
)
|
||||
if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
|
||||
# Required by row parallel
|
||||
raise ValueError(
|
||||
f"The input_size of down's weight = "
|
||||
f"{intermediate_size_per_partition} is not divisible by "
|
||||
f"weight quantization block_k = {block_k}."
|
||||
)
|
||||
|
||||
# WEIGHTS
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# BIASES (for models like GPT-OSS that have biased MoE)
|
||||
if self.moe.has_bias:
|
||||
w13_bias = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
dtype=layer.orig_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_bias", w13_bias)
|
||||
set_weight_attrs(w13_bias, extra_weight_attrs)
|
||||
w2_bias = torch.nn.Parameter(
|
||||
torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_bias", w2_bias)
|
||||
set_weight_attrs(w2_bias, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
if not self.block_quant:
|
||||
# For per-tensor quant, the scales are per expert and weight.
|
||||
w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32)
|
||||
w2_scale_data = torch.ones(num_experts, dtype=torch.float32)
|
||||
else:
|
||||
# For block quant, the scales are per block (typically 128x128).
|
||||
w13_scale_data = torch.ones(
|
||||
num_experts,
|
||||
2 * ((intermediate_size_per_partition + block_n - 1) // block_n),
|
||||
(hidden_size + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w2_scale_data = torch.ones(
|
||||
num_experts,
|
||||
(hidden_size + block_n - 1) // block_n,
|
||||
(intermediate_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False)
|
||||
w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False)
|
||||
# Note: name is weight_scale for tensor, weight_scale_inv for block.
|
||||
layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale)
|
||||
layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale)
|
||||
|
||||
# Add the quantization method used (per tensor/grouped/channel)
|
||||
# to ensure the weight scales are loaded in properly
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
|
||||
if self.block_quant
|
||||
else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
# INPUT_SCALES
|
||||
if self.quant_config.activation_scheme == "static":
|
||||
assert not self.block_quant
|
||||
w13_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w13_input_scale", w13_input_scale)
|
||||
set_weight_attrs(w13_input_scale, extra_weight_attrs)
|
||||
|
||||
w2_input_scale = torch.nn.Parameter(
|
||||
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.register_parameter("w2_input_scale", w2_input_scale)
|
||||
set_weight_attrs(w2_input_scale, extra_weight_attrs)
|
||||
|
||||
else:
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def _setup_kernel(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
w13_input_scale: torch.Tensor | None,
|
||||
w2_input_scale: torch.Tensor | None,
|
||||
) -> None:
|
||||
# Shuffle weights to runtime format.
|
||||
w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
|
||||
fp8_backend=self.fp8_backend,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
w13_input_scale=w13_input_scale,
|
||||
w2_input_scale=w2_input_scale,
|
||||
)
|
||||
|
||||
# Replace parameters with updated versions. Note that this helper
|
||||
# function ensures the replacement is compatible with RL weight reloads.
|
||||
replace_parameter(layer, "w13_weight", w13)
|
||||
replace_parameter(layer, "w2_weight", w2)
|
||||
replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
|
||||
replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
|
||||
|
||||
# AITER backend requires weights to be marked as shuffled.
|
||||
if self.fp8_backend == Fp8MoeBackend.AITER:
|
||||
layer.w13_weight.is_shuffled = True
|
||||
layer.w2_weight.is_shuffled = True
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.moe_quant_config is not None
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_fp8_moe_kernel(
|
||||
moe_quant_config=self.moe_quant_config,
|
||||
moe_config=self.moe,
|
||||
fp8_backend=self.fp8_backend,
|
||||
experts_cls=self.experts_cls,
|
||||
routing_tables=layer._expert_routing_tables(),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
||||
# Allow for accessing weights and scales in standard way.
|
||||
w13 = layer.w13_weight
|
||||
w2 = layer.w2_weight
|
||||
w13_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
w13_input_scale = layer.w13_input_scale
|
||||
w2_input_scale = layer.w2_input_scale
|
||||
|
||||
# MI300x and MI325x use FNUZ format for FP8. Convert if needed.
|
||||
if current_platform.is_fp8_fnuz():
|
||||
w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w13,
|
||||
w13_scale,
|
||||
w13_input_scale,
|
||||
)
|
||||
w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
w2,
|
||||
w2_scale,
|
||||
w2_input_scale,
|
||||
)
|
||||
|
||||
# Per tensor kernels require single activation scale. Use the max.
|
||||
if self.quant_config.activation_scheme == "static":
|
||||
assert not self.block_quant
|
||||
assert w13_input_scale is not None and w2_input_scale is not None
|
||||
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
|
||||
w13_input_scale, w2_input_scale
|
||||
)
|
||||
replace_parameter(layer, "w13_input_scale", w13_input_scale)
|
||||
replace_parameter(layer, "w2_input_scale", w2_input_scale)
|
||||
|
||||
# Per tensor kernels require single weight scale for w13 per expert, but
|
||||
# on disk there is a scale for w1 and w3. Use the max to requantize.
|
||||
if not self.block_quant:
|
||||
shard_size = layer.intermediate_size_per_partition
|
||||
w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
|
||||
w13, w13_scale, shard_size, layer.local_num_experts
|
||||
)
|
||||
|
||||
# Shuffle weights to runtime format and setup kernel.
|
||||
self._setup_kernel(
|
||||
layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
|
||||
)
|
||||
|
||||
def maybe_make_prepare_finalize(
|
||||
self,
|
||||
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||||
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__} uses the new modular kernel initialization "
|
||||
"logic. This function should not be called."
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
|
||||
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
|
||||
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
|
||||
a1_scale = layer.w13_input_scale
|
||||
a2_scale = layer.w2_input_scale
|
||||
|
||||
quant_config = make_fp8_moe_quant_config(
|
||||
fp8_backend=self.fp8_backend,
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
block_shape=self.weight_block_size,
|
||||
swiglu_limit=getattr(layer, "swiglu_limit", None),
|
||||
gemm1_alpha=getattr(layer, "swiglu_alpha", None),
|
||||
gemm1_beta=getattr(layer, "swiglu_beta", None),
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
# Inject biases into the quant config if the model has them
|
||||
# (e.g. GPT-OSS biased MoE)
|
||||
if quant_config is not None and self.moe.has_bias:
|
||||
w13_bias = getattr(layer, "w13_bias", None)
|
||||
w2_bias = getattr(layer, "w2_bias", None)
|
||||
if w13_bias is not None:
|
||||
quant_config._w1.bias = w13_bias
|
||||
if w2_bias is not None:
|
||||
quant_config._w2.bias = w2_bias
|
||||
|
||||
return quant_config
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
return True
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
assert self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply_monolithic(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
router_logits,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
num_expert_group=layer.num_expert_group,
|
||||
topk_group=layer.topk_group,
|
||||
e_score_correction_bias=layer.e_score_correction_bias,
|
||||
routed_scaling_factor=layer.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
assert not self.is_monolithic
|
||||
assert self.moe_kernel is not None
|
||||
return self.moe_kernel.apply(
|
||||
x,
|
||||
layer.w13_weight,
|
||||
layer.w2_weight,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation=layer.activation,
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
apply_router_weight_on_input=layer.apply_router_weight_on_input,
|
||||
shared_experts=shared_experts,
|
||||
shared_experts_input=shared_experts_input,
|
||||
)
|
||||
|
||||
|
||||
class Fp8KVCacheMethod(BaseKVCacheMethod):
|
||||
"""
|
||||
Supports loading kv-cache scaling factors from FP8 checkpoints.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: Fp8Config):
|
||||
super().__init__(quant_config)
|
||||
824
upstream/vllm/models/deepseek_v4/attention.py
Normal file
824
upstream/vllm/models/deepseek_v4/attention.py
Normal file
@@ -0,0 +1,824 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
DeepseekV4 MLA Attention Layer
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, 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 (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
|
||||
from vllm.models.deepseek_v4.common.ops import (
|
||||
fused_indexer_q_rope_quant,
|
||||
fused_q_kv_rmsnorm,
|
||||
)
|
||||
from vllm.utils.deep_gemm import fp8_einsum
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
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.models.utils import extract_layer_index
|
||||
from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
|
||||
from vllm.models.deepseek_v4.compressor import DeepseekCompressor
|
||||
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.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,
|
||||
get_kv_quant_mode,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _resolve_dsv4_kv_cache_dtype(
|
||||
use_fp8_ds_mla_layout: bool,
|
||||
kv_cache_dtype: str,
|
||||
cache_config: CacheConfig | None,
|
||||
) -> tuple[str, torch.dtype]:
|
||||
"""Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
|
||||
|
||||
Both layouts are paged; they differ in the per-token block format. The
|
||||
``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
|
||||
canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
|
||||
page-size specs pick the 576B per-token slot). Plain-row backends store each
|
||||
token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
|
||||
"""
|
||||
if use_fp8_ds_mla_layout:
|
||||
# fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
|
||||
assert kv_cache_dtype.startswith("fp8"), (
|
||||
f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, "
|
||||
f"got {kv_cache_dtype}"
|
||||
)
|
||||
if kv_cache_dtype != "fp8_ds_mla":
|
||||
if 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.")
|
||||
return kv_cache_dtype, torch.uint8
|
||||
|
||||
# Plain bf16 / per-tensor fp8 KV row (FlashInfer).
|
||||
if kv_cache_dtype.startswith("fp8"):
|
||||
return kv_cache_dtype, torch.float8_e4m3fn
|
||||
# auto / bfloat16 -> plain bf16 KV row.
|
||||
return kv_cache_dtype, torch.bfloat16
|
||||
|
||||
|
||||
class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
|
||||
"""DeepseekV4 MLA attention layer.
|
||||
|
||||
The platform-specific sparse-MLA forward (``forward_mqa`` /
|
||||
``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
|
||||
subclass — ``DeepseekV4FlashMLAAttention`` /
|
||||
``DeepseekV4FlashInferSM120Attention`` /
|
||||
``DeepseekV4FlashInferMLAAttention`` (CUDA) or
|
||||
``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
|
||||
deepseek_v4 model module. The base is never instantiated directly.
|
||||
"""
|
||||
|
||||
# Provided by the platform subclass.
|
||||
backend_cls: ClassVar[type[AttentionBackend]]
|
||||
# KV-cache per-token block format (both layouts are paged). True (default)
|
||||
# = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain
|
||||
# bf16 / per-tensor fp8 KV row. Backends can override the instance hook when
|
||||
# a single attention class dispatches across arch-specific layouts.
|
||||
use_fp8_ds_mla_layout: ClassVar[bool] = True
|
||||
# Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
|
||||
# workspace allocated in _forward_prefill and is also read by the dummy-run
|
||||
# path to pre-reserve that workspace.
|
||||
PREFILL_CHUNK_SIZE: ClassVar[int] = 4
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_padded_num_q_heads(cls, num_heads: int) -> int:
|
||||
"""Q head count the q/output buffers are allocated at.
|
||||
|
||||
The layer 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
|
||||
|
||||
@abstractmethod
|
||||
def forward_mqa(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
) -> None:
|
||||
"""Platform-specific sparse MLA forward; writes attention into ``output``."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _uses_fp8_ds_mla_layout(self) -> bool:
|
||||
"""Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
|
||||
return self.use_fp8_ds_mla_layout
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str,
|
||||
topk_indices_buffer: torch.Tensor | None = None,
|
||||
aux_stream_list: list[torch.cuda.Stream] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
cache_config = vllm_config.cache_config
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
layer_id = extract_layer_index(prefix)
|
||||
|
||||
self.prefix = prefix # Alias for compatibility with compressor
|
||||
self.hidden_size = config.hidden_size
|
||||
self.n_heads = config.num_attention_heads
|
||||
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.scale = self.head_dim**-0.5
|
||||
|
||||
# Padded Q head count is dictated by the platform subclass.
|
||||
self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads)
|
||||
# Sink padded to the same head count, initialized to -inf (no sink
|
||||
# effect). Weight loading fills the first n_local_heads slots.
|
||||
self.attn_sink = nn.Parameter(
|
||||
torch.full((self.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)
|
||||
# wo_a uses the standard quant_config for all quantization schemes
|
||||
# (FP8 blockwise, FP8 channelwise, INT8 channelwise). The _o_proj
|
||||
# method in the platform subclass dispatches at runtime based on
|
||||
# weight dtype and scale attribute:
|
||||
# - int8 + weight_scale → INT8 channelwise einsum
|
||||
# - fp8 + weight_scale → FP8 channelwise einsum
|
||||
# - fp8 + weight_scale_inv → FP8 blockwise einsum
|
||||
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",
|
||||
)
|
||||
|
||||
# Initialize rotary embedding before the indexer/compressor consume it.
|
||||
self.rotary_emb = build_deepseek_v4_rope(
|
||||
config,
|
||||
head_dim=self.head_dim,
|
||||
rope_head_dim=self.rope_head_dim,
|
||||
max_position_embeddings=config.max_position_embeddings,
|
||||
compress_ratio=self.compress_ratio,
|
||||
)
|
||||
self.indexer_rotary_emb = self.rotary_emb
|
||||
self.topk_indices_buffer = topk_indices_buffer
|
||||
|
||||
self.indexer = None
|
||||
if self.compress_ratio == 4:
|
||||
# Only C4A uses sparse attention and hence has indexer.
|
||||
# aux_stream_list[2] is free here (outer GEMMs joined) for the inner
|
||||
# overlap of wq_b+fused_indexer_q_rope_quant vs compressor. None on
|
||||
# ROCm, where aux_stream_list is None.
|
||||
indexer_aux_stream = (
|
||||
aux_stream_list[2] if aux_stream_list is not None else None
|
||||
)
|
||||
self.indexer = DeepseekV4Indexer(
|
||||
vllm_config,
|
||||
config=config,
|
||||
hidden_size=self.hidden_size,
|
||||
q_lora_rank=self.q_lora_rank,
|
||||
quant_config=quant_config,
|
||||
cache_config=cache_config,
|
||||
topk_indices_buffer=topk_indices_buffer,
|
||||
compress_ratio=self.compress_ratio,
|
||||
prefix=f"{prefix}.indexer",
|
||||
aux_stream=indexer_aux_stream,
|
||||
)
|
||||
|
||||
# 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"
|
||||
# ---- Attention / KV-cache setup ----
|
||||
self.max_num_batched_tokens = (
|
||||
vllm_config.scheduler_config.max_num_batched_tokens
|
||||
)
|
||||
self.max_model_len = vllm_config.model_config.max_model_len
|
||||
|
||||
# Resolve the kv-cache dtype from this backend's block format. The same
|
||||
# resolution drives the SWA cache tensor dtype below.
|
||||
self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
|
||||
self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
|
||||
)
|
||||
|
||||
self.swa_cache_layer = DeepseekV4SWACache(
|
||||
head_dim=self.head_dim,
|
||||
window_size=self.window_size,
|
||||
dtype=self.kv_cache_torch_dtype,
|
||||
prefix=f"{prefix}.swa_cache",
|
||||
cache_config=cache_config,
|
||||
)
|
||||
|
||||
# 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([])
|
||||
|
||||
# Create the compressor for layers with compress_ratio > 1; after the
|
||||
# attention setup above so its KV-cache prefix (self.prefix) is set.
|
||||
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.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,
|
||||
)
|
||||
|
||||
# Metadata-independent input GEMMs + RMSNorm stay in the captured
|
||||
# graph; the metadata-dependent rest (q up-proj + kv-insert, indexer,
|
||||
# compressor, MLA attention) runs in the eager break.
|
||||
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,
|
||||
)
|
||||
|
||||
# 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,
|
||||
qr,
|
||||
kv,
|
||||
kv_score,
|
||||
indexer_kv_score,
|
||||
indexer_weights,
|
||||
positions,
|
||||
o_padded,
|
||||
)
|
||||
o = o_padded[:, : self.n_local_heads, :]
|
||||
|
||||
# Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
|
||||
return self._o_proj(o, positions)
|
||||
|
||||
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,
|
||||
qr: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
kv_score: torch.Tensor,
|
||||
indexer_kv_score: torch.Tensor,
|
||||
indexer_weights: 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
|
||||
|
||||
# wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
|
||||
# on the default stream so q stays on its consumer stream (forward_mqa
|
||||
# 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.forward_mqa(q, kv, positions, 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
|
||||
# The fused insert ops require int64 position_ids; the runner's positions
|
||||
# buffer is already int64, so no cast is needed.
|
||||
assert positions.dtype == torch.int64
|
||||
cos_sin_cache = self.rotary_emb.cos_sin_cache
|
||||
cache_dtype = swa_kv_cache.dtype
|
||||
|
||||
# kv is unchanged; attention reads kv solely via swa_kv_cache.
|
||||
if cache_dtype == torch.uint8:
|
||||
# fp8_ds_mla UE8M0 paged path. Horizontally fused:
|
||||
# Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
|
||||
# the padding head slots; the kernel allocates and returns
|
||||
# the padded q tensor.
|
||||
# KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
|
||||
swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
|
||||
return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
q,
|
||||
kv,
|
||||
swa_kv_cache_2d,
|
||||
swa_metadata.slot_mapping,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
self.padded_heads,
|
||||
self.eps,
|
||||
swa_metadata.block_size,
|
||||
)
|
||||
|
||||
# Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
|
||||
# row in its element dtype (no Q padding). bf16 rewrites q in place;
|
||||
# per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
|
||||
# KV row.
|
||||
block_size = swa_metadata.block_size
|
||||
swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
|
||||
if cache_dtype == torch.bfloat16:
|
||||
torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
|
||||
q,
|
||||
kv,
|
||||
swa_kv_cache_3d,
|
||||
swa_metadata.slot_mapping,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
self.eps,
|
||||
block_size,
|
||||
)
|
||||
return q
|
||||
|
||||
# per-tensor fp8 (torch.float8_e4m3fn)
|
||||
q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn)
|
||||
torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
q,
|
||||
kv,
|
||||
q_fp8,
|
||||
swa_kv_cache_3d,
|
||||
swa_metadata.slot_mapping,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
self._flashinfer_fp8_kv_scale,
|
||||
self._flashinfer_fp8_q_scale_inv,
|
||||
self.eps,
|
||||
block_size,
|
||||
)
|
||||
return q_fp8
|
||||
|
||||
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
|
||||
# fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
|
||||
# alignment; plain bf16 / per-tensor fp8 rows use natural element-size
|
||||
# pages.
|
||||
uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
|
||||
return MLAAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=self.head_dim,
|
||||
dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
|
||||
compress_ratio=self.compress_ratio,
|
||||
cache_dtype_str=self.kv_cache_dtype,
|
||||
alignment=576 if uses_fp8_ds_mla_layout else 512,
|
||||
model_version="deepseek_v4",
|
||||
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
dtype: torch.dtype,
|
||||
prefix: str,
|
||||
cache_config: CacheConfig,
|
||||
compress_ratio: int = 1,
|
||||
n_head: int | None = None,
|
||||
q_head_dim: int | None = None,
|
||||
):
|
||||
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
|
||||
self.n_head = n_head
|
||||
self.q_head_dim = q_head_dim
|
||||
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.
|
||||
uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
|
||||
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,
|
||||
# 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
|
||||
alignment=576 if uses_fp8_ds_mla_layout else 512,
|
||||
indexer_n_head=self.n_head,
|
||||
indexer_q_head_dim=self.q_head_dim,
|
||||
)
|
||||
|
||||
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,
|
||||
n_head=self.n_head,
|
||||
q_head_dim=self.head_dim,
|
||||
)
|
||||
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)
|
||||
393
upstream/vllm/models/deepseek_v4/ppu/flashmla.py
Normal file
393
upstream/vllm/models/deepseek_v4/ppu/flashmla.py
Normal file
@@ -0,0 +1,393 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.models.deepseek_v4.attention import DeepseekV4Attention
|
||||
from vllm.platforms import current_platform
|
||||
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.models.deepseek_v4.ppu.ops.o_proj import (
|
||||
compute_fp8_einsum_recipe,
|
||||
deep_gemm_fp8_o_proj,
|
||||
deep_gemm_fp8_channel_o_proj,
|
||||
deep_gemm_int8_o_proj,
|
||||
)
|
||||
from vllm.models.deepseek_v4.sparse_mla import (
|
||||
DeepseekV4FlashMLABackend,
|
||||
DeepseekV4FlashMLAMetadata,
|
||||
)
|
||||
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.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
|
||||
|
||||
|
||||
class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
|
||||
"""FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA)."""
|
||||
|
||||
backend_cls = DeepseekV4FlashMLABackend
|
||||
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
|
||||
|
||||
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
||||
if self.wo_a.weight.dtype == torch.int8:
|
||||
return deep_gemm_int8_o_proj(
|
||||
o,
|
||||
positions,
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
self.wo_a,
|
||||
self.wo_b,
|
||||
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,
|
||||
o_lora_rank=self.o_lora_rank,
|
||||
einsum_recipe=self._einsum_recipe,
|
||||
)
|
||||
elif self.wo_a.weight.dtype == torch.float8_e4m3fn and hasattr(
|
||||
self.wo_a, "weight_scale"
|
||||
):
|
||||
# Channelwise FP8 (PPU): wo_a kept as FP8 + per-output-channel
|
||||
# weight_scale (no pre-dequant).
|
||||
return deep_gemm_fp8_channel_o_proj(
|
||||
o,
|
||||
positions,
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
self.wo_a,
|
||||
self.wo_b,
|
||||
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,
|
||||
o_lora_rank=self.o_lora_rank,
|
||||
einsum_recipe=self._einsum_recipe,
|
||||
)
|
||||
return deep_gemm_fp8_o_proj(
|
||||
o,
|
||||
positions,
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
self.wo_a,
|
||||
self.wo_b,
|
||||
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,
|
||||
o_lora_rank=self.o_lora_rank,
|
||||
einsum_recipe=self._einsum_recipe,
|
||||
tma_aligned_scales=self._tma_aligned_scales,
|
||||
)
|
||||
|
||||
@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})."
|
||||
)
|
||||
if current_platform.is_ppu():
|
||||
# PPU FlashMLA supports arbitrary head counts; no padding needed.
|
||||
return num_heads
|
||||
return 64 if num_heads <= 64 else 128
|
||||
|
||||
def forward_mqa(
|
||||
self,
|
||||
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 = self.compress_ratio <= 1
|
||||
N = (
|
||||
0
|
||||
if swa_only
|
||||
else (self.max_model_len + self.compress_ratio - 1)
|
||||
// self.compress_ratio
|
||||
)
|
||||
M = N + self.window_size + self.max_num_batched_tokens
|
||||
current_workspace_manager().get_simultaneous(
|
||||
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
|
||||
)
|
||||
output.zero_()
|
||||
return
|
||||
|
||||
assert isinstance(attn_metadata, dict)
|
||||
flashmla_metadata = cast(
|
||||
DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
|
||||
)
|
||||
swa_metadata = cast(
|
||||
"DeepseekSparseSWAMetadata | None",
|
||||
attn_metadata.get(self.swa_cache_layer.prefix),
|
||||
)
|
||||
assert swa_metadata is not None
|
||||
|
||||
swa_only = self.compress_ratio <= 1
|
||||
# SWA-only layers (compress_ratio <= 1) don't have their own KV cache
|
||||
# allocation, so self.kv_cache may be empty after profiling cleanup.
|
||||
self_kv_cache = self.kv_cache if not swa_only else None
|
||||
swa_kv_cache = self.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:
|
||||
self._forward_prefill(
|
||||
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:
|
||||
self._forward_decode(
|
||||
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],
|
||||
)
|
||||
|
||||
def _forward_decode(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
|
||||
swa_metadata: "DeepseekSparseSWAMetadata",
|
||||
attn_metadata: DeepseekV4FlashMLAMetadata | 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 // self.compress_ratio
|
||||
is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
|
||||
if self.compress_ratio == 4:
|
||||
# C4A: local indices differ per layer (filled by Indexer).
|
||||
assert self.topk_indices_buffer is not None
|
||||
global_indices, topk_lens = compute_global_topk_indices_and_lens(
|
||||
self.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 self.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 = self.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 self.compress_ratio <= 1:
|
||||
tile_metadata = swa_metadata.tile_sched_swaonly
|
||||
elif self.compress_ratio == 4:
|
||||
tile_metadata = swa_metadata.tile_sched_c4a
|
||||
elif self.compress_ratio == 128:
|
||||
tile_metadata = swa_metadata.tile_sched_c128a
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported compress_ratio={self.compress_ratio}; "
|
||||
"expected 1, 4, or 128."
|
||||
)
|
||||
assert tile_metadata is not None, (
|
||||
"swa_metadata missing tile_sched entry for "
|
||||
f"compress_ratio={self.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=self.scale,
|
||||
attn_sink=self.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),
|
||||
)
|
||||
|
||||
def _forward_prefill(
|
||||
self,
|
||||
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: DeepseekV4FlashMLAMetadata | 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 self.compress_ratio == 4:
|
||||
assert self.topk_indices_buffer is not None
|
||||
topk_indices = self.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 = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
|
||||
else:
|
||||
# NOTE(woosuk): topk_indices will not be used for SWA-only layers.
|
||||
assert self.topk_indices_buffer is not None
|
||||
topk_indices = self.topk_indices_buffer[num_decode_tokens:]
|
||||
top_k = 0
|
||||
N = 0
|
||||
|
||||
M = N + self.window_size + self.max_num_batched_tokens
|
||||
chunk_size_const = self.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] // self.compress_ratio,
|
||||
gather_lens=None,
|
||||
block_table=block_table[chunk_start:chunk_end],
|
||||
block_size=attn_metadata.block_size // self.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],
|
||||
self.window_size,
|
||||
self.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=self.scale,
|
||||
attn_sink=self.attn_sink,
|
||||
topk_length=combined_lens,
|
||||
out=output[query_start:query_end],
|
||||
)
|
||||
244
upstream/vllm/models/deepseek_v4/ppu/ops/o_proj.py
Normal file
244
upstream/vllm/models/deepseek_v4/ppu/ops/o_proj.py
Normal file
@@ -0,0 +1,244 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.models.deepseek_v4.common.ops import (
|
||||
fused_inv_rope_fp8_quant,
|
||||
fused_inv_rope_fp8_quant_channelwise,
|
||||
fused_inv_rope_int8_quant_channelwise,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.ppu_deep_gemm import fp8_einsum
|
||||
from vllm.utils.ppu_deep_gemm import int8_einsum
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
|
||||
def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]:
|
||||
"""fp8_einsum recipe + scale layout for the current 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.
|
||||
|
||||
Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``.
|
||||
"""
|
||||
cap = current_platform.get_device_capability()
|
||||
assert cap is not None, "DeepseekV4 attention requires a CUDA device"
|
||||
einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
|
||||
tma_aligned_scales = cap.major >= 10
|
||||
return einsum_recipe, tma_aligned_scales
|
||||
|
||||
|
||||
def deep_gemm_fp8_o_proj(
|
||||
o: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
wo_a: nn.Module,
|
||||
wo_b: nn.Module,
|
||||
*,
|
||||
n_groups: int,
|
||||
heads_per_group: int,
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
o_lora_rank: int,
|
||||
einsum_recipe: tuple[int, int, int],
|
||||
tma_aligned_scales: bool,
|
||||
) -> torch.Tensor:
|
||||
"""O projection: inverse RoPE + FP8 quant + einsum + wo_b.
|
||||
|
||||
Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` /
|
||||
``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``.
|
||||
"""
|
||||
o_fp8, o_scale = fused_inv_rope_fp8_quant(
|
||||
o,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
n_groups=n_groups,
|
||||
heads_per_group=heads_per_group,
|
||||
nope_dim=nope_dim,
|
||||
rope_dim=rope_dim,
|
||||
tma_aligned_scales=tma_aligned_scales,
|
||||
)
|
||||
z = torch.empty(
|
||||
(o.shape[0], n_groups, o_lora_rank),
|
||||
device=o.device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
fp8_einsum(
|
||||
"bhr,hdr->bhd",
|
||||
(o_fp8, o_scale),
|
||||
(wo_a.weight, wo_a.weight_scale_inv),
|
||||
z,
|
||||
recipe=einsum_recipe,
|
||||
)
|
||||
return wo_b(z.flatten(1))
|
||||
|
||||
|
||||
def deep_gemm_int8_o_proj(
|
||||
o: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
wo_a: nn.Module,
|
||||
wo_b: nn.Module,
|
||||
*,
|
||||
n_groups: int,
|
||||
heads_per_group: int,
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
o_lora_rank: int,
|
||||
einsum_recipe: tuple[int, int, int],
|
||||
):
|
||||
# Channelwise INT8 GEMM via PPU DeepGEMM int8_einsum:
|
||||
# - activation: per-(token, group) symmetric INT8 quant.
|
||||
# o_int8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1).
|
||||
# - weight: per-output-channel INT8. compressed-tensors
|
||||
# W8A8 channel scheme stores wo_a as 2D [G*d, R] / scale as
|
||||
# [G*d, 1]; reshape to 3D so int8_bmm sees (h=G, d=lora_rank,
|
||||
# r=R) for `hdr` and a (h, d, 1) scale broadcasting along r.
|
||||
o_int8, o_scale = fused_inv_rope_int8_quant_channelwise(
|
||||
o,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
n_groups,
|
||||
heads_per_group,
|
||||
nope_dim,
|
||||
rope_dim,
|
||||
)
|
||||
wo_a_view = wo_a.weight.reshape(
|
||||
n_groups, o_lora_rank, -1
|
||||
)
|
||||
wo_a_scale_view = wo_a.weight_scale.reshape(
|
||||
n_groups, o_lora_rank, 1
|
||||
)
|
||||
z = torch.empty(
|
||||
(o.shape[0], n_groups, o_lora_rank),
|
||||
device=o.device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
torch.ops.vllm.deepseek_v4_int8_einsum(
|
||||
o_int8,
|
||||
o_scale,
|
||||
wo_a_view,
|
||||
wo_a_scale_view,
|
||||
z,
|
||||
"bhr,hdr->bhd",
|
||||
[1, 1, o_int8.shape[-1]],
|
||||
)
|
||||
return wo_b(z.flatten(1))
|
||||
|
||||
|
||||
def deep_gemm_fp8_channel_o_proj(
|
||||
o: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
wo_a: nn.Module,
|
||||
wo_b: nn.Module,
|
||||
*,
|
||||
n_groups: int,
|
||||
heads_per_group: int,
|
||||
nope_dim: int,
|
||||
rope_dim: int,
|
||||
o_lora_rank: int,
|
||||
einsum_recipe: tuple[int, int, int],
|
||||
):
|
||||
# Channelwise FP8 GEMM via PPU DeepGEMM fp8_einsum:
|
||||
# - activation: per-(token, group) symmetric FP8 quant.
|
||||
# o_fp8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1).
|
||||
# - weight: per-output-channel FP8. Channelwise scheme stores
|
||||
# wo_a as 2D [G*d, R] / scale as [G*d, 1]; reshape to 3D so
|
||||
# fp8_bmm sees (h=G, d=lora_rank, r=R) for `hdr` and a
|
||||
# (h, d, 1) scale broadcasting along r.
|
||||
o_fp8, o_scale = fused_inv_rope_fp8_quant_channelwise(
|
||||
o,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
n_groups,
|
||||
heads_per_group,
|
||||
nope_dim,
|
||||
rope_dim,
|
||||
)
|
||||
wo_a_view = wo_a.weight.reshape(
|
||||
n_groups, o_lora_rank, -1
|
||||
)
|
||||
wo_a_scale_view = wo_a.weight_scale.reshape(
|
||||
n_groups, o_lora_rank, 1
|
||||
)
|
||||
z = torch.empty(
|
||||
(o.shape[0], n_groups, o_lora_rank),
|
||||
device=o.device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
torch.ops.vllm.deepseek_v4_fp8_channel_einsum(
|
||||
o_fp8,
|
||||
o_scale,
|
||||
wo_a_view,
|
||||
wo_a_scale_view,
|
||||
z,
|
||||
"bhr,hdr->bhd",
|
||||
[1, 1, o_fp8.shape[-1]],
|
||||
)
|
||||
return wo_b(z.flatten(1))
|
||||
|
||||
|
||||
def deepseek_v4_int8_einsum(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
b_scale: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
equation: str,
|
||||
recipe: list[int],
|
||||
) -> None:
|
||||
int8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe))
|
||||
|
||||
|
||||
def deepseek_v4_int8_einsum_fake(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
b_scale: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
equation: str,
|
||||
recipe: list[int],
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="deepseek_v4_int8_einsum",
|
||||
op_func=deepseek_v4_int8_einsum,
|
||||
mutates_args=["out"],
|
||||
fake_impl=deepseek_v4_int8_einsum_fake,
|
||||
)
|
||||
|
||||
|
||||
def deepseek_v4_fp8_channel_einsum(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
b_scale: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
equation: str,
|
||||
recipe: list[int],
|
||||
) -> None:
|
||||
fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe))
|
||||
|
||||
|
||||
def deepseek_v4_fp8_channel_einsum_fake(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
b_scale: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
equation: str,
|
||||
recipe: list[int],
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="deepseek_v4_fp8_channel_einsum",
|
||||
op_func=deepseek_v4_fp8_channel_einsum,
|
||||
mutates_args=["out"],
|
||||
fake_impl=deepseek_v4_fp8_channel_einsum_fake,
|
||||
)
|
||||
Reference in New Issue
Block a user