[Feature] block sparse attention (#3209)
Some checks failed
CE Compile Job / ce_job_pre_check (push) Has been cancelled
CE Compile Job / print_ce_job_pre_check_outputs (push) Has been cancelled
CE Compile Job / FD-Clone-Linux (push) Has been cancelled
CE Compile Job / Show Code Archive Output (push) Has been cancelled
CE Compile Job / BUILD_SM8090 (push) Has been cancelled
CE Compile Job / BUILD_SM8689 (push) Has been cancelled
CE Compile Job / CE_UPLOAD (push) Has been cancelled
Deploy GitHub Pages / deploy (push) Has been cancelled

* 支持稀疏attn

* fix bug

* code style

* fix moba attn get kv shape

* 修复a100编译

* codestyle

* code style

* code style

* code style

* fix conflict

* 增加单侧

* code style

* 增加eblite 加载时间

* fix bug

* for ci

* for ci

* for ci

* for ci

* 支持mlp block size 128

* 增加小算子单测

* fix 单测 mlp

* 将环境变量加入到config里面

* fix rollout config
This commit is contained in:
yangjianfengo1
2025-08-26 22:16:04 +08:00
committed by GitHub
parent f0a362af18
commit 646a0c2fd8
31 changed files with 6507 additions and 10 deletions

View File

@@ -34,6 +34,7 @@ from fastdeploy.config import (
FDConfig,
GraphOptimizationConfig,
LoadConfig,
MobaAttentionConfig,
ModelConfig,
ParallelConfig,
SpeculativeConfig,
@@ -541,6 +542,12 @@ def parse_args():
default=None,
help="Configation of Graph optimization backend.",
)
parser.add_argument(
"--moba_attention_config",
type=json.loads,
default=None,
help="Configation of moba attention.",
)
parser.add_argument(
"--guided_decoding_backend",
type=str,
@@ -640,6 +647,8 @@ 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)
early_stop_config = EarlyStopConfig(args.early_stop_config)
# Note(tangbinhan): used for load_checkpoint
@@ -719,6 +728,7 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
cache_config=cache_config,
engine_worker_queue_port=args.engine_worker_queue_port,
ips=args.ips,
moba_attention_config=moba_attention_config,
)
update_fd_config_for_mm(fd_config)