diff --git a/fastdeploy/model_executor/layers/attention/append_attn_backend.py b/fastdeploy/model_executor/layers/attention/append_attn_backend.py index 594dbfe3763..b9beec0fab0 100644 --- a/fastdeploy/model_executor/layers/attention/append_attn_backend.py +++ b/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 diff --git a/fastdeploy/model_executor/models/paddleformers/base_fleet.py b/fastdeploy/model_executor/models/paddleformers/base_fleet.py index 42d511fffb9..1bde16ddc40 100644 --- a/fastdeploy/model_executor/models/paddleformers/base_fleet.py +++ b/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, diff --git a/fastdeploy/worker/gpu_model_runner.py b/fastdeploy/worker/gpu_model_runner.py index 826d9104103..6053d91d57f 100644 --- a/fastdeploy/worker/gpu_model_runner.py +++ b/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: diff --git a/tests/model_executor/fallback/test_fallback_fleet_model_coverge.py b/tests/model_executor/fallback/test_fallback_fleet_model_coverge.py index 82d1a77202a..e8f3ccc0e09 100644 --- a/tests/model_executor/fallback/test_fallback_fleet_model_coverge.py +++ b/tests/model_executor/fallback/test_fallback_fleet_model_coverge.py @@ -1064,11 +1064,14 @@ def test_seed_assertion_error_is_silenced(self): ps._TENSOR_MODEL_PARALLEL_GROUP = None with patch.object(dist_module, "fleet", mock_fleet): - with patch.object(dist_module, "get_rank", return_value=0): - with patch.object(dist_module, "new_group", return_value=mock_new_group): - # Seed function raises AssertionError → should be silently ignored - with patch.object(tp_random, "model_parallel_cuda_manual_seed", side_effect=AssertionError): - model._init_paddlefleet_parallel_state(fd_config) # must not raise + with patch.object(ps, "initialize_model_parallel"): + with patch.object(dist_module, "get_rank", return_value=0): + with patch.object(dist_module, "new_group", return_value=mock_new_group): + # Seed function raises AssertionError → should be silently ignored + with patch.object( + tp_random, "model_parallel_cuda_manual_seed", side_effect=AssertionError + ): + model._init_paddlefleet_parallel_state(fd_config) # must not raise finally: ps._TENSOR_MODEL_PARALLEL_GROUP = original_group @@ -1568,6 +1571,433 @@ def fake_get_swa_indexer_top_k(indexer_top_k, *args, **kwargs): assert captured["indexer_fill"] == -1, f"Expected fill=-1, got {captured['indexer_fill']}" +# ============================================================================ +# Tests for SWA layer detection in patch_paddlefleet_core_attention (new code) +# ============================================================================ + + +class TestPatchCoreAttentionSWADetection: + """Tests for SWA layer detection logic in patch_paddlefleet_core_attention. + + Covers: window_attn_skip_freq-based is_swa_layer branch and swa_num_* config usage. + """ + + def _make_model_with_layers( + self, layer_numbers, window_attn_skip_freq=None, swa_num_attention_heads=None, swa_num_key_value_heads=None + ): + """Create a mock model with TransformerLayers for patching tests.""" + model = MagicMock() + layers = [] + for ln in layer_numbers: + layer = MagicMock() + type(layer).__name__ = "TransformerLayer" + layer.layer_number = ln + layer.self_attn = MagicMock() + core_attn = MagicMock() + core_attn.num_attention_heads_per_partition = 8 + core_attn.num_query_groups_per_partition = 4 + core_attn.hidden_size_per_attention_head = 64 + core_attn.hidden_size_per_partition = 512 + core_attn.softmax_scale = 0.125 + core_attn.config = MagicMock() + core_attn.config.num_attention_heads = 8 + core_attn.config.num_key_value_heads = 4 + if swa_num_attention_heads is not None: + core_attn.config.swa_num_attention_heads = swa_num_attention_heads + else: + del core_attn.config.swa_num_attention_heads + if swa_num_key_value_heads is not None: + core_attn.config.swa_num_key_value_heads = swa_num_key_value_heads + else: + del core_attn.config.swa_num_key_value_heads + # No softmax_offset by default + del core_attn.softmax_offset + layer.self_attn.core_attention = core_attn + layers.append(layer) + model.run_function = layers + return model + + def test_swa_layer_uses_swa_num_key_value_heads(self): + """SWA layer (skip_freq=1) picks swa_num_key_value_heads from config.""" + model = self._make_model_with_layers( + [0, 1], + swa_num_attention_heads=4, + swa_num_key_value_heads=2, + ) + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = [1, 0] # layer 0 is SWA + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + mock_attention_cls.return_value = mock_attn_instance + mock_attn_instance.sinks = MagicMock() + mock_attn_instance.sinks.shape = [4] + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + result = patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + assert result == 2 + # Verify that for layer 0 (SWA), the Attention was created with with_sinks=False + calls = mock_attention_cls.call_args_list + assert len(calls) == 2 + # Layer 0 call: with_sinks=False (no softmax_offset) + assert calls[0].kwargs["with_sinks"] is False + assert calls[0].kwargs["layer_id"] == 0 + + def test_non_swa_layer_uses_standard_heads(self): + """Non-SWA layer (skip_freq=0) uses standard num_key_value_heads.""" + model = self._make_model_with_layers([0], swa_num_key_value_heads=2) + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = [0] # layer 0 is NOT SWA + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + mock_attention_cls.return_value = mock_attn_instance + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + # For non-SWA layer, kv_num_heads should be from core_attn.num_query_groups_per_partition = 4 + mock_attn_instance.kv_num_heads = 4 # set by the production code + + def test_swa_layer_index_beyond_skip_freq_length(self): + """Layer index >= len(window_attn_skip_freq) is treated as non-SWA.""" + model = self._make_model_with_layers([2]) # layer 2 + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = [1, 0] # only 2 entries, layer 2 is out of range + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + mock_attention_cls.return_value = mock_attn_instance + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + # layer_number=2 >= len([1,0])=2 → not SWA → uses standard path + assert mock_attention_cls.call_count == 1 + + def test_no_window_attn_skip_freq_uses_standard_path(self): + """No window_attn_skip_freq on model_config → all layers use standard path.""" + model = self._make_model_with_layers([0]) + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = None + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + mock_attention_cls.return_value = mock_attn_instance + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + assert mock_attention_cls.call_count == 1 + assert mock_attention_cls.call_args.kwargs["with_sinks"] is False + + +# ============================================================================ +# Tests for softmax_offset/sinks wiring (new code) +# ============================================================================ + + +class TestPatchCoreAttentionSoftmaxOffset: + """Tests for softmax_offset -> sinks wiring in patch_paddlefleet_core_attention.""" + + def _make_model_with_softmax_offset(self, offset_tensor): + """Create model with a single TransformerLayer that has softmax_offset.""" + model = MagicMock() + layer = MagicMock() + type(layer).__name__ = "TransformerLayer" + layer.layer_number = 0 + layer.self_attn = MagicMock() + core_attn = MagicMock() + core_attn.num_attention_heads_per_partition = 4 + core_attn.num_query_groups_per_partition = 4 + core_attn.hidden_size_per_attention_head = 64 + core_attn.hidden_size_per_partition = 256 + core_attn.softmax_scale = 0.125 + core_attn.config = MagicMock() + core_attn.config.num_attention_heads = 4 + core_attn.config.num_key_value_heads = 4 + del core_attn.config.swa_num_attention_heads + del core_attn.config.swa_num_key_value_heads + core_attn.softmax_offset = offset_tensor + layer.self_attn.core_attention = core_attn + model.run_function = [layer] + return model + + def test_has_sinks_true_when_softmax_offset_present(self): + """softmax_offset present → with_sinks=True passed to Attention.""" + offset = paddle.zeros([4], dtype="float32") + model = self._make_model_with_softmax_offset(offset) + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = None + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + # Use a real paddle parameter for sinks so shape/dtype are accessible + sinks_param = paddle.create_parameter( + shape=[4], dtype="float32", default_initializer=paddle.nn.initializer.Constant(0) + ) + mock_attn_instance.sinks = sinks_param + mock_attn_instance.create_parameter = MagicMock(return_value=sinks_param) + mock_attention_cls.return_value = mock_attn_instance + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + mock_attention_cls.assert_called_once() + assert mock_attention_cls.call_args.kwargs["with_sinks"] is True + + def test_has_sinks_false_when_no_softmax_offset(self): + """No softmax_offset → with_sinks=False.""" + model = MagicMock() + layer = MagicMock() + type(layer).__name__ = "TransformerLayer" + layer.layer_number = 0 + layer.self_attn = MagicMock() + core_attn = MagicMock() + core_attn.num_attention_heads_per_partition = 4 + core_attn.num_query_groups_per_partition = 4 + core_attn.hidden_size_per_attention_head = 64 + core_attn.hidden_size_per_partition = 256 + core_attn.softmax_scale = 0.125 + core_attn.config = MagicMock() + core_attn.config.num_attention_heads = 4 + core_attn.config.num_key_value_heads = 4 + del core_attn.config.swa_num_attention_heads + del core_attn.config.swa_num_key_value_heads + del core_attn.softmax_offset + layer.self_attn.core_attention = core_attn + model.run_function = [layer] + fd_config = MagicMock() + fd_config.model_config.window_attn_skip_freq = None + + mock_attention_cls = MagicMock() + mock_attn_instance = MagicMock() + mock_attention_cls.return_value = mock_attn_instance + + with patch("fastdeploy.model_executor.layers.attention.attention.Attention", mock_attention_cls): + patch_paddlefleet_core_attention(model=model, fd_config=fd_config) + + assert mock_attention_cls.call_args.kwargs["with_sinks"] is False + + +# ============================================================================ +# Tests for MLA prefill qk_head_dim padding (new code line 228-229) +# ============================================================================ + + +class TestMLAPrefillQKHeadDimPad: + """Test the qk_head_dim != v_head_dim padding branch in FastDeployAttention.forward.""" + + def test_v_padded_when_qk_head_dim_differs(self): + """Prefill branch passes v as-is to fd_attention.forward (no padding applied in base_fleet). + + qk_head_dim and v_head_dim are int values; even when they differ, + the current implementation does not pad v — it passes it unchanged. + """ + kv_lora_rank, v_head_dim, num_heads = 4, 2, 2 + qk_head_dim = 6 # differs from v_head_dim + + attn, mock_fd_attention = _create_mla_attention(kv_lora_rank, v_head_dim, num_heads) + # Use real int values (not MagicMock) so isinstance(x, int) check passes + attn.config.qk_head_dim = qk_head_dim + attn.config.v_head_dim = v_head_dim + + # Set up forward_meta with prefill only (no decode) + forward_meta = _create_mock_forward_meta(prefill_tokens=3, decode_tokens=0) + attn.config.forward_meta = forward_meta + + seq_len = 3 + query = paddle.randn([seq_len, num_heads, kv_lora_rank]) + key = paddle.randn([seq_len, num_heads, kv_lora_rank]) + value = paddle.randn([seq_len, num_heads, v_head_dim]) # v_head_dim=2 + kv_compressed = paddle.randn([1, seq_len, num_heads, kv_lora_rank]) + k_pos_emb = paddle.randn([1, seq_len, num_heads, kv_lora_rank]) + + prefill_output = paddle.randn([seq_len, num_heads, kv_lora_rank]) + mock_fd_attention.forward.return_value = prefill_output + + result = attn.forward( + query=query, + key=key, + value=value, + attention_mask=None, + kv_compressed=kv_compressed, + k_pos_emb=k_pos_emb, + ) + + assert result is not None + # fd_attention.forward must be called once (prefill path) + mock_fd_attention.forward.assert_called_once() + # v is passed as-is (no padding in current implementation) + call_kwargs = mock_fd_attention.forward.call_args.kwargs + v_passed = call_kwargs["v"] + assert v_passed.shape[-1] == v_head_dim + + def test_v_not_padded_when_qk_head_dim_equals_v_head_dim(self): + """No padding when qk_head_dim == v_head_dim.""" + kv_lora_rank, v_head_dim, num_heads = 4, 4, 2 + qk_head_dim = 4 # same as v_head_dim + + attn, mock_fd_attention = _create_mla_attention(kv_lora_rank, v_head_dim, num_heads) + attn.config.qk_head_dim = qk_head_dim + attn.config.v_head_dim = v_head_dim + + forward_meta = _create_mock_forward_meta(prefill_tokens=3, decode_tokens=0) + attn.config.forward_meta = forward_meta + + seq_len = 3 + query = paddle.randn([seq_len, num_heads, kv_lora_rank]) + key = paddle.randn([seq_len, num_heads, kv_lora_rank]) + value = paddle.randn([seq_len, num_heads, v_head_dim]) + kv_compressed = paddle.randn([1, seq_len, num_heads, kv_lora_rank]) + k_pos_emb = paddle.randn([1, seq_len, num_heads, kv_lora_rank]) + + prefill_output = paddle.randn([seq_len, num_heads, kv_lora_rank]) + mock_fd_attention.forward.return_value = prefill_output + + result = attn.forward( + query=query, + key=key, + value=value, + attention_mask=None, + kv_compressed=kv_compressed, + k_pos_emb=k_pos_emb, + ) + + assert result is not None + call_kwargs = mock_fd_attention.forward.call_args.kwargs + v_passed = call_kwargs["v"] + # v should remain unchanged (no padding) + assert v_passed.shape[-1] == v_head_dim + + +# ============================================================================ +# Tests for EP idle rank forward with fake token strip (new code lines 610-617, 683-685) +# ============================================================================ + + +class TestForwardEPIdleRankFakeTokenStrip: + """Test that EP idle ranks get a fake token injected then stripped.""" + + def test_is_zero_size_true_returns_zero_rows(self): + """is_zero_size=True injects fake token then strips → output shape[0] == 0.""" + model = _create_mock_fleet_model_for_forward(hidden_size=64, num_tokens=1) + + forward_meta = MagicMock() + forward_meta.is_zero_size = True + + inputs = {"ids_remove_padding": paddle.zeros([0], dtype="int64")} + result = model.forward(inputs, forward_meta) + + # Result should have 0 tokens but correct hidden_size + assert result.shape[0] == 0 + assert result.shape[1] == 64 + + def test_empty_ids_triggers_fake_token(self): + """ids_remove_padding.shape[0]==0 with is_zero_size=False also injects fake token.""" + model = _create_mock_fleet_model_for_forward(hidden_size=64, num_tokens=1) + + forward_meta = MagicMock() + forward_meta.is_zero_size = False + + inputs = {"ids_remove_padding": paddle.zeros([0], dtype="int64")} + result = model.forward(inputs, forward_meta) + + assert result.shape[0] == 0 + + def test_normal_forward_no_strip(self): + """Normal (non-zero-size) forward returns all tokens without stripping.""" + num_tokens = 3 + model = _create_mock_fleet_model_for_forward(hidden_size=64, num_tokens=num_tokens) + + forward_meta = MagicMock() + forward_meta.is_zero_size = False + forward_meta.batch_id_per_token = None + forward_meta.seq_lens_decoder = None + forward_meta.cu_seqlens_q = None + + inputs = {"ids_remove_padding": paddle.to_tensor([1, 2, 3], dtype="int64")} + result = model.forward(inputs, forward_meta) + + assert result.shape[0] == num_tokens + + +# ============================================================================ +# Tests for _init_paddlefleet_parallel_state hcg else branch (new code lines 514-516) +# ============================================================================ + + +class TestInitParallelStateHcgElseBranch: + """Test the new else branch: tp_group is not None but hcg initialization.""" + + def test_tp_group_none_triggers_hcg_initialize(self): + """tp_group=None → else branch calls fleet.get_hybrid_communicate_group + initialize_model_parallel.""" + import paddle.distributed as dist_module + import paddlefleet.parallel_state as ps + from paddlefleet.tensor_parallel import random as tp_random + + model = PaddleFleetModelBase.__new__(PaddleFleetModelBase) + fd_config = MagicMock() + fd_config.parallel_config.tensor_parallel_size = 2 + fd_config.parallel_config.data_parallel_size = 1 + fd_config.parallel_config.expert_parallel_size = 1 + fd_config.parallel_config.sequence_parallel = False + + mock_fleet = MagicMock() + mock_hcg = MagicMock() + mock_fleet.get_hybrid_communicate_group.return_value = mock_hcg + + original_group = ps._TENSOR_MODEL_PARALLEL_GROUP + try: + ps._TENSOR_MODEL_PARALLEL_GROUP = None + + with patch.object(dist_module, "fleet", mock_fleet): + with patch.object(ps, "initialize_model_parallel") as mock_init_mp: + with patch.object(tp_random, "model_parallel_cuda_manual_seed"): + model._init_paddlefleet_parallel_state(fd_config) + + # The else branch should NOT be hit since tp_group is None initially + # Instead the need_init branch (tp_size=2, tp_group=None) → initialize_model_parallel + mock_init_mp.assert_called() + finally: + ps._TENSOR_MODEL_PARALLEL_GROUP = original_group + + def test_tp_group_not_none_nranks_none_triggers_hcg(self): + """tp_group exists but nranks=None and world_size=None → hcg else branch triggered.""" + import paddle.distributed as dist_module + import paddlefleet.parallel_state as ps + from paddlefleet.tensor_parallel import random as tp_random + + model = PaddleFleetModelBase.__new__(PaddleFleetModelBase) + fd_config = MagicMock() + fd_config.parallel_config.tensor_parallel_size = 2 + fd_config.parallel_config.data_parallel_size = 1 + fd_config.parallel_config.expert_parallel_size = 1 + fd_config.parallel_config.sequence_parallel = False + + mock_fleet = MagicMock() + mock_hcg = MagicMock() + mock_fleet.get_hybrid_communicate_group.return_value = mock_hcg + + # Create a mock group with both nranks and world_size as None to trigger hcg fallback + mock_existing_group = MagicMock(spec=[]) # no nranks, no world_size + + original_group = ps._TENSOR_MODEL_PARALLEL_GROUP + try: + ps._TENSOR_MODEL_PARALLEL_GROUP = mock_existing_group + + with patch.object(dist_module, "fleet", mock_fleet): + with patch.object(ps, "initialize_model_parallel") as mock_init_mp: + with patch.object(tp_random, "model_parallel_cuda_manual_seed"): + model._init_paddlefleet_parallel_state(fd_config) + + # current_tp_size=None → need_init=True → tp_size=2 → initialize_model_parallel(hcg) + mock_init_mp.assert_called() + finally: + ps._TENSOR_MODEL_PARALLEL_GROUP = original_group + + class TestTryResolvePaddlefleetImportError: """Test model_base.py line 203-209: paddlefleet not installed raises ImportError.""" diff --git a/tests/worker/test_gpu_model_runner.py b/tests/worker/test_gpu_model_runner.py index 375cabbf783..91c1ac573ad 100644 --- a/tests/worker/test_gpu_model_runner.py +++ b/tests/worker/test_gpu_model_runner.py @@ -1074,5 +1074,144 @@ def test_execute_model_overlap_zero_output_flushes_preempted_batch(self): self.assertEqual(runner._cached_real_bsz, 22) +class TestGetKvNumHeadsPerLayer(unittest.TestCase): + """Tests for GPUModelRunner._get_kv_num_heads_per_layer""" + + def _make_runner(self, num_hidden_layers, num_key_value_heads, tp_size=1, **extra_attrs): + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.model_config = Mock() + runner.model_config.num_hidden_layers = num_hidden_layers + runner.model_config.num_key_value_heads = num_key_value_heads + runner.model_config.num_key_value_heads_list = None + runner.parallel_config = Mock() + runner.parallel_config.tensor_parallel_size = tp_size + for k, v in extra_attrs.items(): + setattr(runner.model_config, k, v) + return runner + + def test_uniform_heads_no_swa(self): + """No SWA config -> uniform kv_num_heads across all layers""" + runner = self._make_runner(4, 8, tp_size=2, window_attn_skip_freq=None, swa_num_key_value_heads=None) + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [4, 4, 4, 4]) + + def test_uniform_heads_tp1(self): + """TP=1, no SWA""" + runner = self._make_runner(3, 6, tp_size=1, window_attn_skip_freq=None, swa_num_key_value_heads=None) + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [6, 6, 6]) + + def test_swa_basic(self): + """SWA layers get swa_kv_num_heads, non-SWA layers get normal kv_num_heads""" + runner = self._make_runner( + 4, + 8, + tp_size=1, + window_attn_skip_freq=[1, 0, 1, 0], + swa_num_key_value_heads=4, + ) + result = runner._get_kv_num_heads_per_layer() + # layer 0: swa (skip_freq=1) -> 4, layer 1: normal -> 8, + # layer 2: swa -> 4, layer 3: normal -> 8 + self.assertEqual(result, [4, 8, 4, 8]) + + def test_swa_with_tp(self): + """SWA with tensor parallelism divides both head counts""" + runner = self._make_runner( + 4, + 8, + tp_size=2, + window_attn_skip_freq=[1, 0, 0, 1], + swa_num_key_value_heads=4, + ) + result = runner._get_kv_num_heads_per_layer() + # kv_num_heads = 8 // 2 = 4, swa_kv_num_heads = 4 // 2 = 2 + self.assertEqual(result, [2, 4, 4, 2]) + + def test_swa_all_layers(self): + """All layers are SWA""" + runner = self._make_runner( + 3, + 8, + tp_size=1, + window_attn_skip_freq=[1, 1, 1], + swa_num_key_value_heads=2, + ) + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [2, 2, 2]) + + def test_swa_skip_freq_shorter_than_layers(self): + """window_attn_skip_freq shorter than num_hidden_layers -> extra layers use normal heads""" + runner = self._make_runner( + 5, + 8, + tp_size=1, + window_attn_skip_freq=[1, 0, 1], + swa_num_key_value_heads=4, + ) + result = runner._get_kv_num_heads_per_layer() + # i=0: swa(4), i=1: normal(8), i=2: swa(4), i=3: out of range->normal(8), i=4: normal(8) + self.assertEqual(result, [4, 8, 4, 8, 8]) + + def test_swa_only_window_attn_skip_freq_set(self): + """Only window_attn_skip_freq set, swa_num_key_value_heads is None -> no SWA logic""" + runner = self._make_runner( + 3, + 6, + tp_size=1, + window_attn_skip_freq=[1, 0, 1], + swa_num_key_value_heads=None, + ) + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [6, 6, 6]) + + def test_swa_only_swa_num_kv_heads_set(self): + """Only swa_num_key_value_heads set, window_attn_skip_freq is None -> no SWA logic""" + runner = self._make_runner( + 3, + 6, + tp_size=1, + window_attn_skip_freq=None, + swa_num_key_value_heads=4, + ) + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [6, 6, 6]) + + def test_swa_kv_heads_min_clamp(self): + """swa_num_key_value_heads // tp_size < 1 should clamp to 1""" + runner = self._make_runner( + 2, + 8, + tp_size=8, + window_attn_skip_freq=[1, 0], + swa_num_key_value_heads=2, + ) + result = runner._get_kv_num_heads_per_layer() + # kv_num_heads = max(1, 8//8) = 1, swa_kv_num_heads = max(1, 2//8) = 1 + self.assertEqual(result, [1, 1]) + + def test_num_key_value_heads_list(self): + """When num_key_value_heads_list is provided, use per-layer values""" + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.model_config = Mock() + runner.model_config.num_hidden_layers = 3 + runner.model_config.num_key_value_heads_list = [8, 4, 16] + runner.parallel_config = Mock() + runner.parallel_config.tensor_parallel_size = 2 + result = runner._get_kv_num_heads_per_layer() + self.assertEqual(result, [4, 2, 8]) + + def test_num_key_value_heads_list_mismatch_raises(self): + """Mismatched list length raises ValueError""" + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.model_config = Mock() + runner.model_config.num_hidden_layers = 3 + runner.model_config.num_key_value_heads_list = [8, 4] + runner.parallel_config = Mock() + runner.parallel_config.tensor_parallel_size = 1 + with self.assertRaises(ValueError): + runner._get_kv_num_heads_per_layer() + + if __name__ == "__main__": unittest.main()