Opt MHA chunked prefix: merge prefix and extend kv cache to run mha once (#10953)
This commit is contained in:
Yongfei Xu
@@ -57,6 +57,7 @@ from sglang.srt.layers.attention.npu_ops.mla_preprocess import (
|
||||
is_mla_preprocess_enabled,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
||||
from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton
|
||||
from sglang.srt.layers.communicator import (
|
||||
LayerCommunicator,
|
||||
LayerScatterModes,
|
||||
@@ -241,6 +242,10 @@ 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
|
||||
# when the sequence lengths are below the threshold.
|
||||
MHA_ONE_SHOT = auto()
|
||||
|
||||
# Use MLA but with fused RoPE
|
||||
MLA_FUSED_ROPE = auto()
|
||||
|
||||
@@ -306,6 +311,14 @@ def _is_extend_without_speculative(forward_batch):
|
||||
)
|
||||
|
||||
|
||||
def _support_mha_one_shot(attn: DeepseekV2AttentionMLA, forward_batch, backend_name):
|
||||
attn_supported = backend_name in ["fa3", "flashinfer", "flashmla"]
|
||||
sum_seq_lens = (
|
||||
sum(forward_batch.seq_lens_cpu) if forward_batch.seq_lens_cpu is not None else 0
|
||||
)
|
||||
return attn_supported and sum_seq_lens <= forward_batch.get_max_chunk_capacity()
|
||||
|
||||
|
||||
def _handle_attention_backend(
|
||||
attn: DeepseekV2AttentionMLA, forward_batch, backend_name
|
||||
):
|
||||
@@ -325,6 +338,8 @@ def _handle_attention_backend(
|
||||
or sum_extend_prefix_lens == 0
|
||||
)
|
||||
):
|
||||
if _support_mha_one_shot(attn, forward_batch, backend_name):
|
||||
return AttnForwardMethod.MHA_ONE_SHOT
|
||||
return AttnForwardMethod.MHA_CHUNKED_KV
|
||||
else:
|
||||
return _dispatch_mla_subtype(attn, forward_batch)
|
||||
@@ -1062,6 +1077,7 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
self.scaling = self.qk_head_dim**-0.5
|
||||
self.rope_theta = rope_theta
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.kv_cache_dtype = get_global_server_args().kv_cache_dtype
|
||||
|
||||
# NOTE modification to rope_scaling must be done early enough, b/c e.g. Indexer needs it
|
||||
if rope_scaling:
|
||||
@@ -1359,6 +1375,10 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
inner_state = self.forward_normal_chunked_kv_prepare(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
|
||||
inner_state = self.forward_normal_one_shot_prepare(
|
||||
positions, hidden_states, forward_batch, zero_allocator
|
||||
)
|
||||
elif attn_forward_method == AttnForwardMethod.MLA:
|
||||
if not self.is_mla_preprocess_enabled:
|
||||
inner_state = self.forward_absorb_prepare(
|
||||
@@ -1410,6 +1430,8 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
return self.forward_normal_core(*inner_state)
|
||||
elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
|
||||
return self.forward_normal_chunked_kv_core(*inner_state)
|
||||
elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT:
|
||||
return self.forward_normal_one_shot_core(*inner_state)
|
||||
elif attn_forward_method == AttnForwardMethod.MLA:
|
||||
return self.forward_absorb_core(*inner_state)
|
||||
elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
|
||||
@@ -1444,41 +1466,24 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
|
||||
latent_cache = latent_cache.unsqueeze(1)
|
||||
kv_a = self.kv_a_layernorm(kv_a)
|
||||
k_pe = latent_cache[:, :, self.kv_lora_rank :]
|
||||
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
|
||||
):
|
||||
kv_a, k_pe = self._get_mla_kv_buffer(
|
||||
forward_batch.fetch_mha_one_shot_kv_indices(), q.dtype, forward_batch
|
||||
)
|
||||
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_pe = latent_cache[:, :, self.kv_lora_rank :]
|
||||
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
|
||||
q[..., self.qk_nope_head_dim :] = q_pe
|
||||
k = torch.empty_like(q)
|
||||
|
||||
# Temporary for DeepSeek V3/R1 only, but can generalize if needed
|
||||
if (
|
||||
_is_cuda
|
||||
and (self.num_local_heads == 128)
|
||||
and (self.qk_nope_head_dim == 128)
|
||||
and (self.qk_rope_head_dim == 64)
|
||||
):
|
||||
concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe)
|
||||
else:
|
||||
k[..., : self.qk_nope_head_dim] = k_nope
|
||||
k[..., self.qk_nope_head_dim :] = k_pe
|
||||
|
||||
if not _is_npu:
|
||||
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
|
||||
)
|
||||
else:
|
||||
# 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
|
||||
)
|
||||
|
||||
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):
|
||||
@@ -2288,20 +2293,11 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
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
|
||||
latent_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer(
|
||||
self.attn_mha.layer_id
|
||||
kv_a_normed, k_pe = self._get_mla_kv_buffer(
|
||||
kv_indices, q.dtype, forward_batch
|
||||
)
|
||||
latent_cache = (
|
||||
latent_cache_buf[forward_batch.prefix_chunk_kv_indices[i]]
|
||||
.contiguous()
|
||||
.to(q.dtype)
|
||||
)
|
||||
|
||||
kv_a_normed, k_pe = latent_cache.split(
|
||||
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
kv_a_normed = kv_a_normed.squeeze(1).contiguous()
|
||||
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
|
||||
@@ -2376,6 +2372,107 @@ class DeepseekV2AttentionMLA(nn.Module):
|
||||
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:
|
||||
# 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:
|
||||
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 _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)
|
||||
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 get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN"):
|
||||
|
||||
Reference in New Issue
Block a user