🚀[FEA]: Update FLARE++ to the v2 reference design and add token-sharded context parallelism for FLARE/FLARE++ #2024

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Medium

Please provide a clear description of problem you would like to solve.

cc @coreyjadams @mnabian

Following up on #1150 (FLARE), I wanted to share the updated version of FLARE++.

• Paper (v2): FLARE++: Low-rank attention with attention-synthesized routing
• Blog post: FLARE++: Stop compressing every field the same way
• Code: https://github.com/vpuri3/FLARE.py

It's great to see FLARE++ already landed in #1940 (physicsnemo.nn.FLAREPlusPlus, physicsnemo.models.flare.FLAREPlusPlus, and GALE_FPP in GeoTransolver). That implementation follows the v1 preprint, so I want to clarify how v1 and v2 relate.

v1 vs. v2. v1 was preliminary work that we posted so there would be something public before my PhD defense. Because it was premature, we didn't release any code with it. v2 (our ICLR submission) is the polished version and has several improvements. What v1 calls FLARE++ is Simplified FLARE++ in v2 (Appendix A, "Simplified FLARE++"). The model v2 calls FLARE++ is the updated mixer described below.

Two things in v2 should be of interest to the team:

1. The FLARE++ token mixer as a replacement for FLARE

Reference: FLAREPPMixer in pdebench/models/mixer_backbone.py. Appendix A of v2 covers the design in full.

Changes from v1 (= the #1940 implementation) to v2:

• Normalizing the synthesis call: head-wise normalization of the seeds and the synthesis keys/values (Appendix A, "Normalizing the synthesis call").
• Fixed reference and gate: the routing query is a learned fixed query plus a per-head sigmoid-gated synthesized term, q = q_fixed + σ(g) · q_dynamic (Appendix A, "Fixed reference and gate"). With the gate closed, the mixer reduces to FLARE.
• Near-identity feature projections: residual-MLP key/value projections initialized near identity (Appendix A, "Near-identity feature projections").
• Three input projections instead of four: the synthesis keys and values are shared (Appendix A, "Sharing synthesis keys and values"). This saves one $C \to C$ projection per block without hurting accuracy in our experiments.

The normalization and initialization choices in Appendix A drastically improve early-training stability and help FLARE++ converge better. Appendix A also explains why synthesized routing queries are hard to optimize without them ("Optimization difficulty of synthesized routing queries"). "Operations in one block" in Appendix A lists the per-block operation sequence of FLARE, FLARE++, and Simplified FLARE++ side by side.

Two results from Appendix D support this:

• Smaller seed variance (Appendix D.1): over ten seeds on Elasticity, FLARE++ has a smaller spread than FLARE. Its coefficient of variation is about 5% (0.37–0.43% test error), against about 11% for FLARE (0.58–0.86%).
• Lower mixed-precision sensitivity (Appendix D.2): FLARE++ is the least precision-sensitive mixer in the sweep. Its FP16 error stays within 8% of its FP32 error at every depth.

2. Multi-GPU context parallelism for FLARE and FLARE++

Reference: pdebench/distributed/flare_cp.py. The model calls it through flare_encode in pdebench/models/flarepp.py. Appendix B of v2 has the complete derivation.

Point tokens are sharded across ranks, and no rank ever gathers the full sequence. The distributed encoder is exact: each rank runs a fused local SDPA against its shard. The rank then contributes only latent-sized outputs and log-sum-exp statistics to a collective, which rebuilds the globally normalized attention (FlareEncoderCPFlash / FlareFlashCPFunction, with a custom backward). Decoding stays local. FLARE calls this encoder once per block. FLARE++ calls it twice, once for query synthesis and once for the physical-value gather. Communication therefore depends on $M$ and the channel width, not on $N$.

Appendix B covers:

• the token-sharded representation
• the exact distributed encoder
• the numerically stable reductions
• the FLARE and FLARE++ communication schedules
• the backward pass
• measured scaling

#1940 added token-sharded FLARE++ and then reverted it, and it currently rejects context parallelism explicitly. This encoder could fill that gap for both FLARE and FLAREPlusPlus.

I'm happy to help with a PR for either piece, or to review.

Describe any alternatives you have considered

No response