[models] support fleet fallback with SWA

fastdeploy/model_executor/layers/attention/append_attn_backend.py

@@ -319,6 +319,8 @@ def forward_mixed(
         rope_already_applied = getattr(forward_meta, "rope_already_applied", False)
         if rope_already_applied and forward_meta.rotary_embs is not None:
             forward_meta.rotary_embs = self._get_identity_rotary_embs(forward_meta.rotary_embs)
+            if forward_meta.swa_rotary_embs is not None:
+                forward_meta.swa_rotary_embs = self._get_identity_rotary_embs(forward_meta.swa_rotary_embs)
 
         sliding_window = 0
         rotary_embs = forward_meta.rotary_embs

fastdeploy/model_executor/models/paddleformers/base_fleet.py

@@ -17,6 +17,7 @@
 """Generic PaddleFormers modeling backend base class."""
 
 import logging
+import os
 
 from fastdeploy.model_executor.utils import is_paddlefleet_available
 
@@ -45,7 +46,7 @@
 
     from fastdeploy.model_executor.layers.attention.attention import Attention
 
-    USE_ERNIE = False
+    USE_ERNIE = os.environ.get("FD_FALLBACK_FLEET_USE_ERNIE", False)
 
     class FastDeployAttention(FleetLayer):
         """
@@ -214,11 +215,9 @@ def squeeze_to_3d(t: paddle.Tensor, name: str) -> paddle.Tensor:
                             attn_softmax_scale=self.softmax_scale,
                         )
 
-                        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)
-                        output = fmqa_out.transpose([1, 0, 2]).reshape(
-                            [-1, num_attention_heads_tp * self.config.v_head_dim]
-                        )
+                        fmqa_out = fmqa_out.reshape_([-1, num_attention_heads_tp, kv_lora_rank])
+                        fmqa_out = paddle.einsum("thr,hrv->thv", fmqa_out, v_b_proj_weight)
+                        output = fmqa_out.reshape([-1, num_attention_heads_tp * self.config.v_head_dim])
 
                     else:
                         output = None
@@ -260,9 +259,9 @@ def squeeze_to_3d(t: paddle.Tensor, name: str) -> paddle.Tensor:
                             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])
+                            fmqa_out = fmqa_out.reshape([-1, num_heads, kv_lora_rank])
+                            fmqa_out = paddle.einsum("thr,hrv->thv", fmqa_out, v_b_proj_weight)
+                            fmqa_out = fmqa_out.reshape([-1, num_heads * v_head_dim])
                             # Merge prefill and decode outputs if both are present
                             if need_do_prefill:
                                 try:
@@ -357,9 +356,8 @@ def __init__(self, fd_config: "FDConfig", **kwargs):
             self.paddleformers_config.parallel_output = self.paddleformers_config.tensor_model_parallel_size == 1
             self.paddleformers_config.max_seq_len = self.model_config.max_model_len
             self.paddleformers_config.params_dtype = self.model_config.dtype or "bfloat16"
-            # self.paddleformers_config.moe_grouped_gemm = True
             self.paddleformers_config.moe_token_dispatcher_type = "deepep"
-            # self.paddleformers_config.use_cpu_initialization = True
+
             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)
@@ -494,6 +492,9 @@ def _init_paddlefleet_parallel_state(self, fd_config) -> None:
                 current_tp_size = getattr(tp_group, "nranks", None)
                 if current_tp_size is None:
                     current_tp_size = getattr(tp_group, "world_size", None)
+            else:
+                hcg = fleet.get_hybrid_communicate_group()
+                parallel_state.initialize_model_parallel(hcg)
 
             expected_tp_size = parallel_config.tensor_parallel_size
             need_init = tp_group is None or current_tp_size != expected_tp_size
@@ -587,39 +588,41 @@ def forward(
             Returns:
                 hidden_states: [TotalTokens, HiddenDim]
             """
-            # Handle empty batch case (e.g., DP worker with no data in EP mode)
-            if getattr(forward_meta, "is_zero_size", False) or inputs["ids_remove_padding"].shape[0] == 0:
-                # Return zero tensor with correct shape: [0, hidden_size]
-                hidden_size = self.model_config.hidden_size
-                dtype = self.model_config.dtype
-                return paddle.empty([0, hidden_size], dtype=dtype)
-
-            ids_remove_padding = inputs["ids_remove_padding"]
-            num_tokens = ids_remove_padding.shape[0]
-            batch_id_per_token = forward_meta.batch_id_per_token  # [num_tokens]
-            seq_lens_decoder = forward_meta.seq_lens_decoder  # [batch_size, 1]
-
-            if batch_id_per_token is not None and seq_lens_decoder is not None:
-                decoder_offsets = seq_lens_decoder.squeeze(-1)  # [batch_size]
-                # Ensure decoder_offsets is at least 1D tensor
-                if decoder_offsets.ndim == 0:
-                    decoder_offsets = decoder_offsets.reshape([1])
-                token_decoder_offsets = paddle.index_select(
-                    decoder_offsets, batch_id_per_token, axis=0
-                )  # [num_tokens]
-
-                cu_seqlens = forward_meta.cu_seqlens_q  # [batch_size + 1]
-                if cu_seqlens is not None:
-                    token_global_idx = paddle.arange(num_tokens, dtype="int64")
-                    request_start_idx = paddle.index_select(cu_seqlens[:-1], batch_id_per_token, axis=0)
-                    relative_positions = token_global_idx - request_start_idx.astype("int64")
-                else:
-                    relative_positions = paddle.zeros([num_tokens], dtype="int64")
-                position_ids = token_decoder_offsets.astype("int64") + relative_positions
+            # In EP mode, idle ranks (is_zero_size=True) must NOT skip the full forward pass:
+            # MoE layers use alltoall collectives that require ALL EP ranks to participate,
+            # even those with zero tokens. Feed a fake token id so the model can run normally,
+            # then strip the fake output before returning.
+            is_zero_size = getattr(forward_meta, "is_zero_size", False) or inputs["ids_remove_padding"].shape[0] == 0
+            if is_zero_size:
+                ids_remove_padding = paddle.zeros([1], dtype="int64")  # fake token id
+                position_ids = paddle.zeros([1], dtype="int64")
             else:
-                position_ids = paddle.arange(num_tokens, dtype="int64")
-                if seq_lens_decoder is not None:
-                    position_ids = position_ids + seq_lens_decoder[0, 0].astype("int64")
+                ids_remove_padding = inputs["ids_remove_padding"]
+                num_tokens = ids_remove_padding.shape[0]
+                batch_id_per_token = forward_meta.batch_id_per_token  # [num_tokens]
+                seq_lens_decoder = forward_meta.seq_lens_decoder  # [batch_size, 1]
+
+                if batch_id_per_token is not None and seq_lens_decoder is not None:
+                    decoder_offsets = seq_lens_decoder.squeeze(-1)  # [batch_size]
+                    # Ensure decoder_offsets is at least 1D tensor
+                    if decoder_offsets.ndim == 0:
+                        decoder_offsets = decoder_offsets.reshape([1])
+                    token_decoder_offsets = paddle.index_select(
+                        decoder_offsets, batch_id_per_token, axis=0
+                    )  # [num_tokens]
+
+                    cu_seqlens = forward_meta.cu_seqlens_q  # [batch_size + 1]
+                    if cu_seqlens is not None:
+                        token_global_idx = paddle.arange(num_tokens, dtype="int64")
+                        request_start_idx = paddle.index_select(cu_seqlens[:-1], batch_id_per_token, axis=0)
+                        relative_positions = token_global_idx - request_start_idx.astype("int64")
+                    else:
+                        relative_positions = paddle.zeros([num_tokens], dtype="int64")
+                    position_ids = token_decoder_offsets.astype("int64") + relative_positions
+                else:
+                    position_ids = paddle.arange(num_tokens, dtype="int64")
+                    if seq_lens_decoder is not None:
+                        position_ids = position_ids + seq_lens_decoder[0, 0].astype("int64")
             forward_meta.rope_already_applied = True
             # Also set forward_meta on each TransformerLayer's config
             # so that FastDeployAttention can retrieve it from core_attn.config
@@ -658,6 +661,9 @@ def forward(
             # [b, s, h] -> [s, h] (b=1)
             hidden_states = hidden_states.squeeze(0)
 
+            # Strip fake token injected for EP idle ranks
+            if is_zero_size:
+                return hidden_states[:0]  # [0, hidden_size]
             return hidden_states
 
         @paddle.no_grad()
@@ -746,13 +752,50 @@ def patch_paddlefleet_core_attention(
             # Get configuration info
             # Prefer per-partition values (values after TP sharding),
             # because PaddleFleet's QKV output is already per-partition when TP>1
-            num_attention_heads = getattr(
-                core_attn, "num_attention_heads_per_partition", getattr(core_attn.config, "num_attention_heads", None)
+            #
+            # Detect whether this layer is a sliding-window attention (SWA) layer.
+            # SWA layers may have different kv_num_heads / head_dim than full-attention layers.
+            window_attn_skip_freq = getattr(fd_config.model_config, "window_attn_skip_freq", None)
+            is_swa_layer = (
+                window_attn_skip_freq is not None
+                and layer_number < len(window_attn_skip_freq)
+                and window_attn_skip_freq[layer_number] == 1
             )
-            num_key_value_heads = getattr(
-                core_attn,
-                "num_query_groups_per_partition",
-                getattr(core_attn.config, "num_key_value_heads", num_attention_heads),
+
+            if is_swa_layer:
+                # SWA layer: use swa_* config fields when available
+                num_attention_heads = getattr(
+                    core_attn,
+                    "num_attention_heads_per_partition",
+                    getattr(
+                        core_attn.config,
+                        "swa_num_attention_heads",
+                        getattr(core_attn.config, "num_attention_heads", None),
+                    ),
+                )
+                num_key_value_heads = getattr(
+                    core_attn,
+                    "num_query_groups_per_partition",
+                    getattr(
+                        core_attn.config,
+                        "swa_num_key_value_heads",
+                        getattr(core_attn.config, "num_key_value_heads", num_attention_heads),
+                    ),
+                )
+            else:
+                num_attention_heads = getattr(
+                    core_attn,
+                    "num_attention_heads_per_partition",
+                    getattr(core_attn.config, "num_attention_heads", None),
+                )
+                num_key_value_heads = getattr(
+                    core_attn,
+                    "num_query_groups_per_partition",
+                    getattr(core_attn.config, "num_key_value_heads", num_attention_heads),
+                )
+            logger.info(
+                f"Layer {layer_number} is_swa={is_swa_layer}: "
+                f"num_attention_heads={num_attention_heads}, num_key_value_heads={num_key_value_heads}"
             )
             hidden_size_per_attention_head = getattr(core_attn, "hidden_size_per_attention_head", None)
             if hidden_size_per_attention_head is not None:
@@ -767,10 +810,18 @@ def patch_paddlefleet_core_attention(
 
             fd_layer_id = layer_number
 
+            # Detect paddlefleet softmax_offset (a.k.a. attention sink bias).
+            # paddlefleet stores it on DotProductAttention.softmax_offset as either
+            # None (vanilla), a zeros tensor (off-by-one) or a learnable parameter.
+            # FastDeploy's Attention exposes the same math via `with_sinks` / `self.sinks`.
+            softmax_offset = getattr(core_attn, "softmax_offset", None)
+            has_sinks = softmax_offset is not None
+
             # Create Attention instance inside FastDeployAttention
             fd_attn_instance = Attention(
                 fd_config=fd_config,
                 layer_id=fd_layer_id,
+                with_sinks=has_sinks,
             )
 
             # Override Attention instance's head config to match PaddleFleet model
@@ -782,6 +833,31 @@ def patch_paddlefleet_core_attention(
                 f"Overriding Attention config: num_heads={num_attention_heads}, kv_num_heads={num_key_value_heads}, head_dim={hidden_size_per_attention_head}"
             )
 
+            # Wire paddlefleet's softmax_offset -> FastDeploy sinks. Both have
+            # shape [num_heads_per_partition] and identical softmax-off-by-one
+            # semantics: exp(qk_i) / (sum_j exp(qk_j) + exp(offset_h)).
+            if has_sinks:
+                offset_val = softmax_offset.detach()
+                if (
+                    fd_attn_instance.sinks.shape[0] != num_attention_heads
+                    or fd_attn_instance.sinks.dtype != offset_val.dtype
+                ):
+                    # Rebuild sinks parameter so shape/dtype match paddlefleet's
+                    # per-partition softmax_offset (fd_config-derived num_heads
+                    # may differ from paddlefleet's when TP topology differs).
+                    fd_attn_instance.sinks = fd_attn_instance.create_parameter(
+                        shape=[num_attention_heads],
+                        dtype=offset_val.dtype,
+                        is_bias=False,
+                        default_initializer=paddle.nn.initializer.Constant(0),
+                    )
+                fd_attn_instance.sinks.set_value(offset_val.astype(fd_attn_instance.sinks.dtype))
+                logger.info(
+                    f"Wired softmax_offset -> sinks for layer {fd_layer_id} "
+                    f"(shape={list(fd_attn_instance.sinks.shape)}, "
+                    f"dtype={fd_attn_instance.sinks.dtype})"
+                )
+
             # Create FastDeployAttention object and directly replace core_attention
             fast_deploy_core_attn = FastDeployAttention(
                 config=core_attn.config,

fastdeploy/worker/gpu_model_runner.py

@@ -1806,6 +1806,22 @@ def _get_kv_num_heads_per_layer(self) -> list[int]:
                 1,
                 int(self.model_config.num_key_value_heads) // self.parallel_config.tensor_parallel_size,
             )
+            # Check if model has SWA layers with different kv_num_heads
+            window_attn_skip_freq = getattr(self.model_config, "window_attn_skip_freq", None)
+            swa_num_key_value_heads = getattr(self.model_config, "swa_num_key_value_heads", None)
+            if window_attn_skip_freq is not None and swa_num_key_value_heads is not None:
+                swa_kv_num_heads = max(
+                    1,
+                    int(swa_num_key_value_heads) // self.parallel_config.tensor_parallel_size,
+                )
+                return [
+                    (
+                        swa_kv_num_heads
+                        if (i < len(window_attn_skip_freq) and window_attn_skip_freq[i] == 1)
+                        else kv_num_heads
+                    )
+                    for i in range(num_hidden_layers)
+                ]
             return [kv_num_heads] * num_hidden_layers
 
         if len(num_key_value_heads) != num_hidden_layers: