diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py new file mode 100644 index 000000000..26a59a6cc --- /dev/null +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -0,0 +1,232 @@ +# Mixin class for metadata management of Deepseek MHA forward (chunked prefix cache) +# More details can be found in python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py + +from typing import List, Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton + + +class ForwardBatchDeepSeekMHAMixin: + # For MLA chunked prefix cache used in chunked prefill + # Tell attention backend whether the kv cache needs to be attended in current pass + attn_attend_prefix_cache: Optional[bool] = None + # Number of prefix cache chunks + num_prefix_chunks: Optional[int] = None + # Index of current chunk, used by attention backend + prefix_chunk_idx: Optional[int] = None + # Maximum number of tokens in each chunk per sequence. Computed from maximum chunk capacity + prefix_chunk_len: Optional[int] = None + # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size) + prefix_chunk_starts: Optional[torch.Tensor] = None + # Lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size) + prefix_chunk_seq_lens: Optional[torch.Tensor] = None + # Accumulated lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size + 1) + prefix_chunk_cu_seq_lens: Optional[torch.Tensor] = None + # Max lengths of prefix cache for each chunk, (num_prefix_chunks,) + prefix_chunk_max_seq_lens: Optional[List[int]] = None + # Number of tokens in each prefix cache chunk, (num_prefix_chunks,) + prefix_chunk_num_tokens: Optional[List[int]] = None + # KV Indices for each chunk + prefix_chunk_kv_indices: Optional[List[torch.Tensor]] = None + # For MLA chunked prefix cache used in chunked prefill + # Tell attention backend whether lse needs to be returned + mha_return_lse: Optional[bool] = None + # Whether to apply MHA_ONE_SHOT forward method + mha_one_shot: Optional[bool] = None + # KV Indices for MHA_ONE_SHOT forward method + mha_one_shot_kv_indices: Optional[torch.Tensor] = None + + def get_max_chunk_capacity(self): + # Maximum number of tokens in each chunk + # TODO: Should be changed to a better value, maybe passed through server args + return 128 * 1024 + + def set_prefix_chunk_idx(self, idx: int): + self.prefix_chunk_idx = idx + + def set_attn_attend_prefix_cache(self, attn_attend_prefix_cache: bool): + self.attn_attend_prefix_cache = attn_attend_prefix_cache + + def prepare_chunked_kv_indices(self, device: torch.device): + self.prefix_chunk_kv_indices = [] + for idx in range(self.num_prefix_chunks): + chunk_starts = self.prefix_chunk_starts[idx] + chunk_seq_lens = self.prefix_chunk_seq_lens[idx] + chunk_cu_seq_lens = self.prefix_chunk_cu_seq_lens[idx] + num_chunk_tokens = self.prefix_chunk_num_tokens[idx] + + chunk_kv_indices = torch.empty( + num_chunk_tokens, dtype=torch.int32, device=device + ) + + create_chunked_prefix_cache_kv_indices[(self.batch_size,)]( + self.req_to_token_pool.req_to_token, + self.req_pool_indices, + chunk_starts, + chunk_seq_lens, + chunk_cu_seq_lens, + chunk_kv_indices, + self.req_to_token_pool.req_to_token.shape[1], + ) + self.prefix_chunk_kv_indices.append(chunk_kv_indices) + + # Here we suppose the length of each chunk is equal + # For example, if we have 4 sequences with prefix length [256, 512, 768, 1024], prefix_chunk_len = 256 + # num_prefix_chunks = cdiv(1024, 256) = 4 + # prefix_chunk_starts = [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512], [768, 768, 768, 768]] + # prefix_chunk_ends = [[256, 256, 256, 256], [256, 512, 512, 512], [256, 512, 768, 768], [256, 512, 768, 1024]] + # prefix_chunk_seq_lens = [[256, 256, 256, 256], [0, 256, 256, 256], [0, 0, 256, 256], [0, 0, 0, 256]] + # TODO: Implement a better way to allocate chunk lengths that uses memory spaces more efficiently. + def get_prefix_chunk_seq_lens( + self, prefix_lens: torch.Tensor, num_prefix_chunks: int, prefix_chunk_len: int + ): + device = prefix_lens.device + prefix_chunk_starts = ( + torch.arange(num_prefix_chunks, device=device, dtype=torch.int32) + .unsqueeze(1) + .expand(-1, self.batch_size) + * prefix_chunk_len + ) + prefix_chunk_ends = torch.min( + prefix_lens.unsqueeze(0), + prefix_chunk_starts + prefix_chunk_len, + ).to(torch.int32) + + prefix_chunk_seq_lens = ( + (prefix_chunk_ends - prefix_chunk_starts).clamp(min=0).to(torch.int32) + ) + + return prefix_chunk_starts, prefix_chunk_seq_lens + + # Called before each attention module if using chunked kv cache for prefill + # Some of the codes are adapted from https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py + def prepare_chunked_prefix_cache_info(self, device: torch.device): + + from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool + + assert isinstance( + self.token_to_kv_pool, MLATokenToKVPool + ), "Currently chunked prefix cache can only be used by Deepseek models" + + if not any(self.extend_prefix_lens_cpu): + self.num_prefix_chunks = 0 + return + + if self.prefix_chunk_len is not None: + # Chunked kv cache info already prepared by prior modules + return + + self.prefix_chunk_idx = -1 + + # chunk_capacity is the maximum number of tokens in each chunk + chunk_capacity = self.get_max_chunk_capacity() + self.prefix_chunk_len = chunk_capacity // self.batch_size + + self.num_prefix_chunks = ( + max(self.extend_prefix_lens_cpu) + self.prefix_chunk_len - 1 + ) // self.prefix_chunk_len + + # Here we compute chunk lens twice to avoid stream sync, once on gpu and once on cpu. + prefix_chunk_starts_cuda, prefix_chunk_seq_lens_cuda = ( + self.get_prefix_chunk_seq_lens( + self.extend_prefix_lens, + self.num_prefix_chunks, + self.prefix_chunk_len, + ) + ) + _, prefix_chunk_seq_lens_cpu = self.get_prefix_chunk_seq_lens( + torch.tensor(self.extend_prefix_lens_cpu), + self.num_prefix_chunks, + self.prefix_chunk_len, + ) + self.prefix_chunk_starts = prefix_chunk_starts_cuda + self.prefix_chunk_seq_lens = prefix_chunk_seq_lens_cuda + + # Metadata for attention backend + self.prefix_chunk_cu_seq_lens = torch.zeros( + self.num_prefix_chunks, + self.batch_size + 1, + device=device, + dtype=torch.int32, + ) + self.prefix_chunk_cu_seq_lens[:, 1:] = prefix_chunk_seq_lens_cuda.cumsum( + dim=1 + ).to(torch.int32) + self.prefix_chunk_max_seq_lens = prefix_chunk_seq_lens_cpu.max( + dim=1 + ).values.tolist() + + self.prefix_chunk_num_tokens = prefix_chunk_seq_lens_cpu.sum(dim=1).tolist() + assert max(self.prefix_chunk_num_tokens) <= self.get_max_chunk_capacity() + + # Precompute the kv indices for each chunk + self.prepare_chunked_kv_indices(device) + + def fetch_mha_one_shot_kv_indices(self): + if self.mha_one_shot_kv_indices is not None: + return self.mha_one_shot_kv_indices + batch_size = self.batch_size + paged_kernel_lens_sum = sum(self.seq_lens_cpu) + kv_indices = torch.empty( + paged_kernel_lens_sum, + dtype=torch.int32, + device=self.req_pool_indices.device, + ) + kv_indptr = torch.zeros( + batch_size + 1, + dtype=torch.int32, + device=self.req_pool_indices.device, + ) + kv_indptr[1:] = torch.cumsum(self.seq_lens, dim=0) + create_flashinfer_kv_indices_triton[(self.batch_size,)]( + self.req_to_token_pool.req_to_token, + self.req_pool_indices, + self.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token_pool.req_to_token.shape[1], + ) + self.mha_one_shot_kv_indices = kv_indices + return kv_indices + + +@triton.jit +def create_chunked_prefix_cache_kv_indices( + req_to_token_ptr, # (max_batch, max_context_len,) + req_pool_indices_ptr, # (batch_size,) + chunk_start_idx_ptr, # (batch_size,) + chunk_seq_lens_ptr, # (batch_size,) + chunk_cu_seq_lens_ptr, # (batch_size + 1,) + chunk_kv_indices_ptr, # (num_chunk_tokens,) + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + # find the req pool idx, this is for batch to token + req_pool_index = tl.load(req_pool_indices_ptr + pid) + chunk_kv_indices_offset = tl.load(chunk_cu_seq_lens_ptr + pid) + + # get the token positions of current chunk + chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32) + chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < chunk_seq_len + data = tl.load( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + chunk_start_pos + + offset, + mask=mask, + ) + tl.store( + chunk_kv_indices_ptr + chunk_kv_indices_offset + offset, data, mask=mask + ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 4402d0928..ccabb02a2 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -43,7 +43,6 @@ from sglang.srt.distributed.parallel_state import ( get_tensor_model_parallel_world_size, ) from sglang.srt.layers.attention.nsa.utils import NSAContextParallelMetadata -from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.layers.dp_attention import ( DpPaddingMode, get_attention_dp_rank, @@ -52,6 +51,9 @@ from sglang.srt.layers.dp_attention import ( set_dp_buffer_len, set_is_extend_in_batch, ) +from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( + ForwardBatchDeepSeekMHAMixin, +) from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import get_compiler_backend, is_hip, is_npu, support_triton from sglang.srt.utils.common import ceil_align @@ -229,7 +231,7 @@ def compute_local_num_token_non_padded( @dataclass -class ForwardBatch: +class ForwardBatch(ForwardBatchDeepSeekMHAMixin): """Store all inputs of a forward pass.""" # The forward mode @@ -297,33 +299,6 @@ class ForwardBatch: # current split index of layer split_index: int = 0 - # For MLA chunked prefix cache used in chunked prefill - # Tell attention backend whether the kv cache needs to be attended in current pass - attn_attend_prefix_cache: Optional[bool] = None - # Number of prefix cache chunks - num_prefix_chunks: Optional[int] = None - # Index of current chunk, used by attention backend - prefix_chunk_idx: Optional[int] = None - # Maximum number of tokens in each chunk per sequence. Computed from maximum chunk capacity - prefix_chunk_len: Optional[int] = None - # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size) - prefix_chunk_starts: Optional[torch.Tensor] = None - # Lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size) - prefix_chunk_seq_lens: Optional[torch.Tensor] = None - # Accumulated lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size + 1) - prefix_chunk_cu_seq_lens: Optional[torch.Tensor] = None - # Max lengths of prefix cache for each chunk, (num_prefix_chunks,) - prefix_chunk_max_seq_lens: Optional[List[int]] = None - # Number of tokens in each prefix cache chunk, (num_prefix_chunks,) - prefix_chunk_num_tokens: Optional[List[int]] = None - # KV Indices for each chunk - prefix_chunk_kv_indices: Optional[List[torch.Tensor]] = None - # For MLA chunked prefix cache used in chunked prefill - # Tell attention backend whether lse needs to be returned - mha_return_lse: Optional[bool] = None - mha_one_shot_kv_indices: Optional[torch.Tensor] = None - mha_one_shot: Optional[bool] = None - # For multimodal mm_inputs: Optional[List[MultimodalInputs]] = None @@ -738,40 +713,6 @@ class ForwardBatch: dim=1, ).to(dtype=torch.int64, device=model_runner.device, non_blocking=True) - def get_max_chunk_capacity(self): - # Maximum number of tokens in each chunk - # TODO: Should be changed to a better value, maybe passed through server args - return 128 * 1024 - - def set_prefix_chunk_idx(self, idx: int): - self.prefix_chunk_idx = idx - - def set_attn_attend_prefix_cache(self, attn_attend_prefix_cache: bool): - self.attn_attend_prefix_cache = attn_attend_prefix_cache - - def prepare_chunked_kv_indices(self, device: torch.device): - self.prefix_chunk_kv_indices = [] - for idx in range(self.num_prefix_chunks): - chunk_starts = self.prefix_chunk_starts[idx] - chunk_seq_lens = self.prefix_chunk_seq_lens[idx] - chunk_cu_seq_lens = self.prefix_chunk_cu_seq_lens[idx] - num_chunk_tokens = self.prefix_chunk_num_tokens[idx] - - chunk_kv_indices = torch.empty( - num_chunk_tokens, dtype=torch.int32, device=device - ) - - create_chunked_prefix_cache_kv_indices[(self.batch_size,)]( - self.req_to_token_pool.req_to_token, - self.req_pool_indices, - chunk_starts, - chunk_seq_lens, - chunk_cu_seq_lens, - chunk_kv_indices, - self.req_to_token_pool.req_to_token.shape[1], - ) - self.prefix_chunk_kv_indices.append(chunk_kv_indices) - def _pad_tensor_to_size(self, tensor: torch.Tensor, size: int, *, value: int = 0): if value == 0: return torch.cat( @@ -1005,130 +946,10 @@ class ForwardBatch: if logits_output.hidden_states is not None: logits_output.hidden_states = logits_output.hidden_states[:num_tokens] - # Here we suppose the length of each chunk is equal - # For example, if we have 4 sequences with prefix length [256, 512, 768, 1024], prefix_chunk_len = 256 - # num_prefix_chunks = cdiv(1024, 256) = 4 - # prefix_chunk_starts = [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512], [768, 768, 768, 768]] - # prefix_chunk_ends = [[256, 256, 256, 256], [256, 512, 512, 512], [256, 512, 768, 768], [256, 512, 768, 1024]] - # prefix_chunk_seq_lens = [[256, 256, 256, 256], [0, 256, 256, 256], [0, 0, 256, 256], [0, 0, 0, 256]] - # TODO: Implement a better way to allocate chunk lengths that uses memory spaces more efficiently. - def get_prefix_chunk_seq_lens( - self, prefix_lens: torch.Tensor, num_prefix_chunks: int, prefix_chunk_len: int - ): - device = prefix_lens.device - prefix_chunk_starts = ( - torch.arange(num_prefix_chunks, device=device, dtype=torch.int32) - .unsqueeze(1) - .expand(-1, self.batch_size) - * prefix_chunk_len - ) - prefix_chunk_ends = torch.min( - prefix_lens.unsqueeze(0), - prefix_chunk_starts + prefix_chunk_len, - ).to(torch.int32) - - prefix_chunk_seq_lens = ( - (prefix_chunk_ends - prefix_chunk_starts).clamp(min=0).to(torch.int32) - ) - - return prefix_chunk_starts, prefix_chunk_seq_lens - - # Called before each attention module if using chunked kv cache for prefill - # Some of the codes are adapted from https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py - def prepare_chunked_prefix_cache_info(self, device: torch.device): - - from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool - - assert isinstance( - self.token_to_kv_pool, MLATokenToKVPool - ), "Currently chunked prefix cache can only be used by Deepseek models" - - if not any(self.extend_prefix_lens_cpu): - self.num_prefix_chunks = 0 - return - - if self.prefix_chunk_len is not None: - # Chunked kv cache info already prepared by prior modules - return - - self.prefix_chunk_idx = -1 - - # chunk_capacity is the maximum number of tokens in each chunk - chunk_capacity = self.get_max_chunk_capacity() - self.prefix_chunk_len = chunk_capacity // self.batch_size - - self.num_prefix_chunks = ( - max(self.extend_prefix_lens_cpu) + self.prefix_chunk_len - 1 - ) // self.prefix_chunk_len - - # Here we compute chunk lens twice to avoid stream sync, once on gpu and once on cpu. - prefix_chunk_starts_cuda, prefix_chunk_seq_lens_cuda = ( - self.get_prefix_chunk_seq_lens( - self.extend_prefix_lens, - self.num_prefix_chunks, - self.prefix_chunk_len, - ) - ) - _, prefix_chunk_seq_lens_cpu = self.get_prefix_chunk_seq_lens( - torch.tensor(self.extend_prefix_lens_cpu), - self.num_prefix_chunks, - self.prefix_chunk_len, - ) - self.prefix_chunk_starts = prefix_chunk_starts_cuda - self.prefix_chunk_seq_lens = prefix_chunk_seq_lens_cuda - - # Metadata for attention backend - self.prefix_chunk_cu_seq_lens = torch.zeros( - self.num_prefix_chunks, - self.batch_size + 1, - device=device, - dtype=torch.int32, - ) - self.prefix_chunk_cu_seq_lens[:, 1:] = prefix_chunk_seq_lens_cuda.cumsum( - dim=1 - ).to(torch.int32) - self.prefix_chunk_max_seq_lens = prefix_chunk_seq_lens_cpu.max( - dim=1 - ).values.tolist() - - self.prefix_chunk_num_tokens = prefix_chunk_seq_lens_cpu.sum(dim=1).tolist() - assert max(self.prefix_chunk_num_tokens) <= self.get_max_chunk_capacity() - - # Precompute the kv indices for each chunk - self.prepare_chunked_kv_indices(device) - @property def can_run_tbo(self): return self.tbo_split_seq_index is not None - def fetch_mha_one_shot_kv_indices(self): - if self.mha_one_shot_kv_indices is not None: - return self.mha_one_shot_kv_indices - batch_size = self.batch_size - paged_kernel_lens_sum = sum(self.seq_lens_cpu) - kv_indices = torch.empty( - paged_kernel_lens_sum, - dtype=torch.int32, - device=self.req_pool_indices.device, - ) - kv_indptr = torch.zeros( - batch_size + 1, - dtype=torch.int32, - device=self.req_pool_indices.device, - ) - kv_indptr[1:] = torch.cumsum(self.seq_lens, dim=0) - create_flashinfer_kv_indices_triton[(self.batch_size,)]( - self.req_to_token_pool.req_to_token, - self.req_pool_indices, - self.seq_lens, - kv_indptr, - None, - kv_indices, - self.req_to_token_pool.req_to_token.shape[1], - ) - self.mha_one_shot_kv_indices = kv_indices - return kv_indices - def enable_num_token_non_padded(server_args): return get_moe_expert_parallel_world_size() > 1 @@ -1259,40 +1080,3 @@ def compute_position_torch( @torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) def clamp_position(seq_lens): return torch.clamp((seq_lens - 1), min=0).to(torch.int64) - - -@triton.jit -def create_chunked_prefix_cache_kv_indices( - req_to_token_ptr, # (max_batch, max_context_len,) - req_pool_indices_ptr, # (batch_size,) - chunk_start_idx_ptr, # (batch_size,) - chunk_seq_lens_ptr, # (batch_size,) - chunk_cu_seq_lens_ptr, # (batch_size + 1,) - chunk_kv_indices_ptr, # (num_chunk_tokens,) - req_to_token_ptr_stride: tl.constexpr, -): - BLOCK_SIZE: tl.constexpr = 512 - pid = tl.program_id(axis=0) - - # find the req pool idx, this is for batch to token - req_pool_index = tl.load(req_pool_indices_ptr + pid) - chunk_kv_indices_offset = tl.load(chunk_cu_seq_lens_ptr + pid) - - # get the token positions of current chunk - chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32) - chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32) - - num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE) - for i in range(num_loop): - offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE - mask = offset < chunk_seq_len - data = tl.load( - req_to_token_ptr - + req_pool_index * req_to_token_ptr_stride - + chunk_start_pos - + offset, - mask=mask, - ) - tl.store( - chunk_kv_indices_ptr + chunk_kv_indices_offset + offset, data, mask=mask - ) diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 9e99076e9..cc673a9ca 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -7,6 +7,8 @@ from sglang.srt.models.deepseek_common.utils import _is_hip from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import use_intel_amx_backend +MHA_ONE_SHOT_SUPPORTED_BACKENDS = ["fa3", "flashinfer", "flashmla"] + class AttentionBackendRegistry: _handlers = {} @@ -60,7 +62,7 @@ def _get_sum_extend_prefix_lens(forward_batch): def _support_mha_one_shot(attn, forward_batch, backend_name): - attn_supported = backend_name in ["fa3", "flashinfer", "flashmla"] + attn_supported = backend_name in MHA_ONE_SHOT_SUPPORTED_BACKENDS sum_seq_lens = ( sum(forward_batch.seq_lens_cpu) if forward_batch.seq_lens_cpu is not None else 0 ) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py new file mode 100644 index 000000000..e8076e2ad --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/__init__.py @@ -0,0 +1,7 @@ +from .forward_methods import AttnForwardMethod +from .forward_mha import DeepseekMHAForwardMixin + +__all__ = [ + "AttnForwardMethod", + "DeepseekMHAForwardMixin", +] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py index a75424a8d..839ed3a6f 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_methods.py @@ -12,7 +12,7 @@ class AttnForwardMethod(IntEnum): # This method can avoid OOM when prefix lengths are long. MHA_CHUNKED_KV = auto() - # Use multi-head attention, execute the MHA for prefix and extended kv in one shot + # Use multi-head attention, execute the MHA for prefix and extended kv in a single kernel # when the sequence lengths are below the threshold. MHA_ONE_SHOT = auto() diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py new file mode 100644 index 000000000..60b6031e9 --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -0,0 +1,493 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged +from sglang.srt.layers.attention.tbo_backend import TboAttnBackend +from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton +from sglang.srt.layers.communicator import get_attn_tp_context +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.deepseek_common.utils import ( + _is_cuda, + _is_hip, + _is_npu, + _use_aiter_gfx95, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import BumpAllocator + +if TYPE_CHECKING: + from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA + +if _is_cuda: + from sgl_kernel import concat_mla_k, merge_state_v2 + +if _use_aiter_gfx95: + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant + + from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant + +# Configs for DeepSeek-V3: +# num_local_heads = 128 +# qk_nope_head_dim = 128 +# qk_rope_head_dim = 64 +# qk_head_dim = qk_nope_head_dim + qk_rope_head_dim = 192 +# v_head_dim = 128 + +# Configs for kv chunking strategy: +# sum_prefix_length: +# Total number of tokens to be fetched from kv cache for current batch. +# e.g: For batch with 2 sequences, seq_lens_kv = [1024, 2048], seq_lens_q = [512, 1024], then sum_prefix_length = (1024 - 512) + (2048 - 1024) = 1536 +# sum_extended_length: +# Total number of tokens in the extended part of the current batch. (=sum(seq_lens_q)) +# chunked_prefix_cache_threshold: +# The minimum sum_prefix_length to enable mha with kv chunking, 8192 by default (can be changed with SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD) +# For batches with smaller sum_prefix_length > 0, MLA kernel with absorption will be used instead. +# max_kv_chunk_capacity: +# The maximum number of tokens in each kv chunk, 128 * 1024 by default (can be get with forward_batch.get_max_chunk_capacity()) + +# The forward methods for MHA in DeepSeek models: +# +# 1. forward_normal: AttnForwardMethod.MHA +# use multi-head attention with empty kv cache (the first batch of chunked prefill, prefix lens = 0) +# q: [sum_extended_length, num_local_heads, qk_head_dim] +# k: [sum_extended_length, num_local_heads, qk_head_dim] +# v: [sum_extended_length, num_local_heads, v_head_dim] +# +# 2. forward_normal_one_shot: AttnForwardMethod.MHA_ONE_SHOT +# use multi-head attention with short kv prefix length (chunked_prefix_cache_threshold <= sum_prefix_lens <= max_kv_chunk_capacity) +# the kv latent vectors are fetched from memory pool, with combined kv_indices of prefix part and extended part +# q: [batch_size, num_local_heads, qk_head_dim] +# k: [sum_extended_length + sum_prefix_length, num_local_heads, qk_head_dim] +# v: [sum_extended_length + sum_prefix_length, num_local_heads, v_head_dim] +# +# 3. forward_normal_chunked_kv: AttnForwardMethod.MHA_CHUNKED_KV +# multiple phases of multi-head attention with chunked kv cache (sum_prefix_length > max_kv_chunk_capacity) +# For the first phase, it will execute normal forward method, and returns output o_1 and lse_1, +# q_1: [sum_extended_length, num_local_heads, qk_head_dim], +# k_1: [sum_extended_length, num_local_heads, qk_head_dim], +# v_1: [sum_extended_length, num_local_heads, qk_head_dim], +# acc_o_1, acc_lse_1 = o_1, lse_1 +# For i in range(2, n), (n-1 is the number of prefix chunks), kv latent vectors are fetched from memory pool with prefix kv indices +# q_i: [sum_extended_length, num_local_heads, qk_head_dim], +# k_i: [chunk_size, num_local_heads, qk_head_dim], +# v_i: [chunk_size, num_local_heads, v_head_dim], +# acc_o_i, acc_lse_i = merge_state(acc_o_{i-1}, acc_lse_{i-1}, o_i, lse_i) +# The final output is the accumulated output acc_o_n + + +class DeepseekMHAForwardMixin: + + def init_mha_forward(self: DeepseekV2AttentionMLA): + self.disable_chunked_prefix_cache = ( + get_global_server_args().disable_chunked_prefix_cache + ) + + # TODO: Design a finer way to determine the threshold + self.chunked_prefix_cache_threshold = ( + envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.get() + ) + + def forward_normal_prepare( + self: DeepseekV2AttentionMLA, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + if self.q_lora_rank is not None: + q, latent_cache = ( + get_attn_tp_context() + .fetch_qkv_latent() + .split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + dim=-1, + ) + ) + + # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk + + if self.use_nsa: + # NSA requires unquantized q_lora for the indexer. When q_b_proj is FP8 + # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request + # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale) + # tuple. + if ( + _use_aiter_gfx95 + and self.q_b_proj.weight.dtype == torch.float8_e4m3fn + ): + q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant( + q, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + None, + None, + None, + group_size=128, + dtype_quant=torch.float8_e4m3fn, + res1=None, + output_unquantized_inp1=True, + ) + q = self.q_b_proj(q_quanted)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + else: + q_lora = self.q_a_layernorm(q) + q = self.q_b_proj(q_lora)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + _ = self.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=self.layer_id, + return_indices=False, + ) + elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: + # MXFP4: fused RMSNorm + quant + q, _, _, _ = fused_rms_mxfp4_quant( + q, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + None, + None, + None, + ) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn: + + q, _, _, _ = fused_rms_fp8_group_quant( + q, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + None, + None, + None, + group_size=128, + dtype_quant=torch.float8_e4m3fn, + res1=None, + output_unquantized_inp1=False, + ) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + + _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + + if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: + + kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant( + kv_a, + self.kv_a_layernorm.weight, + self.kv_a_layernorm.variance_epsilon, + None, + None, + None, + group_size=128, + dtype_quant=torch.float8_e4m3fn, + res1=None, + output_unquantized_inp1=True, # return unqaunt kv_a + ) + + else: + kv_a = self.kv_a_layernorm(kv_a) + + k_pe = latent_cache[:, :, self.kv_lora_rank :] + if self.rotary_emb is not None: + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q[..., self.qk_nope_head_dim :] = q_pe + + self._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch) + if ( + forward_batch.mha_one_shot + and sum(forward_batch.extend_prefix_lens_cpu) != 0 + ): + if self.use_nsa and self.kv_cache_dtype == "fp8_e4m3": + # FP8 path: dequantize NSA-specific FP8 format to BF16 + kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch) + else: + # BF16/FP16 path: directly fetch from cache + kv_a, k_pe = self._get_mla_kv_buffer( + forward_batch.fetch_mha_one_shot_kv_indices(), + q.dtype, + forward_batch, + ) + if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: + kv = self.kv_b_proj( + kv_a_quanted, + )[0] + else: + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope = kv[..., : self.qk_nope_head_dim] + v = kv[..., self.qk_nope_head_dim :] + + k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch) + return q, k, v, forward_batch + + def forward_normal_core( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + def forward_normal_chunked_kv_prepare( + self: DeepseekV2AttentionMLA, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + # In normal mha, the k and v tensors will become overly large when the prefix length is long. + # To avoid this, we split the kv cache into chunks and process them one after another. + # Since mha is compute friendly, the for loop induced here will not introduce significant overhead. + # The top comments in https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py + # will be helpful for understanding the purpose of this function. + + # First do normal mha forward to get output for extended part + return self.forward_normal_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + + def forward_normal_chunked_kv_core( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + has_extend_prefix = forward_batch.extend_prefix_lens_cpu is not None and any( + forward_batch.extend_prefix_lens_cpu + ) + # Only initialize the info once + if has_extend_prefix and forward_batch.num_prefix_chunks is None: + forward_batch.prepare_chunked_prefix_cache_info(q.device) + if hasattr(forward_batch.attn_backend, "init_mha_chunk_metadata"): + forward_batch.attn_backend.init_mha_chunk_metadata(forward_batch) + + forward_batch.mha_return_lse = has_extend_prefix + # Do mha for extended part without prefix + forward_batch.set_attn_attend_prefix_cache(False) + attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + + # Do mha attention with chunked prefix cache if there are any sequence with prefix + if has_extend_prefix: + attn_output, lse = attn_output + forward_batch.set_attn_attend_prefix_cache(True) + attn_output = self._chunked_prefix_attn_mha( + q=q, + accum_output=attn_output, + accum_lse=lse, + forward_batch=forward_batch, + ) + + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + def forward_normal_one_shot_prepare( + self: DeepseekV2AttentionMLA, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ): + forward_batch.mha_one_shot = True + return self.forward_normal_prepare( + positions, hidden_states, forward_batch, zero_allocator + ) + + def forward_normal_one_shot_core( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + has_extend_prefix = any(forward_batch.extend_prefix_lens_cpu) + # Only initialize the info once + if has_extend_prefix and forward_batch.num_prefix_chunks is None: + forward_batch.num_prefix_chunks = 0 + if hasattr(forward_batch.attn_backend, "init_mha_chunk_metadata"): + forward_batch.attn_backend.init_mha_chunk_metadata(forward_batch) + forward_batch.mha_return_lse = False + # Do mha for extended part without prefix + forward_batch.set_attn_attend_prefix_cache(False) + return self.forward_normal_core(q, k, v, forward_batch) + + def _chunked_prefix_attn_mha( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + accum_output: torch.Tensor, + accum_lse: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + + assert forward_batch.num_prefix_chunks is not None + for i in range(forward_batch.num_prefix_chunks): + forward_batch.set_prefix_chunk_idx(i) + + kv_indices = forward_batch.prefix_chunk_kv_indices[i] + # Fetch latent cache from memory pool with precomputed chunked kv indices + kv_a_normed, k_pe = self._get_mla_kv_buffer( + kv_indices, q.dtype, forward_batch + ) + kv = self.kv_b_proj(kv_a_normed)[0] + kv = kv.view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + v = kv[..., self.qk_nope_head_dim :] + k_nope = kv[..., : self.qk_nope_head_dim] + + k = torch.empty( + ( + k_nope.shape[0], + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + ), + dtype=v.dtype, + device=v.device, + ) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + tmp_output = torch.empty_like(accum_output) + tmp_lse = torch.empty_like(accum_lse) + merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse) + accum_output, accum_lse = tmp_output, tmp_lse + del kv, k, v, output, lse, tmp_output, tmp_lse + + return accum_output + + def _set_mla_kv_buffer( + self: DeepseekV2AttentionMLA, + latent_cache: torch.Tensor, + kv_a: torch.Tensor, + k_pe: torch.Tensor, + forward_batch: ForwardBatch, + ): + if _is_cuda or _use_aiter_gfx95: + # Save latent cache + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe + ) + elif _is_npu: + # To reduce a time-costing split operation + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe + ) + else: + latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) + latent_cache[:, :, self.kv_lora_rank :] = k_pe + + # Save latent cache + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mha, forward_batch.out_cache_loc, latent_cache, None + ) + + def _get_mla_kv_buffer( + self: DeepseekV2AttentionMLA, + kv_indices: torch.Tensor, + dst_dtype: torch.dtype, + forward_batch: ForwardBatch, + ): + if _is_cuda or _use_aiter_gfx95: + kv_a, k_pe = forward_batch.token_to_kv_pool.get_mla_kv_buffer( + self.attn_mha, kv_indices, dst_dtype + ) + kv_a = kv_a.squeeze(1) + else: + latent_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer( + self.attn_mha.layer_id + ) + latent_cache = latent_cache_buf[kv_indices].contiguous().to(dst_dtype) + + kv_a, k_pe = latent_cache.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv_a = kv_a.squeeze(1).contiguous() + return kv_a, k_pe + + def _get_mla_kv_buffer_from_fp8_for_nsa( + self: DeepseekV2AttentionMLA, + forward_batch: ForwardBatch, + ): + """ + Dequantize FP8 KV cache to BF16 for MLA attention (NSA-specific format). + + Returns: (kv_a, k_pe) both in BF16 + """ + backend = forward_batch.attn_backend + if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend + backend = backend.primary + kv_indices = backend.forward_metadata.page_table_1_flattened + assert ( + kv_indices is not None + ), "page_table_1_flattened should have been generated for FP8 MHA path" + + kv_cache_fp8 = forward_batch.token_to_kv_pool.get_key_buffer( + self.attn_mha.layer_id + ) + + kv_latent_bf16 = dequantize_k_cache_paged(kv_cache_fp8, kv_indices) + + kv_a = kv_latent_bf16[:, :, : self.kv_lora_rank].squeeze(1).contiguous() + k_pe = kv_latent_bf16[:, :, self.kv_lora_rank :] + + return kv_a, k_pe + + def _concat_and_cast_mha_k( + self: DeepseekV2AttentionMLA, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + forward_batch: ForwardBatch, + ): + # Temporary for DeepSeek V3/R1 only, but can generalize if needed + k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim) + if ( + _is_cuda + and (self.num_local_heads == 128) + and (self.qk_nope_head_dim == 128) + and (self.qk_rope_head_dim == 64) + ): + k = k_nope.new_empty(*k_shape) + concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe) + elif _is_cuda: + # fa3 mha support fp8 inputs + if ( + self.current_attention_backend == "fa3" + and self.kv_cache_dtype != "auto" + ): + attn_dtype = forward_batch.token_to_kv_pool.dtype + else: + attn_dtype = k_nope.dtype + k = k_nope.new_empty(*k_shape, dtype=attn_dtype) + concat_and_cast_mha_k_triton(k, k_nope, k_pe) + elif _is_hip and self.current_attention_backend == "aiter": + k = k_nope.new_empty(*k_shape) + concat_and_cast_mha_k_triton(k, k_nope, k_pe) + else: + k = k_nope.new_empty(*k_shape) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + return k diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3aa447f78..2fb748077 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -56,7 +56,6 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.amx_utils import PackWeightMethod -from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer from sglang.srt.layers.attention.nsa.utils import ( can_cp_split, @@ -67,8 +66,6 @@ from sglang.srt.layers.attention.nsa.utils import ( nsa_use_prefill_cp, prepare_input_dp_with_cp_dsa, ) -from sglang.srt.layers.attention.tbo_backend import TboAttnBackend -from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, @@ -142,8 +139,9 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_common.attention_backend_handler import ( AttentionBackendRegistry, ) -from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import ( +from sglang.srt.models.deepseek_common.attention_forward_methods import ( AttnForwardMethod, + DeepseekMHAForwardMixin, ) from sglang.srt.models.deepseek_common.utils import ( _device_sm, @@ -195,14 +193,7 @@ if _use_aiter_gfx95: ) if _is_cuda: - from sgl_kernel import ( - awq_dequantize, - bmm_fp8, - concat_mla_k, - dsv3_fused_a_gemm, - dsv3_router_gemm, - merge_state_v2, - ) + from sgl_kernel import awq_dequantize, bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: @@ -255,12 +246,6 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [ ] -def add_forward_absorb_core_attention_backend(backend_name): - if backend_name not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS: - FORWARD_ABSORB_CORE_ATTENTION_BACKENDS.append(backend_name) - logger.info(f"Added {backend_name} to FORWARD_ABSORB_CORE_ATTENTION_BACKENDS.") - - class DeepseekV2MLP(nn.Module): def __init__( self, @@ -1074,7 +1059,7 @@ def _get_llama_4_scaling( return scaling[..., None, None] -class DeepseekV2AttentionMLA(nn.Module): +class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): def __init__( self, @@ -1264,9 +1249,6 @@ class DeepseekV2AttentionMLA(nn.Module): self.flashinfer_mla_disable_ragged = ( get_global_server_args().flashinfer_mla_disable_ragged ) - self.disable_chunked_prefix_cache = ( - get_global_server_args().disable_chunked_prefix_cache - ) self.current_attention_backend = ( None # Attention backend used by current forward batch @@ -1275,11 +1257,6 @@ class DeepseekV2AttentionMLA(nn.Module): "SGLANG_ROCM_FUSED_DECODE_MLA", "false" ) - # TODO: Design a finer way to determine the threshold - self.chunked_prefix_cache_threshold = ( - envs.SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD.get() - ) - # If we have self.fused_qkv_a_proj_with_mqa and we're running on CPU, we will choose the torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight kernel # which requires self.w_kc and self.w_vc to be packed. # If not, we will use torch.bmm and weight shouldn't be packed in this case @@ -1334,6 +1311,8 @@ class DeepseekV2AttentionMLA(nn.Module): self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size ) + self.init_mha_forward() + def dispatch_attn_forward_method( self, forward_batch: ForwardBatch ) -> AttnForwardMethod: @@ -1503,161 +1482,6 @@ class DeepseekV2AttentionMLA(nn.Module): qkv_latent = self.fused_qkv_a_proj_with_mqa(hidden_states)[0] return qkv_latent - def forward_normal_prepare( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - ): - if self.q_lora_rank is not None: - q, latent_cache = ( - get_attn_tp_context() - .fetch_qkv_latent() - .split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], - dim=-1, - ) - ) - - # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk - - if self.use_nsa: - # NSA requires unquantized q_lora for the indexer. When q_b_proj is FP8 - # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request - # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale) - # tuple. - if ( - _use_aiter_gfx95 - and self.q_b_proj.weight.dtype == torch.float8_e4m3fn - ): - q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant( - q, - self.q_a_layernorm.weight, - self.q_a_layernorm.variance_epsilon, - None, - None, - None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=True, - ) - q = self.q_b_proj(q_quanted)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) - else: - q_lora = self.q_a_layernorm(q) - q = self.q_b_proj(q_lora)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) - _ = self.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=self.layer_id, - return_indices=False, - ) - elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: - # MXFP4: fused RMSNorm + quant - q, _, _, _ = fused_rms_mxfp4_quant( - q, - self.q_a_layernorm.weight, - self.q_a_layernorm.variance_epsilon, - None, - None, - None, - ) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) - elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn: - - q, _, _, _ = fused_rms_fp8_group_quant( - q, - self.q_a_layernorm.weight, - self.q_a_layernorm.variance_epsilon, - None, - None, - None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=False, - ) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) - else: - q = self.q_a_layernorm(q) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) - - else: - q = self.q_proj(hidden_states)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) - latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] - - _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - latent_cache = latent_cache.unsqueeze(1) - - if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: - - kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant( - kv_a, - self.kv_a_layernorm.weight, - self.kv_a_layernorm.variance_epsilon, - None, - None, - None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=True, # return unqaunt kv_a - ) - - else: - kv_a = self.kv_a_layernorm(kv_a) - - # kv_a = self.kv_a_layernorm(kv_a) - - k_pe = latent_cache[:, :, self.kv_lora_rank :] - if self.rotary_emb is not None: - q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) - q[..., self.qk_nope_head_dim :] = q_pe - - self._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch) - if ( - forward_batch.mha_one_shot - and sum(forward_batch.extend_prefix_lens_cpu) != 0 - ): - if self.use_nsa and self.kv_cache_dtype == "fp8_e4m3": - # FP8 path: dequantize NSA-specific FP8 format to BF16 - kv_a, k_pe = self._get_mla_kv_buffer_from_fp8(forward_batch) - else: - # BF16/FP16 path: directly fetch from cache - kv_a, k_pe = self._get_mla_kv_buffer( - forward_batch.fetch_mha_one_shot_kv_indices(), - q.dtype, - forward_batch, - ) - if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: - kv = self.kv_b_proj( - kv_a_quanted, - )[0] - else: - kv = self.kv_b_proj(kv_a)[0] - kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) - k_nope = kv[..., : self.qk_nope_head_dim] - v = kv[..., self.qk_nope_head_dim :] - - k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch) - return q, k, v, forward_batch - - def forward_normal_core(self, q, k, v, forward_batch): - attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) - attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) - output, _ = self.o_proj(attn_output) - return output - def _fuse_rope_for_trtllm_mla(self, forward_batch: ForwardBatch) -> bool: """ Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path. @@ -2377,231 +2201,6 @@ class DeepseekV2AttentionMLA(nn.Module): return output - def _chunked_prefix_attn_mha( - self, - q: torch.Tensor, - accum_output: torch.Tensor, - accum_lse: torch.Tensor, - forward_batch: ForwardBatch, - ) -> torch.Tensor: - - assert forward_batch.num_prefix_chunks is not None - for i in range(forward_batch.num_prefix_chunks): - forward_batch.set_prefix_chunk_idx(i) - - kv_indices = forward_batch.prefix_chunk_kv_indices[i] - # Fetch latent cache from memory pool with precomputed chunked kv indices - kv_a_normed, k_pe = self._get_mla_kv_buffer( - kv_indices, q.dtype, forward_batch - ) - kv = self.kv_b_proj(kv_a_normed)[0] - kv = kv.view( - -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim - ) - v = kv[..., self.qk_nope_head_dim :] - k_nope = kv[..., : self.qk_nope_head_dim] - - k = torch.empty( - ( - k_nope.shape[0], - self.num_local_heads, - self.qk_nope_head_dim + self.qk_rope_head_dim, - ), - dtype=v.dtype, - device=v.device, - ) - k[..., : self.qk_nope_head_dim] = k_nope - k[..., self.qk_nope_head_dim :] = k_pe - - output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) - tmp_output = torch.empty_like(accum_output) - tmp_lse = torch.empty_like(accum_lse) - merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse) - accum_output, accum_lse = tmp_output, tmp_lse - del kv, k, v, output, lse, tmp_output, tmp_lse - - return accum_output - - def forward_normal_chunked_kv_prepare( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - ): - # In normal mha, the k and v tensors will become overly large when the prefix length is long. - # To avoid this, we split the kv cache into chunks and process them one after another. - # Since mha is compute friendly, the for loop induced here will not introduce significant overhead. - # The top comments in https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py - # will be helpful for understanding the purpose of this function. - - # First do normal mha forward to get output for extended part - return self.forward_normal_prepare( - positions, hidden_states, forward_batch, zero_allocator - ) - - def forward_normal_chunked_kv_core(self, q, k, v, forward_batch): - has_extend_prefix = forward_batch.extend_prefix_lens_cpu is not None and any( - forward_batch.extend_prefix_lens_cpu - ) - # Only initialize the info once - if has_extend_prefix and forward_batch.num_prefix_chunks is None: - forward_batch.prepare_chunked_prefix_cache_info(q.device) - if hasattr(forward_batch.attn_backend, "init_mha_chunk_metadata"): - forward_batch.attn_backend.init_mha_chunk_metadata(forward_batch) - - forward_batch.mha_return_lse = has_extend_prefix - # Do mha for extended part without prefix - forward_batch.set_attn_attend_prefix_cache(False) - attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) - - # Do mha attention with chunked prefix cache if there are any sequence with prefix - if has_extend_prefix: - attn_output, lse = attn_output - forward_batch.set_attn_attend_prefix_cache(True) - attn_output = self._chunked_prefix_attn_mha( - q=q, - accum_output=attn_output, - accum_lse=lse, - forward_batch=forward_batch, - ) - - attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) - output, _ = self.o_proj(attn_output) - return output - - def forward_normal_one_shot_prepare( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - ): - forward_batch.mha_one_shot = True - return self.forward_normal_prepare( - positions, hidden_states, forward_batch, zero_allocator - ) - - def forward_normal_one_shot_core(self, q, k, v, forward_batch): - has_extend_prefix = any(forward_batch.extend_prefix_lens_cpu) - # Only initialize the info once - if has_extend_prefix and forward_batch.num_prefix_chunks is None: - forward_batch.num_prefix_chunks = 0 - if hasattr(forward_batch.attn_backend, "init_mha_chunk_metadata"): - forward_batch.attn_backend.init_mha_chunk_metadata(forward_batch) - forward_batch.mha_return_lse = False - # Do mha for extended part without prefix - forward_batch.set_attn_attend_prefix_cache(False) - return self.forward_normal_core(q, k, v, forward_batch) - - def _set_mla_kv_buffer( - self, - latent_cache: torch.Tensor, - kv_a: torch.Tensor, - k_pe: torch.Tensor, - forward_batch: ForwardBatch, - ): - if _is_cuda or _use_aiter_gfx95: - # Save latent cache - forward_batch.token_to_kv_pool.set_mla_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe - ) - elif _is_npu: - # To reduce a time-costing split operation - forward_batch.token_to_kv_pool.set_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe - ) - else: - latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) - latent_cache[:, :, self.kv_lora_rank :] = k_pe - - # Save latent cache - forward_batch.token_to_kv_pool.set_kv_buffer( - self.attn_mha, forward_batch.out_cache_loc, latent_cache, None - ) - - def _get_mla_kv_buffer( - self, - kv_indices: torch.Tensor, - dst_dtype: torch.dtype, - forward_batch: ForwardBatch, - ): - if _is_cuda or _use_aiter_gfx95: - kv_a, k_pe = forward_batch.token_to_kv_pool.get_mla_kv_buffer( - self.attn_mha, kv_indices, dst_dtype - ) - kv_a = kv_a.squeeze(1) - else: - latent_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer( - self.attn_mha.layer_id - ) - latent_cache = latent_cache_buf[kv_indices].contiguous().to(dst_dtype) - - kv_a, k_pe = latent_cache.split( - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - kv_a = kv_a.squeeze(1).contiguous() - return kv_a, k_pe - - def _get_mla_kv_buffer_from_fp8( - self, - forward_batch: ForwardBatch, - ): - """ - Dequantize FP8 KV cache to BF16 for MLA attention (NSA-specific format). - - Returns: (kv_a, k_pe) both in BF16 - """ - backend = forward_batch.attn_backend - if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend - backend = backend.primary - kv_indices = backend.forward_metadata.page_table_1_flattened - assert ( - kv_indices is not None - ), "page_table_1_flattened should have been generated for FP8 MHA path" - - kv_cache_fp8 = forward_batch.token_to_kv_pool.get_key_buffer( - self.attn_mha.layer_id - ) - - kv_latent_bf16 = dequantize_k_cache_paged(kv_cache_fp8, kv_indices) - - kv_a = kv_latent_bf16[:, :, : self.kv_lora_rank].squeeze(1).contiguous() - k_pe = kv_latent_bf16[:, :, self.kv_lora_rank :] - - return kv_a, k_pe - - def _concat_and_cast_mha_k(self, k_nope, k_pe, forward_batch): - # Temporary for DeepSeek V3/R1 only, but can generalize if needed - k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim) - if ( - _is_cuda - and (self.num_local_heads == 128) - and (self.qk_nope_head_dim == 128) - and (self.qk_rope_head_dim == 64) - ): - k = k_nope.new_empty(*k_shape) - concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe) - elif _is_cuda: - # fa3 mha support fp8 inputs - if ( - self.current_attention_backend == "fa3" - and self.kv_cache_dtype != "auto" - ): - attn_dtype = forward_batch.token_to_kv_pool.dtype - else: - attn_dtype = k_nope.dtype - k = k_nope.new_empty(*k_shape, dtype=attn_dtype) - concat_and_cast_mha_k_triton(k, k_nope, k_pe) - elif _is_hip and self.current_attention_backend == "aiter": - k = k_nope.new_empty(*k_shape) - concat_and_cast_mha_k_triton(k, k_nope, k_pe) - else: - k = k_nope.new_empty(*k_shape) - k[..., : self.qk_nope_head_dim] = k_nope - k[..., self.qk_nope_head_dim :] = k_pe - return k - @staticmethod def _get_q_b_proj_quant_config(quant_config): if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():