[2/n] deepseek_v2.py Refactor: Migrate MHA forward method in deepseek_v2.py (#16817)
This commit is contained in:
Baizhou Zhang
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
from .forward_methods import AttnForwardMethod
|
||||
from .forward_mha import DeepseekMHAForwardMixin
|
||||
|
||||
__all__ = [
|
||||
"AttnForwardMethod",
|
||||
"DeepseekMHAForwardMixin",
|
||||
]
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user