Skip to content

[Models] support swa for fleet fallback model #8090

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chang-wenbin merged 4 commits into from
Jul 6, 2026
Merged

[Models] support swa for fleet fallback model #8090

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ed7f39d
update swa
xiaoguoguo626807 Jul 2, 2026
bf0419a
updtae test
xiaoguoguo626807 Jul 2, 2026
4e0edca
clear test fix ci
xiaoguoguo626807 Jul 2, 2026
4ed6cab
fix
xiaoguoguo626807 Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
227 changes: 150 additions & 77 deletions fastdeploy/model_executor/models/paddleformers/base_fleet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,6 @@
from paddlefleet.models.gpt.lm_head import GPTLMHead
from paddlefleet.transformer.layer import FleetLayer
from paddlefleet.transformer.transformer_config import TransformerConfig
from paddleformers.transformers import AutoConfig
from paddleformers.transformers.auto.modeling import AutoModelForCausalLM
from paddleformers.utils.log import logger

from fastdeploy.model_executor.forward_meta import ForwardMeta # noqa: F401
Expand All @@ -47,6 +45,8 @@

from fastdeploy.model_executor.layers.attention.attention import Attention

USE_ERNIE = False

class FastDeployAttention(FleetLayer):
"""
FastDeploy version of DotProductAttention, holding an internal FastDeploy Attention module.
Expand All @@ -65,6 +65,8 @@ def __init__(
hidden_size_per_attention_head: int,
hidden_size_per_partition: int,
layer_id: int,
window_attn_skip_freq=None,
sliding_window: int = 0,
):
"""
Initialize FastDeployAttention.
Expand All @@ -86,6 +88,8 @@ def __init__(
self.hidden_size_per_attention_head = hidden_size_per_attention_head
self.hidden_size_per_partition = hidden_size_per_partition
self.layer_id = layer_id
self.window_attn_skip_freq = window_attn_skip_freq
self.sliding_window = sliding_window

def forward(
self,
Expand Down Expand Up @@ -167,85 +171,123 @@ def squeeze_to_3d(t: paddle.Tensor, name: str) -> paddle.Tensor:
need_do_prefill = forward_meta.max_len_tensor_cpu[1] > 0
need_do_decode = forward_meta.max_len_tensor_cpu[2] > 0

# MLA mode: pass q, k, v, compressed_kv, k_pe separately
# Reference: deepseek_v3.py line 389
#
# Note:
# - Prefill (flash_attn_func): expects 3D tensors [seq, heads, dim]
# - Decode (multi_head_latent_attention): expects 2D tensors [seq, heads*dim]
# So we need to flatten q for decode phase only

# Process compressed_kv and k_pe

assert kv_compressed is not None, "kv_compressed must be provided when use"
compressed_kv = kv_compressed.squeeze(0)
k_pos_emb = k_pos_emb.squeeze(0)

output = None
fmqa_out = None
if need_do_prefill:
# Prefill: keep 3D tensors for flash_attn_func
output = self.fd_attention.forward(
q=q,
k=k,
v=v,
qkv=None,
compressed_kv=compressed_kv,
k_pe=k_pos_emb,
forward_meta=forward_meta,
k_pos_emb_sq = k_pos_emb.squeeze(0)

if self.window_attn_skip_freq is not None and self.window_attn_skip_freq[self.layer_id] == 1:

@PaddlePaddle-bot PaddlePaddle-bot Jul 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 建议 这里新增了 window_attn_skip_freq 命中的 MLA/SWA 层直接走 DSAAttentionBackend.forward_static 的前向路径，但 PR 只提供了 mock/shape 级单测，Accuracy Tests 也标为 N/A。该分支会改变 prefill/decode 的 mask、indexer、cache 写入和 softmax scale 使用，输出直接影响 logits；没有 fallback logits 或参考实现对齐，后续 get_swa_indexer_top_k、cache layout、V de-absorb 任一错误都不会被发现。

建议修复方式：
补一个小模型/固定 seed 的 PaddleFleet fallback SWA 对齐用例，至少覆盖 window_attn_skip_freq[layer_id] == 1 的 prefill-only 和 mixed prefill+decode，并与 DeepSeek/PaddleFleet 参考路径比较 logits 容差；如果 CI 无法跑真实 GPU kernel，请在 PR 中附可复现命令和对齐结果。

kv_lora_rank = self.config.kv_lora_rank

q_input = squeeze_to_3d(q_absorbed, "q_absorbed") if q_absorbed.ndim == 4 else q_absorbed
Comment thread
xiaoguoguo626807 marked this conversation as resolved.
num_attention_heads_tp = q_input.shape[1]

"""DSA sliding-window attention path, mirroring DeepseekV3MLAAttention.forward_swa_static."""
from fastdeploy.model_executor.layers.attention import (
DSAAttentionBackend,
)
from fastdeploy.model_executor.models.deepseek_v3 import (
get_swa_indexer_top_k,
)

window_size = (
self.sliding_window[0]
if isinstance(self.sliding_window, (list, tuple))
else self.sliding_window
)
output.reshape_([output.shape[0], output.shape[1] * output.shape[2]])

if need_do_decode:
# Decode: use absorbed q for multi_head_latent_attention C++ kernel
# q_absorbed: [s, heads, kv_lora_rank + qk_rope_head_dim] (after squeeze_to_3d)
# C++ kernel expects: [token_num, heads * (kv_lora_rank + qk_rope_head_dim)]
q_abs = squeeze_to_3d(q_absorbed, "q_absorbed") if q_absorbed.ndim == 4 else q_absorbed
seq_len = int(q_abs.shape[0])
q_input = q_abs.reshape([seq_len, -1])

fmqa_out = self.fd_attention.forward(
q=q_input,
k=None,
v=None,
qkv=None,
indexer_top_k = paddle.full([q_input.shape[0], 1, window_size], -1, dtype="int32")

get_swa_indexer_top_k(
indexer_top_k,
forward_meta.block_tables,
forward_meta.cu_seqlens_q,
forward_meta.seq_lens_encoder,
forward_meta.seq_lens_decoder,
forward_meta.batch_id_per_token,
)
fmqa_out = DSAAttentionBackend.forward_static(
q=q_input.contiguous(),
indexer_topk=indexer_top_k,
compressed_kv=compressed_kv,
k_pe=k_pos_emb,
k_pe=k_pos_emb_sq,
latent_cache=forward_meta.caches[self.layer_id],
forward_meta=forward_meta,
attn_softmax_scale=self.softmax_scale,
)

# V de-absorption: kernel output [token, heads * kv_lora_rank]
# -> [heads, token, kv_lora_rank] @ wv_b [heads, kv_lora_rank, v_head_dim]
# -> [token, heads * v_head_dim]
kv_lora_rank = self.config.kv_lora_rank
v_head_dim = self.config.v_head_dim
num_heads = fmqa_out.shape[-1] // kv_lora_rank
fmqa_out = fmqa_out.reshape([-1, num_heads, kv_lora_rank]).transpose([1, 0, 2])
fmqa_out = fmqa_out.reshape_([-1, num_attention_heads_tp, kv_lora_rank]).transpose([1, 0, 2])
fmqa_out = paddle.bmm(fmqa_out, v_b_proj_weight)
fmqa_out = fmqa_out.transpose([1, 0, 2]).reshape([-1, num_heads * v_head_dim])
# Merge prefill and decode outputs if both are present
output = fmqa_out.transpose([1, 0, 2]).reshape(
[-1, num_attention_heads_tp * self.config.v_head_dim]
)

else:
output = None
fmqa_out = None
if need_do_prefill:
try:
from fastdeploy.model_executor.ops.gpu import (
merge_prefill_decode_output,
)

merge_prefill_decode_output(
output,
fmqa_out,
forward_meta.seq_lens_encoder,
forward_meta.seq_lens_decoder,
forward_meta.seq_lens_this_time,
forward_meta.cu_seqlens_q,
num_heads,
v_head_dim,
1,
)
except (ImportError, AttributeError):
logger.warning("merge_prefill_decode_output not available, using decode output only")
# Prefill: keep 3D tensors for flash_attn_func
output = self.fd_attention.forward(
q=q,
k=k,
v=v,
qkv=None,
compressed_kv=compressed_kv,
k_pe=k_pos_emb_sq,
forward_meta=forward_meta,
)
output.reshape_([output.shape[0], output.shape[1] * output.shape[2]])

if need_do_decode:
# Decode: use absorbed q for multi_head_latent_attention C++ kernel
# q_absorbed: [s, heads, kv_lora_rank + qk_rope_head_dim] (after squeeze_to_3d)
# C++ kernel expects: [token_num, heads * (kv_lora_rank + qk_rope_head_dim)]
q_abs = squeeze_to_3d(q_absorbed, "q_absorbed") if q_absorbed.ndim == 4 else q_absorbed
seq_len = int(q_abs.shape[0])
q_input = q_abs.reshape([seq_len, -1])

fmqa_out = self.fd_attention.forward(
q=q_input,
k=None,
v=None,
qkv=None,
compressed_kv=compressed_kv,
k_pe=k_pos_emb_sq,
forward_meta=forward_meta,
)

# V de-absorption: kernel output [token, heads * kv_lora_rank]
# -> [heads, token, kv_lora_rank] @ wv_b [heads, kv_lora_rank, v_head_dim]
# -> [token, heads * v_head_dim]
kv_lora_rank = self.config.kv_lora_rank
v_head_dim = self.config.v_head_dim
num_heads = fmqa_out.shape[-1] // kv_lora_rank
fmqa_out = fmqa_out.reshape([-1, num_heads, kv_lora_rank]).transpose([1, 0, 2])
fmqa_out = paddle.bmm(fmqa_out, v_b_proj_weight)
fmqa_out = fmqa_out.transpose([1, 0, 2]).reshape([-1, num_heads * v_head_dim])
# Merge prefill and decode outputs if both are present
if need_do_prefill:
try:
from fastdeploy.model_executor.ops.gpu import (
merge_prefill_decode_output,
)

merge_prefill_decode_output(
output,
fmqa_out,
forward_meta.seq_lens_encoder,
forward_meta.seq_lens_decoder,
forward_meta.seq_lens_this_time,
forward_meta.cu_seqlens_q,
num_heads,
v_head_dim,
1,
)
except (ImportError, AttributeError):
logger.warning(
"merge_prefill_decode_output not available, using decode output only"
)
output = fmqa_out
else:
output = fmqa_out
else:
output = fmqa_out
else:
# Standard mode: concatenate QKV
seq_len = int(q.shape[0])
Expand Down Expand Up @@ -286,7 +328,22 @@ def __init__(self, fd_config: "FDConfig", **kwargs):
logger.info("Initializing PaddleFormers backend.")
self.fd_config = fd_config # FastDeploy's top-level FDConfig
self.model_config = fd_config.model_config # FastDeploy's ModelConfig
self.paddleformers_config = AutoConfig.from_pretrained(self.model_config.model)
if USE_ERNIE:
from paddleformers.transformers.configuration_utils import (
PretrainedConfig,
)

_config_dict, _ = PretrainedConfig.get_config_dict(
self.model_config.model, _configuration_file="model_config.json"
)
from ernie5.pretrain import Ernie5V2Config

self.paddleformers_config = Ernie5V2Config.from_dict(_config_dict)
self.paddleformers_config.moe_dequant_input = True
else:
from paddleformers.transformers import AutoConfig

self.paddleformers_config = AutoConfig.from_pretrained(self.model_config.model)

# Assign parallel config from fd_config.parallel_config to paddleformers_config
parallel_config = fd_config.parallel_config
Expand All @@ -306,6 +363,7 @@ def __init__(self, fd_config: "FDConfig", **kwargs):
self.paddleformers_config.use_cpu_initialization = True
self.paddleformers_config.perform_initialization = False
self.paddleformers_config.gated_attention = getattr(self.paddleformers_config, "use_gated_attn", False)
self.paddleformers_config.moe_layer_interval = getattr(self.paddleformers_config, "moe_layer_freq", 1)
if getattr(self.paddleformers_config, "multi_latent_attention", False):
self.paddleformers_config.qk_head_dim = (
self.paddleformers_config.qk_rope_head_dim + self.paddleformers_config.qk_nope_head_dim
Expand Down Expand Up @@ -334,11 +392,24 @@ def __init__(self, fd_config: "FDConfig", **kwargs):
"load_via_cpu": True,
"load_checkpoint_format": "flex_checkpoint",
}
# Set random seed before model construction for reproducibility
self.model = AutoModelForCausalLM.from_pretrained(
self.model_config.model,
**model_load_kwargs,
)
if USE_ERNIE:
from fleet_bridge import AutoModelForCausalLM

self.model = AutoModelForCausalLM.from_pretrained(
self.model_config.model,
config=self.paddleformers_config,
dtype=self.model_config.dtype,
)
else:
from paddleformers.transformers.auto.modeling import (
AutoModelForCausalLM,
)

# Set random seed before model construction for reproducibility
self.model = AutoModelForCausalLM.from_pretrained(
self.model_config.model,
**model_load_kwargs,
)

self.model.eval()
# Patch PaddleFleet core_attention with FastDeploy attention
Expand Down Expand Up @@ -721,6 +792,8 @@ def patch_paddlefleet_core_attention(
hidden_size_per_attention_head=hidden_size_per_attention_head,
hidden_size_per_partition=hidden_size_per_partition,
layer_id=fd_layer_id,
window_attn_skip_freq=getattr(fd_config.model_config, "window_attn_skip_freq", None),
sliding_window=getattr(fd_config.model_config, "sliding_window", 0),
)

# Replace core_attention object
Expand Down
Loading
Loading