【FIX】Change the name of sparse attn from moba to plas (#4006) (#4076)

* 【FIX】Change the name of sparse attn from moba to plas (#4006)

* 更新文档

* 【docs】 update readme (#4000)

* 更新文档

* update readme

* update docs

* 【FIX】Change the name of sparse attn from moba to plas (#3845)

* 更新文档

* 更新文档

* 更新文档

* 更新文档

* 修改moba为plas

* code style

* update ci

* code style

* update ci

* code style

---------

Co-authored-by: Jiang-Jia-Jun <163579578+Jiang-Jia-Jun@users.noreply.github.com>

* fix max_num_seqs

* fix test load attn

---------

Co-authored-by: Jiang-Jia-Jun <163579578+Jiang-Jia-Jun@users.noreply.github.com>
This commit is contained in:
yangjianfengo1
2025-09-23 10:26:40 +08:00
committed by GitHub
parent 2c34a557f4
commit 4325b737e7
14 changed files with 152 additions and 152 deletions

View File

@@ -196,7 +196,7 @@ We selected a subset (longbook_sum_eng) from InfiniteBench as the performance ev
## Usage
```
export FD_ATTENTION_BACKEND="MOBA_ATTN"
export FD_ATTENTION_BACKEND="PLAS_ATTN"
python -m fastdeploy.entrypoints.openai.api_server
--model baidu/ERNIE-4.5-300B-A47B-Paddle \
@@ -207,13 +207,13 @@ python -m fastdeploy.entrypoints.openai.api_server
--max-num-batched-tokens 8192 \
--max-model-len 131072 \
--max-num-seqs 32 \
--moba-attention-config '{"moba_encoder_top_k_left": 50, "moba_encoder_top_k_right": 60, "moba_decoder_top_k_left": 100, "moba_decoder_top_k_right": 120}'
--plas-attention-config '{"plas_encoder_top_k_left": 50, "plas_encoder_top_k_right": 60, "plas_decoder_top_k_left": 100, "plas_decoder_top_k_right": 120}'
```
**Note**: If sparse attention is enabled, the system will automatically load the MLP weights from `moba_mlp_weight.safetensors` in the weight directory. If the MLP weight file is not found, mean pooling will be applied to the key representations.
**Note**: If sparse attention is enabled, the system will automatically load the MLP weights from `plas_attention_mlp_weight.safetensors` in the weight directory. If the MLP weight file is not found, mean pooling will be applied to the key representations.
**Parameter Description:**
* Setting `FD_ATTENTION_BACKEND="MOBA_ATTN"` enables MOBA sparse attention.
* `moba_encoder_top_k_left=50, moba_encoder_top_k_right=60` indicates that the range of top-k is between 50 and 60 when the encoder is sparse.
* `moba_decoder_top_k_left=100, moba_decoder_top_k_right=120` indicates that the range of top-k is between 100 and 120 when the decoder is sparse.
* Setting `FD_ATTENTION_BACKEND="PLAS_ATTN"` enables PLAS sparse attention.
* `plas_encoder_top_k_left=50, plas_encoder_top_k_right=60` indicates that the range of top-k is between 50 and 60 when the encoder is sparse.
* `plas_decoder_top_k_left=100, plas_decoder_top_k_right=120` indicates that the range of top-k is between 100 and 120 when the decoder is sparse.

View File

@@ -200,7 +200,7 @@
## 使用方式
```
export FD_ATTENTION_BACKEND="MOBA_ATTN"
export FD_ATTENTION_BACKEND="PLAS_ATTN"
python -m fastdeploy.entrypoints.openai.api_server
--model baidu/ERNIE-4.5-300B-A47B-Paddle \
@@ -211,13 +211,13 @@ python -m fastdeploy.entrypoints.openai.api_server
--max-num-batched-tokens 8192 \
--max-model-len 131072 \
--max-num-seqs 32 \
--moba-attention-config '{"moba_encoder_top_k_left": 50, "moba_encoder_top_k_right": 60, "moba_decoder_top_k_left": 100, "moba_decoder_top_k_right": 120}'
--plas-attention-config '{"plas_encoder_top_k_left": 50, "plas_encoder_top_k_right": 60, "plas_decoder_top_k_left": 100, "plas_decoder_top_k_right": 120}'
```
**Note**: 如果启用了稀疏注意力机制,系统将自动从权重目录中的`moba_mlp_weight.safetensors`文件加载 MLP 权重。如果未找到 MLP 权重文件,则将对关键表示应用均值池化
**Note**: 如果启用了稀疏注意力机制,系统将自动从权重目录中的`plas_attention_mlp_weight.safetensors`文件加载 MLP 权重。如果未找到 MLP 权重文件,则将对关键表示应用均值池化
**Parameter Description:**
* `FD_ATTENTION_BACKEND="MOBA_ATTN"` 启用 MOBA sparse attention.
* `moba_encoder_top_k_left=50, moba_encoder_top_k_right=60` 表示当encoder时top-k的范围在50到60之间。
* `moba_decoder_top_k_left=100, moba_decoder_top_k_right=120` 表示当decoder时top-k的范围在100到120之间。
* `FD_ATTENTION_BACKEND="PLAS_ATTN"` 启用 PLAS sparse attention.
* `plas_encoder_top_k_left=50, plas_encoder_top_k_right=60` 表示当encoder时top-k的范围在50到60之间。
* `plas_decoder_top_k_left=100, plas_decoder_top_k_right=120` 表示当decoder时top-k的范围在100到120之间。

View File

@@ -945,63 +945,63 @@ class GraphOptimizationConfig:
argument = self.use_cudagraph
class MobaAttentionConfig:
class PlasAttentionConfig:
def __init__(
self,
args,
):
self.moba_encoder_top_k_left: int = None
self.moba_encoder_top_k_right: int = None
"The sparse topk of encoder attention is located at [moba_encoder_top_k_left, moba_encoder top_k_right]"
self.moba_decoder_top_k_left: int = None
self.moba_decoder_top_k_right: int = None
"The sparse topk of decoder attention is located at [moba_decoder_top_k_left, moba_decoder top_k_right]"
self.moba_use_encoder_seq_limit: int = None
"When the number of encdoer token is less than moba_use_encoder_seq_limit, it is not sparse"
self.moba_use_decoder_seq_limit: int = None
"When the number of decdoer token is less than moba_use_decoder_seq_limit, it is not sparse"
self.moba_block_size: int = 128
self.mlp_weight_name: str = "moba_mlp_weight.safetensors"
self.moba_max_seq_length: int = 128 * 1024
self.plas_encoder_top_k_left: int = None
self.plas_encoder_top_k_right: int = None
"The sparse topk of encoder attention is located at [plas_encoder_top_k_left, plas_encoder top_k_right]"
self.plas_decoder_top_k_left: int = None
self.plas_decoder_top_k_right: int = None
"The sparse topk of decoder attention is located at [plas_decoder_top_k_left, plas_decoder top_k_right]"
self.plas_use_encoder_seq_limit: int = None
"When the number of encdoer token is less than plas_use_encoder_seq_limit, it is not sparse"
self.plas_use_decoder_seq_limit: int = None
"When the number of decdoer token is less than plas_use_decoder_seq_limit, it is not sparse"
self.plas_block_size: int = 128
self.mlp_weight_name: str = "plas_attention_mlp_weight.safetensors"
self.plas_max_seq_length: int = 128 * 1024
if args is not None:
for key, value in args.items():
if hasattr(self, key):
setattr(self, key, value)
if self.moba_use_encoder_seq_limit is None and self.moba_encoder_top_k_left is not None:
self.moba_use_encoder_seq_limit = self.moba_encoder_top_k_left * self.moba_block_size
if self.moba_use_decoder_seq_limit is None and self.moba_decoder_top_k_left is not None:
self.moba_use_decoder_seq_limit = self.moba_decoder_top_k_left * self.moba_block_size
if self.plas_use_encoder_seq_limit is None and self.plas_encoder_top_k_left is not None:
self.plas_use_encoder_seq_limit = self.plas_encoder_top_k_left * self.plas_block_size
if self.plas_use_decoder_seq_limit is None and self.plas_decoder_top_k_left is not None:
self.plas_use_decoder_seq_limit = self.plas_decoder_top_k_left * self.plas_block_size
self.check_legality_parameters()
def check_legality_parameters(
self,
) -> None:
if self.moba_encoder_top_k_left is not None:
assert self.moba_encoder_top_k_left > 0, "moba_encoder_top_k_left must large than 0"
if self.plas_encoder_top_k_left is not None:
assert self.plas_encoder_top_k_left > 0, "plas_encoder_top_k_left must large than 0"
if self.moba_encoder_top_k_right is not None:
assert self.moba_encoder_top_k_right > 0, "moba_encoder_top_k_right must large than 0"
if self.plas_encoder_top_k_right is not None:
assert self.plas_encoder_top_k_right > 0, "plas_encoder_top_k_right must large than 0"
assert (
self.moba_encoder_top_k_right >= self.moba_encoder_top_k_left
), "moba_encoder_top_k_right must large than moba_encoder_top_k_left"
self.plas_encoder_top_k_right >= self.plas_encoder_top_k_left
), "plas_encoder_top_k_right must large than plas_encoder_top_k_left"
if self.moba_decoder_top_k_left is not None:
assert self.moba_decoder_top_k_left > 0, "moba_decoder_top_k_left must large than 0"
if self.plas_decoder_top_k_left is not None:
assert self.plas_decoder_top_k_left > 0, "plas_decoder_top_k_left must large than 0"
if self.moba_decoder_top_k_right is not None:
assert self.moba_decoder_top_k_right > 0, "moba_decoder_top_k_right must large than 0"
if self.plas_decoder_top_k_right is not None:
assert self.plas_decoder_top_k_right > 0, "plas_decoder_top_k_right must large than 0"
assert (
self.moba_decoder_top_k_right >= self.moba_decoder_top_k_left
), "moba_decoder_top_k_right must large than moba_decoder_top_k_left"
self.plas_decoder_top_k_right >= self.plas_decoder_top_k_left
), "plas_decoder_top_k_right must large than plas_decoder_top_k_left"
if self.moba_use_encoder_seq_limit is not None and self.moba_encoder_top_k_left is not None:
assert self.moba_use_encoder_seq_limit >= self.moba_encoder_top_k_left * self.moba_block_size
if self.moba_use_decoder_seq_limit is not None and self.moba_decoder_top_k_left is not None:
assert self.moba_use_decoder_seq_limit >= self.moba_decoder_top_k_left * self.moba_block_size
if self.plas_use_encoder_seq_limit is not None and self.plas_encoder_top_k_left is not None:
assert self.plas_use_encoder_seq_limit >= self.plas_encoder_top_k_left * self.plas_block_size
if self.plas_use_decoder_seq_limit is not None and self.plas_decoder_top_k_left is not None:
assert self.plas_use_decoder_seq_limit >= self.plas_decoder_top_k_left * self.plas_block_size
def to_json_string(self):
"""
Convert moba_attention_config to json string.
Convert plas_attention_config to json string.
"""
return json.dumps({key: value for key, value in self.__dict__.items() if value is not None})
@@ -1396,7 +1396,7 @@ class FDConfig:
decoding_config: DecodingConfig = None,
quant_config: QuantConfigBase = None,
graph_opt_config: GraphOptimizationConfig = None,
moba_attention_config: MobaAttentionConfig = None,
plas_attention_config: PlasAttentionConfig = None,
speculative_config: SpeculativeConfig = None,
tokenizer: str = None,
max_model_len: int = 8192,
@@ -1427,7 +1427,7 @@ class FDConfig:
self.early_stop_config: Optional[EarlyStopConfig] = early_stop_config
self.decoding_config: DecodingConfig = decoding_config # type: ignore
self.cache_config: CacheConfig = cache_config # type: ignore
self.moba_attention_config: Optional[MobaAttentionConfig] = moba_attention_config
self.plas_attention_config: Optional[PlasAttentionConfig] = plas_attention_config
# Initialize cuda graph capture list
if self.graph_opt_config.cudagraph_capture_sizes is None:
self.graph_opt_config._set_cudagraph_sizes(max_num_seqs=self.scheduler_config.max_num_seqs)

View File

@@ -30,9 +30,9 @@ from fastdeploy.config import (
FDConfig,
GraphOptimizationConfig,
LoadConfig,
MobaAttentionConfig,
ModelConfig,
ParallelConfig,
PlasAttentionConfig,
PoolerConfig,
RunnerOption,
SpeculativeConfig,
@@ -361,9 +361,9 @@ class EngineArgs:
"""
Configuration for graph optimization backend execution.
"""
moba_attention_config: Optional[Dict[str, Any]] = None
plas_attention_config: Optional[Dict[str, Any]] = None
"""
Configuration for moba attention.
Configuration for plas attention.
"""
enable_logprob: bool = False
@@ -601,9 +601,9 @@ class EngineArgs:
help="",
)
model_group.add_argument(
"--moba-attention-config",
"--plas-attention-config",
type=json.loads,
default=EngineArgs.moba_attention_config,
default=EngineArgs.plas_attention_config,
help="",
)
model_group.add_argument(
@@ -993,17 +993,17 @@ class EngineArgs:
graph_optimization_args[k] = v
return GraphOptimizationConfig(graph_optimization_args)
def create_moba_attention_config(self) -> MobaAttentionConfig:
def create_plas_attention_config(self) -> PlasAttentionConfig:
"""
Create and retuan a MobaAttentionConfig object based on the current settings.
Create and retuan a PlasAttentionConfig object based on the current settings.
"""
attention_args = asdict(self)
if self.moba_attention_config is not None:
for k, v in self.moba_attention_config.items():
if self.plas_attention_config is not None:
for k, v in self.plas_attention_config.items():
attention_args[k] = v
return MobaAttentionConfig(attention_args)
return PlasAttentionConfig(attention_args)
else:
return MobaAttentionConfig(None)
return PlasAttentionConfig(None)
def create_early_stop_config(self) -> EarlyStopConfig:
"""
@@ -1064,7 +1064,7 @@ class EngineArgs:
scheduler_cfg = self.create_scheduler_config()
graph_opt_cfg = self.create_graph_optimization_config()
graph_opt_cfg.update_use_cudagraph(self.use_cudagraph)
moba_attention_config = self.create_moba_attention_config()
plas_attention_config = self.create_plas_attention_config()
early_stop_cfg = self.create_early_stop_config()
early_stop_cfg.update_enable_early_stop(self.enable_early_stop)
@@ -1093,7 +1093,7 @@ class EngineArgs:
max_long_partial_prefills=self.max_long_partial_prefills,
long_prefill_token_threshold=self.long_prefill_token_threshold,
graph_opt_config=graph_opt_cfg,
moba_attention_config=moba_attention_config,
plas_attention_config=plas_attention_config,
guided_decoding_backend=self.guided_decoding_backend,
disable_any_whitespace=self.guided_decoding_disable_any_whitespace,
early_stop_config=early_stop_cfg,

View File

@@ -501,7 +501,7 @@ class LLMEngine:
f" --early_stop_config '{self.cfg.early_stop_config.to_json_string()}'"
f" --reasoning_parser {self.cfg.reasoning_parser}"
f" --load_choices {self.cfg.load_config.load_choices}"
f" --moba_attention_config '{self.cfg.moba_attention_config.to_json_string()}'"
f" --plas_attention_config '{self.cfg.plas_attention_config.to_json_string()}'"
f" --ips {ips}"
f" --cache-transfer-protocol {self.cfg.cache_config.cache_transfer_protocol}"
f" --runner {self.cfg.model_config.runner}"

View File

@@ -20,7 +20,7 @@ from .block_multihead_attn_backend import BlockAttentionBackend
from .flash_attn_backend import FlashAttentionBackend
from .iluvatar_attn_backend import IluvatarAttnBackend
from .mla_attention_backend import MLAAttentionBackend
from .moba_attention_backend import MobaAttentionBackend
from .moba_attention_backend import PlasAttentionBackend
from .native_paddle_backend import PaddleNativeAttnBackend
from .xpu_attn_backend import XPUAttentionBackend
@@ -35,5 +35,5 @@ __all__ = [
"IluvatarAttnBackend",
"BlockAttentionBackend",
"Attention",
"MobaAttentionBackend",
"PlasAttentionBackend",
]

View File

@@ -119,19 +119,19 @@ class Attention(nn.Layer):
self.init_weight()
if (
fd_config.moba_attention_config is not None
and fd_config.moba_attention_config.moba_encoder_top_k_left is not None
and fd_config.moba_attention_config.moba_encoder_top_k_right is not None
and fd_config.moba_attention_config.moba_decoder_top_k_left is not None
and fd_config.moba_attention_config.moba_decoder_top_k_right is not None
fd_config.plas_attention_config is not None
and fd_config.plas_attention_config.plas_encoder_top_k_left is not None
and fd_config.plas_attention_config.plas_encoder_top_k_right is not None
and fd_config.plas_attention_config.plas_decoder_top_k_left is not None
and fd_config.plas_attention_config.plas_decoder_top_k_right is not None
):
mlp_weight_path = os.path.join(
fd_config.model_config.model, fd_config.moba_attention_config.mlp_weight_name
fd_config.model_config.model, fd_config.plas_attention_config.mlp_weight_name
)
self.moba_use_mlp = mlp_weight_path is not None and os.path.exists(mlp_weight_path)
moba_block_size = fd_config.moba_attention_config.moba_block_size
moba_max_seq_length = fd_config.moba_attention_config.moba_max_seq_length
if self.moba_use_mlp:
self.plas_use_mlp = mlp_weight_path is not None and os.path.exists(mlp_weight_path)
plas_block_size = fd_config.plas_attention_config.plas_block_size
plas_max_seq_length = fd_config.plas_attention_config.plas_max_seq_length
if self.plas_use_mlp:
mlp_weight = {}
with safe_open(mlp_weight_path, framework="np", device="cpu") as f:
for key_name in f.keys():
@@ -148,12 +148,12 @@ class Attention(nn.Layer):
* self.kv_num_heads : (fd_config.parallel_config.tensor_parallel_rank + 1)
* self.kv_num_heads
]
assert self.attn_gate_weight.shape[1] % moba_block_size == 0
assert self.attn_gate_weight.shape[1] % plas_block_size == 0
self.cache_k_block_means = paddle.zeros(
[
fd_config.scheduler_config.max_num_seqs,
moba_max_seq_length // moba_block_size,
plas_max_seq_length // plas_block_size,
self.kv_num_heads,
self.head_dim,
],

View File

@@ -39,7 +39,7 @@ from fastdeploy.model_executor.layers.attention.base_attention_backend import (
@dataclass
class MobaAttentionMetadata(AttentionMetadata):
class PlasAttentionMetadata(AttentionMetadata):
"""
AppendAttentionMetadata
"""
@@ -54,7 +54,7 @@ class MobaAttentionMetadata(AttentionMetadata):
max_dec_len_this_time: int = 0
class MobaAttentionBackend(AttentionBackend):
class PlasAttentionBackend(AttentionBackend):
"""
The backend class that uses paddle native attention implementation.
Which is used only for testing purpose.
@@ -70,11 +70,11 @@ class MobaAttentionBackend(AttentionBackend):
decoder_block_shape_q: int = -1,
) -> None:
"""
MobaAttentionBackend __init__
PlasAttentionBackend __init__
"""
super().__init__()
self.attention_metadata: MobaAttentionMetadata = None
assert fd_config.moba_attention_config is not None, "moba_attention_config is None"
self.attention_metadata: PlasAttentionMetadata = None
assert fd_config.plas_attention_config is not None, "plas_attention_config is None"
self.block_size = fd_config.parallel_config.block_size
self.max_seq_len = fd_config.parallel_config.max_model_len
self.max_num_seqs = fd_config.scheduler_config.max_num_seqs
@@ -83,18 +83,18 @@ class MobaAttentionBackend(AttentionBackend):
self.head_dim = fd_config.model_config.head_dim
self.num_layers: int = fd_config.model_config.num_hidden_layers
self.attn_block_m = 128
self.moba_block_size = fd_config.moba_attention_config.moba_block_size
self.moba_encoder_top_k_left = int(fd_config.moba_attention_config.moba_encoder_top_k_left)
self.moba_encoder_top_k_right = int(fd_config.moba_attention_config.moba_encoder_top_k_right)
self.moba_use_encoder_seq_limit = int(fd_config.moba_attention_config.moba_use_encoder_seq_limit)
self.moba_decoder_top_k_left = int(fd_config.moba_attention_config.moba_decoder_top_k_left)
self.moba_decoder_top_k_right = int(fd_config.moba_attention_config.moba_decoder_top_k_right)
self.moba_use_decoder_seq_limit = int(fd_config.moba_attention_config.moba_use_decoder_seq_limit)
self.moba_max_seq_length = fd_config.moba_attention_config.moba_max_seq_length
self.plas_block_size = fd_config.plas_attention_config.plas_block_size
self.plas_encoder_top_k_left = int(fd_config.plas_attention_config.plas_encoder_top_k_left)
self.plas_encoder_top_k_right = int(fd_config.plas_attention_config.plas_encoder_top_k_right)
self.plas_use_encoder_seq_limit = int(fd_config.plas_attention_config.plas_use_encoder_seq_limit)
self.plas_decoder_top_k_left = int(fd_config.plas_attention_config.plas_decoder_top_k_left)
self.plas_decoder_top_k_right = int(fd_config.plas_attention_config.plas_decoder_top_k_right)
self.plas_use_decoder_seq_limit = int(fd_config.plas_attention_config.plas_use_decoder_seq_limit)
self.plas_max_seq_length = fd_config.plas_attention_config.plas_max_seq_length
def init_attention_metadata(self, forward_meta: ForwardMeta):
"""Init the metadata for a forward pass."""
metadata = MobaAttentionMetadata()
metadata = PlasAttentionMetadata()
metadata._dtype = paddle.get_default_dtype()
metadata.cu_seq_q_pack, metadata.cu_seqlens_k, metadata.q_pack_tokens = get_cur_cu_seq_len_k(
forward_meta.seq_lens_encoder,
@@ -116,7 +116,7 @@ class MobaAttentionBackend(AttentionBackend):
[k_token_num + self.attn_block_m, self.kv_num_heads * self.head_dim], dtype=metadata._dtype
)
self.attention_metadata = metadata
assert self.max_seq_len <= self.moba_max_seq_length
assert self.max_seq_len <= self.plas_max_seq_length
def get_kv_cache_shape(
self,
@@ -186,13 +186,13 @@ class MobaAttentionBackend(AttentionBackend):
self.max_seq_len,
attention_metadata.max_enc_len_this_time,
attention_metadata.max_dec_len_this_time,
self.moba_encoder_top_k_left,
self.moba_encoder_top_k_right,
self.moba_use_encoder_seq_limit,
self.moba_decoder_top_k_left,
self.moba_decoder_top_k_right,
self.moba_use_decoder_seq_limit,
layer.moba_use_mlp,
self.plas_encoder_top_k_left,
self.plas_encoder_top_k_right,
self.plas_use_encoder_seq_limit,
self.plas_decoder_top_k_left,
self.plas_decoder_top_k_right,
self.plas_use_decoder_seq_limit,
layer.plas_use_mlp,
getattr(layer, "cache_quant_type_str", "none"),
)[0]
return out

View File

@@ -26,7 +26,7 @@ class _Backend(enum.Enum):
MLA_ATTN = enum.auto()
FLASH_ATTN = enum.auto()
BLOCK_ATTN = enum.auto()
MOBA_ATTN = enum.auto()
PLAS_ATTN = enum.auto()
class Platform:

View File

@@ -64,9 +64,9 @@ class CUDAPlatform(Platform):
elif selected_backend == _Backend.FLASH_ATTN:
logger.info("Using FLASH ATTN backend.")
return "fastdeploy.model_executor.layers.attention.FlashAttentionBackend"
elif selected_backend == _Backend.MOBA_ATTN:
logger.info("Using MOBA ATTN backend.")
return "fastdeploy.model_executor.layers.attention.MobaAttentionBackend"
elif selected_backend == _Backend.PLAS_ATTN:
logger.info("Using PLAS ATTN backend.")
return "fastdeploy.model_executor.layers.attention.PlasAttentionBackend"
else:
raise ValueError(
"Invalid attention backend you specified.\n"

View File

@@ -61,7 +61,7 @@ class RolloutModelConfig:
graph_optimization_config: str = None,
early_stop_config: str = None,
local_rank: int = 0,
moba_attention_config: str = None,
plas_attention_config: str = None,
data_parallel_size: int = 1,
num_nextn_predict_layers: int = 0,
):
@@ -109,7 +109,7 @@ class RolloutModelConfig:
self.local_rank = local_rank
self.early_stop_config = early_stop_config
self.ips = None
self.moba_attention_config = moba_attention_config
self.plas_attention_config = plas_attention_config
self.num_nextn_predict_layers = num_nextn_predict_layers
def __str__(self):

View File

@@ -35,9 +35,9 @@ from fastdeploy.config import (
FDConfig,
GraphOptimizationConfig,
LoadConfig,
MobaAttentionConfig,
ModelConfig,
ParallelConfig,
PlasAttentionConfig,
SpeculativeConfig,
)
from fastdeploy.input.ernie4_5_tokenizer import Ernie4_5Tokenizer
@@ -577,10 +577,10 @@ def parse_args():
help="Configuration of Graph optimization backend.",
)
parser.add_argument(
"--moba_attention_config",
"--plas_attention_config",
type=json.loads,
default=None,
help="Configuration of moba attention.",
help="Configation of plas attention.",
)
parser.add_argument(
"--guided_decoding_backend",
@@ -723,7 +723,7 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
graph_opt_config = GraphOptimizationConfig(args.graph_optimization_config)
moba_attention_config = MobaAttentionConfig(args.moba_attention_config)
plas_attention_config = PlasAttentionConfig(args.plas_attention_config)
early_stop_config = EarlyStopConfig(args.early_stop_config)
@@ -795,7 +795,7 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
cache_config=cache_config,
scheduler_config=scheduler_config,
ips=args.ips,
moba_attention_config=moba_attention_config,
plas_attention_config=plas_attention_config,
)
update_fd_config_for_mm(fd_config)

View File

@@ -57,7 +57,7 @@ def naive_attn(q_input, k_input, v_input, mask):
return out
class TestMobaAttention(unittest.TestCase):
class TestPlasAttention(unittest.TestCase):
def setUp(self):
paddle.seed(0)
self.seq_len = int(8 * 1024)
@@ -65,15 +65,15 @@ class TestMobaAttention(unittest.TestCase):
self.num_kv_heads = int(1)
self.head_dim = int(128)
self.max_num_seqs = 1
self.moba_max_seq_length = int(128 * 1024)
self.moba_block_size = int(128)
self.moba_encoder_top_k_left = 2
self.moba_encoder_top_k_right = 3
self.moba_use_encoder_seq_limit = int(4 * 1024)
self.plas_max_seq_length = int(128 * 1024)
self.plas_block_size = int(128)
self.plas_encoder_top_k_left = 2
self.plas_encoder_top_k_right = 3
self.plas_use_encoder_seq_limit = int(4 * 1024)
self.cache_k_block_means = paddle.zeros(
[
self.max_num_seqs,
self.moba_max_seq_length // self.moba_block_size,
self.plas_max_seq_length // self.plas_block_size,
self.num_kv_heads,
self.head_dim,
],
@@ -96,12 +96,12 @@ class TestMobaAttention(unittest.TestCase):
self.rotary_embs = paddle.ones([2, self.seq_len, self.head_dim // 2], dtype="float32")
self.attn_gate_weight = paddle.randn(
[self.num_kv_heads, self.moba_block_size, self.head_dim], dtype="bfloat16"
[self.num_kv_heads, self.plas_block_size, self.head_dim], dtype="bfloat16"
)
self.gqa_group_size = self.num_heads // self.num_kv_heads
self.num_blocks = (self.seq_len + self.moba_block_size - 1) // self.moba_block_size
self.num_blocks = (self.seq_len + self.plas_block_size - 1) // self.plas_block_size
self.sparse_step = 4
@@ -115,38 +115,38 @@ class TestMobaAttention(unittest.TestCase):
for i in range(self.max_num_seqs):
k_padding = paddle.zeros(
[
(self.seq_len + self.moba_block_size - 1) // self.moba_block_size * self.moba_block_size,
(self.seq_len + self.plas_block_size - 1) // self.plas_block_size * self.plas_block_size,
self.num_kv_heads,
self.head_dim,
],
dtype="bfloat16",
)
k_padding[0 : self.seq_len] = self.k_input[i * self.seq_len : (i + 1) * self.seq_len]
real_k_block_means = k_padding.reshape([-1, self.moba_block_size, self.num_kv_heads, self.head_dim])
real_k_block_means = k_padding.reshape([-1, self.plas_block_size, self.num_kv_heads, self.head_dim])
real_k_block_means = real_k_block_means.mean(axis=1)
compute_k_block_means = self.cache_k_block_means[i, 0 : real_k_block_means.shape[0]]
assert (compute_k_block_means - real_k_block_means).abs().max() < 0.003
print("[consistency]Moba attention: split_qkv_rope matches.")
print("[consistency]plas attention: split_qkv_rope matches.")
def compare_mlp_einsum(self, k_gate_weight):
for i in range(self.max_num_seqs):
k_padding = paddle.zeros(
[
(self.seq_len + self.moba_block_size - 1) // self.moba_block_size * self.moba_block_size,
(self.seq_len + self.plas_block_size - 1) // self.plas_block_size * self.plas_block_size,
self.num_kv_heads,
self.head_dim,
],
dtype="bfloat16",
)
k_padding[0 : self.seq_len] = self.k_input[i * self.seq_len : (i + 1) * self.seq_len]
k_padding = k_padding.reshape([-1, self.moba_block_size, self.num_kv_heads, self.head_dim])
k_padding = k_padding.reshape([-1, self.plas_block_size, self.num_kv_heads, self.head_dim])
real_result = paddle.einsum("nbhd,hbd->nhd", k_padding, self.attn_gate_weight)
compute_result = k_gate_weight[i][0 : real_result.shape[0]]
assert (real_result - compute_result).abs().max() < 0.5
print("[consistency]Moba attention: MLP einsum matches.")
print("[consistency]plas attention: MLP einsum matches.")
def compare_qk_gemm(self, qk_gate_weight):
for i in range(self.max_num_seqs):
@@ -170,10 +170,10 @@ class TestMobaAttention(unittest.TestCase):
conpute_result = qk_gate_weight[i * self.seq_len : (i + 1) * self.seq_len, :, 0 : self.num_blocks]
assert (qk_gemm_out - conpute_result).abs().max() < 1e-4
print("[consistency]Moba attention: qk_gemm matches.")
print("[consistency]plas attention: qk_gemm matches.")
def compare_qk_gate_topk(self, qk_gate_topk_idx):
limit_topk = self.moba_use_encoder_seq_limit // self.moba_block_size
limit_topk = self.plas_use_encoder_seq_limit // self.plas_block_size
for i in range(self.max_num_seqs):
qk_gate_topk_idx_batch = qk_gate_topk_idx[i * self.num_blocks : (i + 1) * self.num_blocks]
qk_gate_topk_idx_batch_no_sparse = qk_gate_topk_idx_batch[0 : limit_topk - 1]
@@ -191,40 +191,40 @@ class TestMobaAttention(unittest.TestCase):
- paddle.ones(qk_gate_topk_idx_batch_sparse.shape, qk_gate_topk_idx_batch_sparse.dtype)
* self.sparse_step
).abs().max() < 1e-6
print("[consistency]Moba attention: qk_gate_topk matches.")
print("[consistency]plas attention: qk_gate_topk matches.")
def compare_attn(self, attn_out, qk_gate_topk_idx):
x = (
paddle.tensor.triu(paddle.ones([self.moba_block_size, self.moba_block_size], dtype="bfloat16"), 1)
paddle.tensor.triu(paddle.ones([self.plas_block_size, self.plas_block_size], dtype="bfloat16"), 1)
* -1000000
)
limit_topk = self.moba_use_encoder_seq_limit // self.moba_block_size
limit_topk = self.plas_use_encoder_seq_limit // self.plas_block_size
for i in range(self.max_num_seqs):
q_input = self.q_input[i * self.seq_len : (i + 1) * self.seq_len].unsqueeze(axis=0)
k_input = self.k_input[i * self.seq_len : (i + 1) * self.seq_len].unsqueeze(axis=0)
v_input = self.v_input[i * self.seq_len : (i + 1) * self.seq_len].unsqueeze(axis=0)
mask = paddle.tensor.triu(paddle.ones([self.seq_len, self.seq_len], dtype="bfloat16"), 1) * -1000000
mask[self.moba_use_encoder_seq_limit - self.moba_block_size :] = -1000000
mask[self.plas_use_encoder_seq_limit - self.plas_block_size :] = -1000000
for i in range(limit_topk - 1, self.num_blocks):
n_block = i
mask[
i * self.moba_block_size : i * self.moba_block_size + self.moba_block_size,
n_block * self.moba_block_size : n_block * self.moba_block_size + self.moba_block_size,
i * self.plas_block_size : i * self.plas_block_size + self.plas_block_size,
n_block * self.plas_block_size : n_block * self.plas_block_size + self.plas_block_size,
] = x
idx = 0
n_block -= int(qk_gate_topk_idx[i, 0, idx])
idx += 1
while n_block >= 0:
mask[
i * self.moba_block_size : i * self.moba_block_size + self.moba_block_size,
n_block * self.moba_block_size : n_block * self.moba_block_size + self.moba_block_size,
i * self.plas_block_size : i * self.plas_block_size + self.plas_block_size,
n_block * self.plas_block_size : n_block * self.plas_block_size + self.plas_block_size,
] = 0
n_block -= int(qk_gate_topk_idx[i, 0, idx])
idx += 1
naive_attn_out = naive_attn(q_input, k_input, v_input, mask).squeeze(axis=0).transpose([1, 0, 2])
assert (attn_out - naive_attn_out).abs().max() < 0.016
def test_moba_attention(self):
def test_plas_attention(self):
qkv_out = paddle.randn([self.tokens, self.num_heads + 2 * self.num_kv_heads, self.head_dim], dtype="bfloat16")
seq_len_encoder = paddle.to_tensor([self.seq_len] * self.max_num_seqs, dtype="int32")
@@ -255,7 +255,7 @@ class TestMobaAttention(unittest.TestCase):
self.num_heads,
self.num_kv_heads,
self.head_dim,
self.moba_max_seq_length,
self.plas_max_seq_length,
self.seq_len,
self.seq_len,
"none",
@@ -307,9 +307,9 @@ class TestMobaAttention(unittest.TestCase):
self.seq_len,
self.num_heads,
self.num_kv_heads,
self.moba_encoder_top_k_left,
self.moba_encoder_top_k_right,
self.moba_use_encoder_seq_limit,
self.plas_encoder_top_k_left,
self.plas_encoder_top_k_right,
self.plas_use_encoder_seq_limit,
)
self.compare_qk_gate_topk(qk_gate_topk_idx)
@@ -332,7 +332,7 @@ class TestMobaAttention(unittest.TestCase):
self.num_heads,
self.num_kv_heads,
self.head_dim,
self.moba_max_seq_length,
self.plas_max_seq_length,
)
self.compare_attn(attn_out, qk_gate_topk_idx)
@@ -340,18 +340,18 @@ class TestMobaAttention(unittest.TestCase):
def test_server(self):
if get_cur_cu_seq_len_k is None:
return
os.environ["FD_ATTENTION_BACKEND"] = "MOBA_ATTN"
os.environ["FD_ATTENTION_BACKEND"] = "PLAS_ATTN"
base_path = os.getenv("MODEL_PATH")
if base_path:
model_path = os.path.join(base_path, "./ernie-4_5-21b-a3b-bf16-paddle")
else:
model_path = "./ernie-4_5-21b-a3b-bf16-paddle"
moba_attention_config = {
"moba_encoder_top_k_left": 50,
"moba_encoder_top_k_right": 60,
"moba_decoder_top_k_left": 100,
"moba_decoder_top_k_right": 120,
plas_attention_config = {
"plas_encoder_top_k_left": 50,
"plas_encoder_top_k_right": 60,
"plas_decoder_top_k_left": 100,
"plas_decoder_top_k_right": 120,
}
# 加载模型
@@ -365,7 +365,7 @@ class TestMobaAttention(unittest.TestCase):
quantization="wint4",
enable_chunked_prefill=True,
max_num_batched_tokens=8192,
moba_attention_config=moba_attention_config,
plas_attention_config=plas_attention_config,
)
prompts = ["Hello world!"]

View File

@@ -65,7 +65,7 @@ class TestAttentionInitWeight(unittest.TestCase):
self.fd_config.parallel_config = self.parallel_config
self.fd_config.cache_config = self.cache_config
self.fd_config.quant_config = None
self.fd_config.moba_attention_config = None
self.fd_config.plas_attention_config = None
def test_init_weight_without_quantization(self):
"""Test init_weight without quantization."""
@@ -141,7 +141,7 @@ class TestAttentionWeightLoader(unittest.TestCase):
self.fd_config.model_config = self.model_config
self.fd_config.parallel_config = self.parallel_config
self.fd_config.cache_config = self.cache_config
self.fd_config.moba_attention_config = None
self.fd_config.plas_attention_config = None
# Create mock quant method
self.mock_quant_method = MockQuantMethod()