[v1] support GDN Ulysses cp

[XuanyuChen-SEU]

What does this PR do?

Summary

Add Ulysses-style Context Parallel (CP) support for GDN (Gated Delta Network) linear attention layers in Qwen3.5 models, including Pack (cu_seqlens) support for improved GPU utilization.

Motivation

Qwen3.5 models mix two attention types: full_attention (standard Flash Attention) and linear_attention (GDN). The existing Ulysses CP implementation only handles full_attention layers. GDN layers were silently running without CP, causing:

• Incorrect computation when cp_size > 1 (GDN layers processed only 1/cp of the sequence without all_to_all)
• Inability to scale GDN models to longer sequences via CP

Additionally, GDN layers lacked Pack (cu_seqlens) support, preventing multi-sample packing for higher GPU utilization under CP.

What This PR Does

New file: gdn_attention.py

Core implementation with 4 functions:

• is_gdn_layer() — identifies GDN layers by layer_type == "linear_attention" or block_type == "linear_attention"
• _get_gdn_module() — resolves the actual GDN module from either a Qwen3_5GatedDeltaNet or a Qwen3_5DecoderLayer
• get_parameter_local_cp() — slices model parameters (conv1d weight/bias, A_log, dt_bias) for the current CP rank, with split_sections support to ensure proportional slicing of each Q/K/V sub-group
• gdn_forward_with_cp() — replaces the GDN forward when cp_size > 1, falls back to self.original_forward when cp_size <= 1. Supports Pack mode via cu_seqlens derived from position_ids.

Modified file: sequence_parallel.py

Added GDN forward registration in apply_sequence_parallel(): when cp_size > 1, iterates over model modules, identifies GDN layers, and replaces their forward method with gdn_forward_with_cp.

Design

Reuses SeqAllToAll4D (same all_to_all primitive as UlyssesAttention). Each GDN component (Q/K/V/z/b/a) is independently reshaped to 4D and all_to_all'd with scatter heads / gather seq, then Q/K/V are concatenated for conv1d. After chunk_gated_delta_rule + norm in HP layout, the output is all_to_all'd back to CP layout.

Key decisions:

1. Per-component all_to_all — avoids the bug where uniform hidden-split on merged qkv non-proportionally distributes Q/K/V sub-groups across ranks
2. Parameter slicing — get_parameter_local_cp with split_sections ensures conv1d weight/bias, A_log, dt_bias are proportionally sliced for each CP rank
3. Zero overhead when cp_size=1 — directly falls back to self.original_forward
4. cu_seqlens from position_ids — Pack support derives cu_seqlens from position_ids via prepare_fa_kwargs_from_position_ids after all_gather, consistent with UlyssesAttention's position_ids handling
5. conv1d cu_seqlens — passes cu_seqlens to causal_conv1d_fn (FLA) to prevent cross-sample information leakage; F.conv1d fallback is disabled when cu_seqlens is not None
6. chunk_gated_delta_rule cu_seqlens — passes cu_seqlens conditionally to enable FLA Triton varlen mode; only active when batch_size == 1
7. GPU-only Pack — both causal_conv1d_fn and chunk_gated_delta_rule cu_seqlens support depend on FLA Triton kernels, which are GPU-only

Testing

Model

Qwen3.5-4B with num_hidden_layers reduced from 32 to 8 for fast single-node verification. Model weights fully downloaded.

Precision

• CP=1 (2 NPUs) vs CP=2 (4 Npus) loss alignment over 2000 steps — relative precision within 2%

• CP+Pack vs CP+non-Pack loss alignment — relative deviation < 2%

Memory

• CP=2 single-GPU memory usage lower than CP=1, confirming CP memory reduction effect

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?
  No more tests needed

[frozenleaves]

gdn适配好fla，如果还没准备好，可以先转draft