Improve long-context training for Hybrid-SWA models by avoiding A2A/AllGather on Sliding Window Attention layers

## Summary

I would like to discuss whether Megatron-Core / Transformer Engine should support an SWA-specific Context Parallel communication path for efficient long-context training of Hybrid-SWA models.

Current Context Parallel backends work well for full attention, but they can be inefficient or constrained for Sliding Window Attention (SWA) layers in modern hybrid architectures. In particular, when training long-context models with a high ratio of SWA layers, the current practical options often force SWA layers into either:

* `a2a`, which is constrained by attention head / GQA group divisibility and becomes problematic for GQA models with a small number of KV heads, or
* `all_gather`, which gathers much more K/V than SWA actually needs and therefore loses the locality benefit of sliding-window attention.

A related issue has already reported severe slowdown and timeout problems when using the `all_gather` strategy for SWA training with CP: https://github.com/NVIDIA/Megatron-LM/issues/2222

This can be a practical blocker for long context (128K~) training of recent Hybrid-SWA architectures (gpt-oss, Gemma3~, MiMo-V2-Flash, etc.), where most layers may be SWA and only a subset of layers use full attention.

This issue is also related Megatron's 2026 Q1 roadmap: https://github.com/NVIDIA/Megatron-LM/issues/4003
> SWA mixing with full attention control - Sliding window attention improvements

## Problem

For long-context training, Context Parallelism is needed to split the sequence dimension across ranks. For example, if we want to keep the per-rank local sequence length around 8K, then 256K context requires CP=32.

However, for SWA layers, the current practical backend choice is often `a2a`. This is problematic for GQA models:

```text
global sequence length: 256K
target local sequence: 8K
required CP size: 32

num_kv_heads / GQA groups: 8
A2A limit: CP <= 8
```

This means that a model with 8 KV heads cannot easily use CP=32 with A2A-based SWA. Reducing CP to 8 makes the local sequence length 32K, which can be too expensive in memory and activation footprint.

Using `all_gather` for SWA layers is also not ideal. SWA only needs a local window of K/V tokens. Gathering full-sequence K/V for every SWA layer makes communication scale with the global context length instead of the SWA window size. This is especially inefficient and reported that it is too slow: https://github.com/NVIDIA/Megatron-LM/issues/2222

## Why this matters

Recent long-context LLM architectures increasingly use Hybrid-SWA patterns:

* many layers use sliding-window attention,
* only a smaller number of layers use full attention,
* GQA is commonly used, often with a small number of KV heads.

This makes SWA-layer efficiency critical for end-to-end training performance.

For such models, optimizing only the full-attention CP path is not enough. If SWA layers are forced to use A2A or full-KV AllGather, the dominant layer type can become the scaling bottleneck.

This is particularly important for:

* long-context training at 128K / 256K and beyond,
* GQA models with small KV head counts,
* Hybrid-SWA models where most layers are SWA,
* layer-wise `cp_comm_type` configurations where full-attention layers can use `p2p` or `a2a+p2p`, but SWA layers need a different efficient communication path.

## Possible direction

One possible direction is to add an SWA-aware CP communication path, for example `swa_p2p`, that exchanges only the boundary K/V tokens required by the sliding window.

For causal SWA, each CP rank only needs a limited K/V boundary from previous sequence-neighbor ranks near the CP boundary.
If the sliding-window size is smaller than the local sequence length, this is typically just the last window_size K/V tokens from the previous rank.
If the window is larger than the local sequence length, the same logic can extend to multiple predecessor ranks.

The forward pass would communicate only the required K/V halo. (The term "halo" comes from domain decomposition / stencil computation).
The backward pass would need to send the corresponding dK/dV gradients back to the owner rank and accumulate them correctly.

Ideally, this communication path should scale with the SWA window size rather than the global sequence length:

- all_gather SWA communication: O(global_sequence_length)
- SWA boundary exchange: O(sliding_window_size)

It would also avoid the KV-head divisibility constraint of A2A, allowing SWA layers to use CP sizes larger than the number of KV heads.

For Hybrid-SWA models, this would fit naturally with layer-wise `cp_comm_type`.

For example:

- SWA layers: `swa_p2p`
- Full attention layers: `a2a+p2p` or `p2p`

This would allow full-attention layers to continue using the existing optimized CP paths, while SWA layers use a locality-aware communication path.