[Model] Support DeepSeek-V4
This commit is contained in:
chenxb002
934
vllm_mlu/v1/attention/backends/mla/flashmla.py
Normal file
934
vllm_mlu/v1/attention/backends/mla/flashmla.py
Normal file
@@ -0,0 +1,934 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-MLU project
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.attention.backends.abstract import (AttentionType,
|
||||
is_quantized_kv_cache)
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.utils.math_utils import cdiv, round_down
|
||||
from vllm.attention.backends.utils import MLADims
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.v1.attention.backends.mla.common import (
|
||||
MLACommonBackend, MLACommonPrefillMetadata,
|
||||
MLACommonDecodeMetadata, MLACommonMetadata,
|
||||
MLACommonMetadataBuilder, M, QueryLenSupport,
|
||||
use_cudnn_prefill, use_flashinfer_prefill,
|
||||
use_trtllm_ragged_deepseek_prefill,
|
||||
FlashInferPrefillMetadata,
|
||||
CudnnPrefillMetadata,
|
||||
MLACommonImpl,
|
||||
CUDNN_WORKSPACE_SIZE
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
AttentionCGSupport, split_decodes_and_prefills,
|
||||
infer_global_hyperparameters, get_per_layer_parameters,
|
||||
)
|
||||
from vllm.attention.backends.abstract import (
|
||||
AttentionBackend,
|
||||
AttentionLayer,
|
||||
MLAAttentionImpl,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.worker.gpu_input_batch import InputBatch
|
||||
|
||||
import vllm_mlu._mlu_utils as mlu_envs
|
||||
from vllm_mlu import _mlu_ops as mlu_ops
|
||||
from vllm_mlu.v1.attention.backends.flash_attn import MLUFlashAttentionImpl
|
||||
from vllm_mlu.v1.attention.backends.utils import (
|
||||
MLUCommonAttentionMetadata, get_common_metadata,
|
||||
MLUInferMode)
|
||||
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
|
||||
from vllm.platforms import current_platform
|
||||
from vllm import envs
|
||||
try:
|
||||
from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
|
||||
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache # noqa: F401
|
||||
|
||||
flashinfer_available = True
|
||||
except ImportError:
|
||||
BatchPrefillWithRaggedKVCacheWrapper = object
|
||||
|
||||
flashinfer_available = False
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
|
||||
from vllm_mlu.mlu_hijack_utils import MluHijackObject
|
||||
|
||||
|
||||
class MLACommonBackend_MluHijack(MLACommonBackend):
|
||||
|
||||
@classmethod
|
||||
def get_supported_head_sizes(cls) -> list[int]:
|
||||
return [576, 512]
|
||||
|
||||
|
||||
def get_mla_dims(model_config: ModelConfig) -> MLADims:
|
||||
hf_text_config = model_config.hf_text_config
|
||||
|
||||
if model_config.hf_text_config.model_type == "deepseek_v4":
|
||||
return MLADims(
|
||||
q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
|
||||
kv_lora_rank=hf_text_config.head_dim,
|
||||
qk_nope_head_dim=hf_text_config.head_dim - hf_text_config.rope_head_dim,
|
||||
qk_rope_head_dim=hf_text_config.rope_head_dim,
|
||||
v_head_dim=hf_text_config.head_dim,
|
||||
)
|
||||
|
||||
return MLADims(
|
||||
q_lora_rank=getattr(hf_text_config, "q_lora_rank", None),
|
||||
kv_lora_rank=hf_text_config.kv_lora_rank,
|
||||
qk_nope_head_dim=hf_text_config.qk_nope_head_dim,
|
||||
qk_rope_head_dim=hf_text_config.qk_rope_head_dim,
|
||||
v_head_dim=hf_text_config.v_head_dim,
|
||||
)
|
||||
|
||||
|
||||
class MLACommonMetadataBuilder_MluHijack(MLACommonMetadataBuilder):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kv_cache_spec: AttentionSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config: VllmConfig,
|
||||
device: torch.device,
|
||||
metadata_cls: type[M] | None = None,
|
||||
supports_dcp_with_varlen: bool = False,
|
||||
):
|
||||
self.metadata_cls = (
|
||||
metadata_cls if metadata_cls is not None else MLACommonMetadata
|
||||
)
|
||||
self.kv_cache_spec = kv_cache_spec
|
||||
scheduler_config = vllm_config.scheduler_config
|
||||
self.model_config = vllm_config.model_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
self.compilation_config = vllm_config.compilation_config
|
||||
self.vllm_config = vllm_config
|
||||
self.device = device
|
||||
|
||||
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
|
||||
self.mla_dims = get_mla_dims(self.model_config)
|
||||
self.aot_schedule = current_platform.is_cuda()
|
||||
try:
|
||||
self.dcp_world_size = get_dcp_group().world_size
|
||||
self.dcp_rank = get_dcp_group().rank_in_group
|
||||
except AssertionError:
|
||||
# DCP might not be initialized in testing
|
||||
self.dcp_world_size = 1
|
||||
self.dcp_rank = 0
|
||||
self.dcp_local_block_size = parallel_config.dcp_kv_cache_interleave_size
|
||||
self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
|
||||
|
||||
# Don't try to access the runner on AMD
|
||||
if self.aot_schedule:
|
||||
self.page_size = self.kv_cache_spec.block_size
|
||||
|
||||
self.chunked_prefill_workspace_size = (
|
||||
self.determine_chunked_prefill_workspace_size(vllm_config)
|
||||
)
|
||||
|
||||
if self.dcp_world_size > 1:
|
||||
# Note(hc): The local kvcache is incomplete when DCP is triggered,
|
||||
# an additional kvcache allgather across the DCP group is therefore
|
||||
# required, so the workspace has to be enlarged by 1/DCP relative
|
||||
# to the original TP allocation.
|
||||
assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
|
||||
self.chunked_prefill_workspace = torch.empty(
|
||||
(
|
||||
self.chunked_prefill_workspace_size
|
||||
+ self.chunked_prefill_workspace_size // self.dcp_world_size,
|
||||
self.model_config.get_head_size(),
|
||||
),
|
||||
dtype=self.model_config.dtype,
|
||||
device=device,
|
||||
)
|
||||
else:
|
||||
self.chunked_prefill_workspace = torch.empty(
|
||||
(
|
||||
self.chunked_prefill_workspace_size,
|
||||
self.model_config.get_head_size(),
|
||||
),
|
||||
dtype=self.model_config.dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self._use_cudnn_prefill = use_cudnn_prefill()
|
||||
self._use_fi_prefill = use_flashinfer_prefill()
|
||||
self._use_trtllm_ragged_prefill = use_trtllm_ragged_deepseek_prefill()
|
||||
self.prefill_metadata_cls = (
|
||||
FlashInferPrefillMetadata
|
||||
if self._use_fi_prefill
|
||||
else CudnnPrefillMetadata
|
||||
if self._use_cudnn_prefill
|
||||
else MLACommonPrefillMetadata
|
||||
)
|
||||
|
||||
if self._use_fi_prefill:
|
||||
self._workspace_buffer = torch.empty(
|
||||
envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
|
||||
self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
|
||||
|
||||
self._global_hyperparameters = infer_global_hyperparameters(
|
||||
get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl)
|
||||
)
|
||||
|
||||
if self._use_trtllm_ragged_prefill:
|
||||
self._workspace_buffer = torch.empty(
|
||||
envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
|
||||
if self._use_cudnn_prefill:
|
||||
self.cudnn_workspace = torch.empty(
|
||||
CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,
|
||||
dtype=torch.int8,
|
||||
device=device,
|
||||
)
|
||||
|
||||
supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
|
||||
self._init_reorder_batch_threshold(
|
||||
self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen
|
||||
)
|
||||
|
||||
# Validate consistency between query_len_support and reorder_batch_threshold
|
||||
if self.query_len_support == QueryLenSupport.SINGLE_ONLY:
|
||||
assert self.reorder_batch_threshold == 1, (
|
||||
f"reorder_batch_threshold must be 1 when query_len_support is "
|
||||
f"SINGLE_ONLY, got {self.reorder_batch_threshold}"
|
||||
)
|
||||
|
||||
|
||||
MluHijackObject.apply_hijack(MLACommonBackend,
|
||||
MLACommonBackend.get_supported_head_sizes,
|
||||
MLACommonBackend_MluHijack.get_supported_head_sizes)
|
||||
MluHijackObject.apply_hijack(MLACommonMetadataBuilder,
|
||||
MLACommonMetadataBuilder.__init__,
|
||||
MLACommonMetadataBuilder_MluHijack.__init__)
|
||||
|
||||
|
||||
class FlashMLABackend(MLACommonBackend):
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASHMLA_VLLM_V1"
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["FlashMLAMetadata"]:
|
||||
return FlashMLAMetadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["FlashMLAMetadataBuilder"]:
|
||||
return FlashMLAMetadataBuilder
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["FlashMLAImpl"]:
|
||||
return FlashMLAImpl
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_shape(
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
num_kv_heads: int, # assumed to be 1 for MLA
|
||||
head_size: int,
|
||||
cache_dtype_str: str = "auto",
|
||||
) -> tuple[int, ...]:
|
||||
return (1, num_blocks, num_kv_heads, block_size, head_size)
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_scale_shape(
|
||||
num_blocks: int,
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
) -> tuple[int, ...]:
|
||||
return (1, num_blocks, num_kv_heads, block_size)
|
||||
|
||||
@classmethod
|
||||
def get_supported_head_sizes(cls) -> list[int]:
|
||||
return [576, 512]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashMLAPrefillMetadata(MLACommonPrefillMetadata):
|
||||
num_prefills: int = -1 # for gather_cache
|
||||
max_seq_len: int = -1 # for attn forward
|
||||
|
||||
@property
|
||||
def block_tables(self):
|
||||
return self.block_table
|
||||
|
||||
@property
|
||||
def context_chunk_cu_seq_lens(self):
|
||||
if self.chunked_context is None:
|
||||
return None
|
||||
return self.chunked_context.cu_seq_lens
|
||||
|
||||
@property
|
||||
def context_chunk_starts(self):
|
||||
if self.chunked_context is None:
|
||||
return None
|
||||
return self.chunked_context.starts
|
||||
|
||||
@property
|
||||
def context_chunk_seq_tot(self):
|
||||
if self.chunked_context is None:
|
||||
return None
|
||||
return self.chunked_context.seq_tot
|
||||
|
||||
@property
|
||||
def context_chunk_max_seq_lens(self):
|
||||
if self.chunked_context is None:
|
||||
return None
|
||||
return self.chunked_context.max_seq_lens
|
||||
|
||||
@property
|
||||
def context_chunk_workspace(self):
|
||||
if self.chunked_context is None:
|
||||
return None
|
||||
return self.chunked_context.workspace
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
|
||||
tile_scheduler_metadata: torch.Tensor
|
||||
num_splits: torch.Tensor
|
||||
|
||||
# add for mlu rope and attn forward
|
||||
query_start_loc: torch.Tensor # for rope
|
||||
max_query_len: int # for rope
|
||||
max_seq_len:int = -1 # for attn forward
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlashMLAMetadata(MLACommonMetadata):
|
||||
num_prefill_tokens: Optional[int] = None
|
||||
|
||||
|
||||
class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
|
||||
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
|
||||
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
|
||||
reorder_batch_threshold: int = 128 # process small prefills with decode pathway
|
||||
# ^ TODO(matt): tune this
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kv_cache_spec: AttentionSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config: VllmConfig,
|
||||
device: torch.device,
|
||||
):
|
||||
super().__init__(
|
||||
kv_cache_spec, layer_names, vllm_config, device, FlashMLAMetadata
|
||||
)
|
||||
|
||||
self.num_q_heads = vllm_config.model_config.get_num_attention_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
|
||||
self.cg_buf_tile_scheduler_metadata = None
|
||||
self.cg_buf_num_splits = None
|
||||
self.is_fp8_kvcache = vllm_config.cache_config.cache_dtype.startswith("fp8")
|
||||
|
||||
self.cg_buf_tile_scheduler_metadata = None
|
||||
self.cg_buf_num_splits = None
|
||||
|
||||
'''
|
||||
=============================
|
||||
Modify by vllm_mlu
|
||||
=============================
|
||||
@brief: 1. set decoder_query_len for mtp
|
||||
@brief: 2. init chunk workspace for prefix_caching only
|
||||
@brief: 3. set prefill_metadata_cls
|
||||
@brief: 4. add deepseek v3.2 infos
|
||||
'''
|
||||
cache_config = vllm_config.cache_config
|
||||
scheduler_config = vllm_config.scheduler_config
|
||||
speculative_config = vllm_config.speculative_config
|
||||
self.num_speculative_tokens = (speculative_config.num_speculative_tokens
|
||||
if speculative_config is not None else 0)
|
||||
self.decoder_query_len = 1 + self.num_speculative_tokens
|
||||
|
||||
self.max_model_len = self.model_config.max_model_len
|
||||
self.is_deepseek_v32 = self.model_config.hf_text_config.model_type == "deepseek_v32"
|
||||
|
||||
self.enable_caching = cache_config.enable_prefix_caching
|
||||
self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
|
||||
if (not self.is_deepseek_v32 and not self.chunked_prefill_enabled and
|
||||
(mlu_envs.VLLM_V1_USE_UNCHUNK_SCHED and self.enable_caching)):
|
||||
self.chunked_prefill_workspace_size = min(
|
||||
# Max sure there is enough for 8 full length request or at least
|
||||
# 4 pages of cache per request
|
||||
max(
|
||||
8 * self.model_config.max_model_len, 4 *
|
||||
scheduler_config.max_num_seqs * cache_config.block_size),
|
||||
# For long-context models try not to over-allocate limiting
|
||||
# kv-cache space, limiting it to 64k tokens,
|
||||
# which would result in the workspace being:
|
||||
# 2*(576)*(64*1024) = 144mb
|
||||
# (assuming 576 MLA head dim, and fp16)
|
||||
# which would result in up-projected context being
|
||||
# 2*(192*128)*(64*1024) = 3gb
|
||||
# (assuming 192 QK head dim, 128 heads, and fp16)
|
||||
128 * 1024)
|
||||
assert self.chunked_prefill_workspace_size >= \
|
||||
scheduler_config.max_num_seqs * cache_config.block_size
|
||||
self.chunked_prefill_workspace = torch.empty(
|
||||
(self.chunked_prefill_workspace_size,
|
||||
self.model_config.get_head_size()),
|
||||
dtype=self.model_config.dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.prefill_metadata_cls = FlashMLAPrefillMetadata
|
||||
'''
|
||||
==================
|
||||
End of MLU Hijack
|
||||
==================
|
||||
'''
|
||||
|
||||
def reorder_batch(self, input_batch: "InputBatch",
|
||||
scheduler_output: "SchedulerOutput") -> bool:
|
||||
# We now want to reorder the batch so that the "decode" requests are and
|
||||
# the front and the "prefill" requests are at the using the least amount
|
||||
# swaps possible. (NOTE for now we loosely use "decode" to mean requests
|
||||
# where attention is likely memory-bound and "prefill" to mean requests
|
||||
# where attention is likely compute-bound, TODO(lucas): figure out a
|
||||
# better naming here)
|
||||
decodes = []
|
||||
prefills = []
|
||||
num_decode_tokens = 0
|
||||
num_prefill_tokens = 0
|
||||
# mlu v1 mtp forces decoder_query_len = 1 for k > 1, so we should set again
|
||||
self.decoder_query_len = 1 + self.num_speculative_tokens
|
||||
|
||||
for i, req_id in enumerate(input_batch.req_ids):
|
||||
num_tokens = scheduler_output.num_scheduled_tokens[req_id]
|
||||
# for now treat 1 scheduled token as "decode" even if its not,
|
||||
# we should update this to something like < 8 in the future but
|
||||
# currently the TritonMLA._forward_decode only supports
|
||||
# num_tokens = 1
|
||||
'''
|
||||
=============================
|
||||
Modify by vllm_mlu
|
||||
=============================
|
||||
@brief: record prefill and decode requests and token nums to call
|
||||
chunked fa and single-query attn respectively in forward.
|
||||
@Notes: decodes need all prompt tokens are computed.
|
||||
'''
|
||||
req_index = input_batch.req_id_to_index.get(req_id)
|
||||
all_prompt_tokens_has_computed = (
|
||||
input_batch.num_computed_tokens_cpu[req_index] >=
|
||||
input_batch.num_prompt_tokens[req_index])
|
||||
if num_tokens <= self.decoder_query_len and all_prompt_tokens_has_computed:
|
||||
decodes.append(i)
|
||||
num_decode_tokens += num_tokens
|
||||
else:
|
||||
prefills.append(i)
|
||||
num_prefill_tokens += num_tokens
|
||||
'''
|
||||
==================
|
||||
End of MLU Hijack
|
||||
==================
|
||||
'''
|
||||
|
||||
# We hope that this is fairly minimal since decodes
|
||||
# should be around for a number of iterations so hopefully they are
|
||||
# relatively stationary (and new request are generally appended to the
|
||||
# persistent batch so already should be at the back)
|
||||
# To achieve this we loop over the decodes in descending order and
|
||||
# the prefills in ascending order. We swap decodes from the "back"
|
||||
# i.e. past where the last decode should be in the reodorered with
|
||||
# prefills from the front of the batch.
|
||||
# `decodes` and `prefills` are already in ascending order just based on
|
||||
# the above loop
|
||||
num_decodes = len(decodes)
|
||||
num_prefills = len(prefills)
|
||||
modified_batch = False
|
||||
|
||||
for i in range(1, min(num_decodes, num_prefills) + 1):
|
||||
# If the decode is at the "back" of the batch, i, we can swap it
|
||||
# with the prefill closest to the front of the batch
|
||||
decode_idx = decodes[num_decodes - i]
|
||||
if decode_idx < num_decodes:
|
||||
break
|
||||
|
||||
input_batch.swap_states(prefills[i - 1], decode_idx)
|
||||
modified_batch = True
|
||||
|
||||
return modified_batch
|
||||
|
||||
def _build_decode(
|
||||
self,
|
||||
block_table_tensor: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
query_start_loc: torch.Tensor,
|
||||
max_query_len: int,
|
||||
max_seq_len: int,
|
||||
) -> FlashMLADecodeMetadata:
|
||||
'''
|
||||
=============================
|
||||
Modify by vllm_mlu
|
||||
=============================
|
||||
@brief: set tile_scheduler_metadata and num_splits to None.
|
||||
@brief: set dcp_tot_seq_lens_device.
|
||||
'''
|
||||
return FlashMLADecodeMetadata(
|
||||
block_table=block_table_tensor,
|
||||
seq_lens=seq_lens,
|
||||
tile_scheduler_metadata=None,
|
||||
num_splits=None,
|
||||
dcp_tot_seq_lens=None,
|
||||
# for mlu
|
||||
max_seq_len=max_seq_len,
|
||||
query_start_loc=query_start_loc,
|
||||
max_query_len=max_query_len
|
||||
)
|
||||
'''
|
||||
==================
|
||||
End of MLU Hijack
|
||||
==================
|
||||
'''
|
||||
|
||||
def build_for_cudagraph_capture(
|
||||
self, common_attn_metadata: MLUCommonAttentionMetadata) -> M:
|
||||
"""
|
||||
This method builds the metadata for full cudagraph capture.
|
||||
Currently, only decode is supported for full cudagraphs with MLA.
|
||||
"""
|
||||
m = common_attn_metadata
|
||||
if m.infer_mode == MLUInferMode.DECODE_ONLY:
|
||||
assert m.num_reqs * m.max_query_len == m.num_actual_tokens, \
|
||||
"MLA only supports decode-only full CUDAGraph capture. " \
|
||||
"Make sure all cudagraph capture sizes <= max_num_seq."
|
||||
|
||||
return self.build(0, m)
|
||||
|
||||
def build(self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: MLUCommonAttentionMetadata,
|
||||
fast_build: bool = False,
|
||||
input_batch: "InputBatch" = None) -> M:
|
||||
num_reqs = common_attn_metadata.num_reqs
|
||||
num_tokens = common_attn_metadata.num_actual_tokens
|
||||
max_query_len = common_attn_metadata.max_query_len
|
||||
|
||||
# Note(simon): be careful about the CPU <> GPU memory movement in this
|
||||
# function. We should avoid GPU -> CPU sync as much as possible because
|
||||
# it blocks on all previous kernels.
|
||||
device = self.device
|
||||
block_table_tensor = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
|
||||
query_start_loc = common_attn_metadata.query_start_loc
|
||||
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
|
||||
seq_lens = common_attn_metadata.seq_lens
|
||||
|
||||
query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
|
||||
|
||||
num_computed_tokens_cpu = (common_attn_metadata.seq_lens_cpu -
|
||||
query_seq_lens_cpu)
|
||||
'''
|
||||
=============================
|
||||
Modify by vllm_mlu
|
||||
=============================
|
||||
@brief: support normal and mtp input split
|
||||
'''
|
||||
if input_batch is None:
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = \
|
||||
split_decodes_and_prefills(common_attn_metadata,
|
||||
self.decoder_query_len)
|
||||
else:
|
||||
num_decodes, num_prefills = input_batch.split_decodes_and_prefills()
|
||||
num_decode_tokens = common_attn_metadata.query_start_loc_cpu[num_decodes].item()
|
||||
num_prefill_tokens = num_tokens - num_decode_tokens
|
||||
'''
|
||||
==================
|
||||
End of MLU Hijack
|
||||
==================
|
||||
'''
|
||||
assert num_decodes + num_prefills == num_reqs
|
||||
assert num_decode_tokens + num_prefill_tokens == num_tokens
|
||||
|
||||
prefill_metadata = None
|
||||
if num_prefills > 0:
|
||||
reqs_start = num_decodes # prefill_start
|
||||
|
||||
context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
|
||||
max_context_len_cpu = context_lens_cpu.max().item()
|
||||
num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
|
||||
'''
|
||||
=============================
|
||||
Modify by vllm_mlu
|
||||
=============================
|
||||
@brief: avoid buffer missing when prefill_only + mlugraph
|
||||
'''
|
||||
if num_decodes > 0:
|
||||
prefill_query_start_loc = query_start_loc[
|
||||
reqs_start:] - query_start_loc[reqs_start]
|
||||
else:
|
||||
prefill_query_start_loc= query_start_loc
|
||||
'''
|
||||
==================
|
||||
End of MLU Hijack
|
||||
==================
|
||||
'''
|
||||
|
||||
chunked_context_metadata = None
|
||||
if ((self.chunked_prefill_enabled or
|
||||
(mlu_envs.VLLM_V1_USE_UNCHUNK_SCHED and
|
||||
self.enable_caching and
|
||||
common_attn_metadata.is_chunked)
|
||||
) and num_prefills > 0 and max_context_len_cpu > 0):
|
||||
# NOTE: it is recommend you read the `Chunked Prefill` section
|
||||
# in the comment at the top of the file before trying to
|
||||
# understand the following code
|
||||
|
||||
# currently we allocate an equal amount of workspace for each
|
||||
# prefill in the batch, we could probably use a more advanced
|
||||
# algorithm here and allocate more workspace to prefills with
|
||||
# longer context lengths
|
||||
if self.is_deepseek_v32:
|
||||
max_context_chunk = self.max_model_len
|
||||
else:
|
||||
max_context_chunk = (self.chunked_prefill_workspace_size //
|
||||
num_prefills_with_context_cpu)
|
||||
|
||||
if self.aot_schedule:
|
||||
# align max_context_chunk to page_size by rounding down,
|
||||
# currently the `gather_cache` kernel cannot handle
|
||||
# `context_chunk_starts` that are not aligned to page_size
|
||||
max_context_chunk = round_down(max_context_chunk,
|
||||
self.page_size)
|
||||
|
||||
assert max_context_chunk > 0
|
||||
num_chunks = cdiv(max_context_len_cpu, max_context_chunk)
|
||||
|
||||
# if `max_context_chunk = 256`, `num_chunks = 3`, and
|
||||
# `num_prefills_with_context = 4`, create a tensor that looks
|
||||
# like
|
||||
# [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
|
||||
# Note(simon): this is done in CPU because of downstream's
|
||||
# of `to_list`.
|
||||
chunk_starts = \
|
||||
torch.arange(num_chunks, dtype=torch.int32) \
|
||||
.unsqueeze(1).expand(-1, num_prefills) \
|
||||
* max_context_chunk
|
||||
chunk_ends = torch.min(context_lens_cpu.unsqueeze(0),
|
||||
chunk_starts + max_context_chunk)
|
||||
chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
|
||||
|
||||
cu_seq_lens_cpu = torch.zeros(num_chunks,
|
||||
num_prefills + 1,
|
||||
dtype=torch.int32,
|
||||
pin_memory=True)
|
||||
torch.cumsum(chunk_seq_lens,
|
||||
dim=1,
|
||||
out=cu_seq_lens_cpu[:, 1:],
|
||||
dtype=torch.int32)
|
||||
|
||||
chunked_context_metadata_cls = \
|
||||
FlashMLAPrefillMetadata.ChunkedContextMetadata
|
||||
|
||||
chunked_context_metadata = \
|
||||
chunked_context_metadata_cls(
|
||||
cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
|
||||
starts=chunk_starts.to(device, non_blocking=True),
|
||||
seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
|
||||
max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
|
||||
seq_lens=chunk_seq_lens,
|
||||
workspace=getattr(self, "chunked_prefill_workspace", None),
|
||||
)
|
||||
|
||||
if not self.is_deepseek_v32:
|
||||
assert max(chunked_context_metadata.max_seq_lens) <= \
|
||||
self.chunked_prefill_workspace_size
|
||||
|
||||
prefill_metadata = self.prefill_metadata_cls(
|
||||
block_table=block_table_tensor[reqs_start:, ...],
|
||||
query_start_loc=prefill_query_start_loc,
|
||||
max_query_len=max_query_len,
|
||||
chunked_context=chunked_context_metadata,
|
||||
# for mlu
|
||||
num_prefills=num_prefills,
|
||||
max_seq_len=common_attn_metadata.seq_lens_cpu[reqs_start:].max().item(),
|
||||
)
|
||||
|
||||
decode_metadata = None
|
||||
if num_decodes > 0:
|
||||
decode_metadata = self._build_decode(
|
||||
block_table_tensor=block_table_tensor[:num_decodes, ...],
|
||||
seq_lens=seq_lens[:num_decodes],
|
||||
query_start_loc=query_start_loc[:num_decodes + 1],
|
||||
max_query_len=query_seq_lens_cpu[:num_decodes].max().item(),
|
||||
max_seq_len=common_attn_metadata.seq_lens_cpu[:num_decodes].max().item(),
|
||||
)
|
||||
|
||||
attn_metadata = self.metadata_cls(
|
||||
num_reqs=common_attn_metadata.num_reqs,
|
||||
max_query_len=common_attn_metadata.max_query_len,
|
||||
max_seq_len=common_attn_metadata.max_seq_len,
|
||||
num_actual_tokens=num_tokens,
|
||||
query_start_loc=query_start_loc,
|
||||
slot_mapping=slot_mapping,
|
||||
head_dim=self.model_config.get_head_size(),
|
||||
# MLACommonMetadata Chunk prefill specific
|
||||
num_decodes=num_decodes,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_prefills=num_prefills,
|
||||
prefill=prefill_metadata,
|
||||
decode=decode_metadata,
|
||||
)
|
||||
|
||||
return attn_metadata
|
||||
|
||||
def can_run_in_cudagraph(
|
||||
self, common_attn_metadata: MLUCommonAttentionMetadata) -> bool:
|
||||
return common_attn_metadata.max_query_len == self.decoder_query_len
|
||||
|
||||
def use_cascade_attention(self, *args, **kwargs) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class FlashMLAImpl(MLUFlashAttentionImpl):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
scale: float,
|
||||
num_kv_heads: int,
|
||||
alibi_slopes: Optional[list[float]],
|
||||
sliding_window: Optional[int],
|
||||
kv_cache_dtype: str,
|
||||
logits_soft_cap: Optional[float],
|
||||
attn_type: str,
|
||||
kv_sharing_target_layer_name: Optional[str],
|
||||
# MLA Specific Arguments
|
||||
**mla_args) -> None:
|
||||
super().__init__(num_heads, head_size, scale, num_kv_heads,
|
||||
alibi_slopes, sliding_window, kv_cache_dtype,
|
||||
logits_soft_cap, attn_type,
|
||||
kv_sharing_target_layer_name, **mla_args)
|
||||
|
||||
unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
|
||||
if any(unsupported_features):
|
||||
raise NotImplementedError(
|
||||
"FlashMLAImpl does not support one of the following: "
|
||||
"alibi_slopes, sliding_window, logits_soft_cap")
|
||||
|
||||
if attn_type != AttentionType.DECODER:
|
||||
raise NotImplementedError("Encoder self-attention and "
|
||||
"encoder/decoder cross-attention "
|
||||
"are not implemented for "
|
||||
"FlashMLAImpl")
|
||||
|
||||
def forward(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_metadata: FlashMLAMetadata,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
output_scale: Optional[torch.Tensor] = None,
|
||||
kwargs: Optional[dict[str, Any]] = {},
|
||||
) -> torch.Tensor:
|
||||
assert output is not None, "Output tensor must be provided."
|
||||
|
||||
if output_scale is not None:
|
||||
raise NotImplementedError(
|
||||
"fused output quantization is not yet supported"
|
||||
" for FlashAttentionImpl")
|
||||
|
||||
if attn_metadata is None:
|
||||
# Profiling run.
|
||||
return output
|
||||
|
||||
out_lse = None
|
||||
|
||||
# use default common metadata if kwargs does not have common_metadata
|
||||
common_metadata: MLUCommonAttentionMetadata = kwargs.get("common_metadata", None)
|
||||
if common_metadata is None:
|
||||
common_metadata = get_common_metadata()
|
||||
|
||||
only_prefill = kwargs.get("only_prefill", False)
|
||||
only_decode = kwargs.get("only_decode", False)
|
||||
attn_bias = kwargs.get("attn_bias", None)
|
||||
|
||||
assert only_prefill != only_decode, "only_prefill and only_decode cannot be True and False at the same time."
|
||||
|
||||
if only_prefill:
|
||||
cu_seqlens_q = attn_metadata.prefill.query_start_loc
|
||||
cu_seqlens_kv = common_metadata.query_start_loc
|
||||
seqused_k = common_metadata.seq_lens[attn_metadata.num_decodes:]
|
||||
max_seqlen_q = attn_metadata.prefill.max_query_len
|
||||
max_seqlen_k = attn_metadata.prefill.max_seq_len
|
||||
block_table = attn_metadata.prefill.block_table
|
||||
num_actual_tokens = attn_metadata.num_prefill_tokens
|
||||
else:
|
||||
cu_seqlens_q = None # nouse
|
||||
cu_seqlens_kv = None # nouse
|
||||
seqused_k = common_metadata.seq_lens[:attn_metadata.num_decodes]
|
||||
max_seqlen_q = None # nouse
|
||||
max_seqlen_k = common_metadata.max_seq_len
|
||||
block_table = attn_metadata.decode.block_table
|
||||
num_actual_tokens = attn_metadata.num_decode_tokens
|
||||
|
||||
skip_process_cache = ((self.use_mla
|
||||
and (common_metadata.is_prefill_only
|
||||
or self.use_fused_mla_qkv
|
||||
or only_prefill))
|
||||
or self.kv_sharing_target_layer_name is not None)
|
||||
|
||||
kv_cache_, kv_cache_scale_, kv_cache_index_ = kv_cache
|
||||
key_cache = kv_cache_[0]
|
||||
value_cache = None if self.use_mla else kv_cache_[1]
|
||||
key_cache_scale, value_cache_scale = None, None
|
||||
if kv_cache_scale_.numel() > 0:
|
||||
key_cache_scale = kv_cache_scale_[0]
|
||||
value_cache_scale = None if self.use_mla else kv_cache_scale_[1]
|
||||
if not skip_process_cache:
|
||||
if is_quantized_kv_cache(self.kv_cache_dtype):
|
||||
mlu_ops.quant_to_paged_cache(
|
||||
k=key[:num_actual_tokens],
|
||||
v=(None if self.use_mla else value[:num_actual_tokens]),
|
||||
k_cache=key_cache,
|
||||
v_cache=value_cache,
|
||||
k_cache_quant_scale=key_cache_scale,
|
||||
v_cache_quant_scale=value_cache_scale,
|
||||
slot_mapping=attn_metadata.slot_mapping.flatten(),
|
||||
)
|
||||
else:
|
||||
mlu_ops.reshape_paged_cache(
|
||||
k=key[:num_actual_tokens],
|
||||
v=(None if self.use_mla else value[:num_actual_tokens]),
|
||||
k_cache=key_cache,
|
||||
v_cache=value_cache,
|
||||
slot_mapping=attn_metadata.slot_mapping.flatten()
|
||||
)
|
||||
|
||||
alibi_slopes = None if self.alibi_slopes is None else \
|
||||
self.alibi_slopes.repeat(seqused_k.shape[0], 1)
|
||||
|
||||
if kwargs.get("model_type", "") == "deepseek_v32":
|
||||
from vllm_mlu.model_executor.models.sp_utils import get_sp_forward_context
|
||||
sp_context = get_sp_forward_context()
|
||||
if sp_context is not None and sp_context.is_v32:
|
||||
num_actual_tokens = sp_context.sp_attn_metadata.num_prefill_tokens
|
||||
decode_query = query[:num_actual_tokens].view(-1, self.num_heads, self.head_size)
|
||||
head_size_v = value.shape[-1] if self.use_mla else self.head_size
|
||||
decode_output = output[:num_actual_tokens].view(-1, self.num_heads, head_size_v)
|
||||
decode_query = query.unsqueeze(1) # see tokens as batch dim
|
||||
decode_output = decode_output.unsqueeze(1)
|
||||
q_quant_scale = kwargs.get("q_quant_scale", None)
|
||||
if q_quant_scale is not None:
|
||||
q_quant_scale = q_quant_scale[:num_actual_tokens].view(-1, self.num_heads)
|
||||
q_quant_scale = q_quant_scale.unsqueeze(1)
|
||||
mlu_ops.single_query_cached_kv_attn(
|
||||
q=decode_query,
|
||||
k_cache=key_cache,
|
||||
v_cache=value_cache,
|
||||
out=decode_output,
|
||||
block_tables=kwargs.get("new_block_tables", None),
|
||||
context_lens=kwargs.get("new_context_lens", None),
|
||||
k_cache_quant_scale=key_cache_scale,
|
||||
v_cache_quant_scale=value_cache_scale,
|
||||
alibi_slopes=alibi_slopes,
|
||||
max_contxt_len=kwargs.get("index_topk", None),
|
||||
windows_size_left=(-1 if self.sliding_window is None else self.sliding_window[0]),
|
||||
windows_size_right=(-1 if self.sliding_window is None else self.sliding_window[0]),
|
||||
softmax_scale=self.scale,
|
||||
head_size_v=(-1 if not self.use_mla else head_size_v),
|
||||
compute_dtype=compute_dtype,
|
||||
q_quant_scale=q_quant_scale,
|
||||
decoder_attn_dtype=self.decoder_attn_dtype,
|
||||
mask=attn_bias,
|
||||
)
|
||||
return output
|
||||
|
||||
if common_metadata.is_prefill_only or only_prefill:
|
||||
# prefill only
|
||||
prefill_causal = kwargs.get("prefill_causal", True)
|
||||
cu_seqlens_q = kwargs.get("cu_seq_lens_q", cu_seqlens_q)
|
||||
cu_seqlens_kv = kwargs.get("cu_seq_lens_kv", cu_seqlens_kv)
|
||||
max_seqlen_q = kwargs.get("max_seq_len_q", max_seqlen_q)
|
||||
max_seqlen_k = kwargs.get("max_seq_len_kv", max_seqlen_k)
|
||||
return_lse = kwargs.get("return_lse", False)
|
||||
num_prefill_query_tokens = common_metadata.num_prefill_query_tokens
|
||||
num_prefill_kv_tokens = common_metadata.num_prefill_kv_tokens
|
||||
use_f32 = attn_bias is not None and attn_bias.dtype == torch.float32
|
||||
if use_f32:
|
||||
f32_output = torch.empty_like(output, dtype=torch.float32)
|
||||
attn_output_list = mlu_ops.flash_attention(
|
||||
q=query[:num_prefill_query_tokens].to(torch.float32) if use_f32 else query[:num_prefill_query_tokens],
|
||||
k=key[:num_prefill_kv_tokens].to(torch.float32) if use_f32 else key[:num_prefill_kv_tokens],
|
||||
v=value[:num_prefill_kv_tokens].to(torch.float32) if use_f32 else value[:num_prefill_kv_tokens],
|
||||
out=f32_output[:num_prefill_query_tokens] if use_f32 else output[:num_prefill_query_tokens],
|
||||
cu_seq_lens_q=cu_seqlens_q,
|
||||
cu_seq_lens_kv=cu_seqlens_kv,
|
||||
alibi_slope=alibi_slopes,
|
||||
attn_bias=attn_bias,
|
||||
max_seq_len_q=max_seqlen_q,
|
||||
max_seq_len_kv=max_seqlen_k,
|
||||
softmax_scale=self.scale,
|
||||
is_causal=prefill_causal,
|
||||
window_size_left=(-1 if self.sliding_window is None else self.sliding_window[0]),
|
||||
window_size_right=(-1 if self.sliding_window is None else self.sliding_window[1]),
|
||||
compute_dtype=self.prefill_compute_dtype,
|
||||
return_lse=return_lse,
|
||||
q_quant_dtype=self.prefill_q_dtype,
|
||||
k_quant_dtype=self.prefill_k_dtype,
|
||||
v_quant_dtype=self.prefill_v_dtype
|
||||
)
|
||||
if use_f32:
|
||||
output[:num_prefill_query_tokens].copy_(f32_output[:num_prefill_query_tokens])
|
||||
|
||||
if return_lse:
|
||||
out_lse = attn_output_list[1]
|
||||
else:
|
||||
batch_size = block_table.shape[0]
|
||||
# decode only
|
||||
decode_query = query[:num_actual_tokens].view(batch_size, -1, self.num_heads, self.head_size)
|
||||
head_size_v = value.shape[-1] if self.use_mla else self.head_size
|
||||
decode_output = output[:num_actual_tokens].view(batch_size, -1, self.num_heads, head_size_v)
|
||||
q_quant_scale = kwargs.get("q_quant_scale", None)
|
||||
if q_quant_scale is not None:
|
||||
q_quant_scale = q_quant_scale[:num_actual_tokens].view(batch_size, -1, self.num_heads)
|
||||
mlu_ops.single_query_cached_kv_attn(
|
||||
q=decode_query,
|
||||
k_cache=key_cache,
|
||||
v_cache=value_cache,
|
||||
out=decode_output,
|
||||
block_tables=block_table,
|
||||
context_lens=seqused_k,
|
||||
k_cache_quant_scale=key_cache_scale,
|
||||
v_cache_quant_scale=value_cache_scale,
|
||||
alibi_slopes=alibi_slopes,
|
||||
max_contxt_len=max_seqlen_k,
|
||||
windows_size_left=(-1 if self.sliding_window is None else self.sliding_window[0]),
|
||||
windows_size_right=(-1 if self.sliding_window is None else self.sliding_window[0]),
|
||||
softmax_scale=self.scale,
|
||||
head_size_v=(-1 if not self.use_mla else head_size_v),
|
||||
compute_dtype=attn_metadata.decode.compute_dtype,
|
||||
q_quant_scale=q_quant_scale,
|
||||
decoder_attn_dtype=self.decoder_attn_dtype,
|
||||
mask=attn_bias,
|
||||
)
|
||||
|
||||
return output if out_lse is None else (output, out_lse)
|
||||
Reference in New Issue
Block a user