feat: hybrid CP for Qwen3.5 MoE (DeltaNet + attention)

[samsja]

Summary

Adds context parallelism (CP) support for Qwen3.5 hybrid models that mix GatedDeltaNet (linear attention) with standard softmax attention.

The challenge: CP works via ring attention for standard attention layers, but DeltaNet is recurrent — it maintains a compact state matrix S ∈ R^{K×V} that's updated sequentially. Ring attention doesn't apply.

Approach

We use different CP strategies per layer type:

• Standard attention layers (25% of layers): ring attention, same as before
• DeltaNet layers (75% of layers): fla library's native CP, which communicates only the recurrent state between ranks

This follows the LASP-2H approach from the literature (paper) and uses fla's merged PR #691 which implements CP for chunk_gated_delta_rule.

Communication cost per DeltaNet layer

Approach	Volume per layer
All-gather full sequence	O(N × d) — e.g. 256 MB at seq_len=32K, d=4096
fla native state passing	O(H × K × (K+V)) — e.g. 4 MB for 64 heads × 128 × 256

~64x reduction in communication volume for DeltaNet layers.

Changes

• flash-linear-attention added as dependency — provides Triton kernels for GatedDeltaNet with CP support
• modeling_qwen3_5_moe.py: Qwen3_5MoeGatedDeltaNet now builds FLACPContext and passes it to fla's chunk_gated_delta_rule when CP is enabled
• cp.py: setup_hybrid_cp() configures DeltaNet modules with CP group/rank info. Also adds cp_gather/cp_scatter autograd primitives as general-purpose fallback utilities
• rl/train.py and sft/train.py: call setup_hybrid_cp() after ring attention setup when CP is enabled
• model.py: fix debug.num_layers for VLM configs (nested text_config)
• Test configs: rl_cp_test.toml (Qwen3.5-35B-A3B, CP=2, 1 node, fake data) for quick validation

Test plan

• Dry run validates for both RL and SFT configs
• Qwen3.5-35B-A3B with CP=2, fake data, random init — 5 steps on 1 node (8 GPUs)
• Qwen3.5-397B-A17B with CP=4, EP=8, LoRA, fake data — forward/backward passes on 4 nodes (32 GPUs)
• Correctness validation: compare CP vs non-CP outputs on small model

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches distributed context-parallel training paths and introduces a new CUDA/Triton dependency (flash-linear-attention), which could impact correctness/performance across CP setups and hardware environments.

Overview
Enables hybrid context parallelism for Qwen3.5 MoE hybrid models by using ring-attention CP for standard attention layers while using flash-linear-attention’s native CP/state-passing for GatedDeltaNet (linear attention) layers.

This wires CP group/rank info into DeltaNet modules via new setup_hybrid_cp() (called from both rl/train.py and sft/train.py) and updates the DeltaNet fast path to build/pass an FLACPContext into chunk_gated_delta_rule when available, falling back to the existing PyTorch kernel otherwise.

Also tightens VLM training constraints (requires bf16 optimization/reduce dtypes) and fixes debug.num_layers handling for VLM configs by applying it to nested text_config.

^{Written by Cursor Bugbot for commit fe4bdb5. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

setup_hybrid_cp uses model before it's defined in SFT

High Severity

setup_hybrid_cp(model, ...) is called at line 108 inside the if parallel_dims.cp_enabled block, but model is not created until line 124 via setup_model(...). This will raise a NameError at runtime whenever context parallelism is enabled for SFT training. The equivalent code in rl/train.py correctly places setup_hybrid_cp after model creation.

Additional Locations (1)

• src/prime_rl/trainer/sft/train.py#L123-L126