Support MiMo V2.5 with zigzag context parallelism

[Fridge003]

Summary

Adds text-only MiMo V2.5 support to CP-v2 with FA4 and zigzag prefill context parallelism.

Key changes:

• Registers MiMo V2.5 for CP-v2 and enables FA4 for its CP attention path.
• Keeps logical token lengths separate from global attention-TP padding and CP-local physical padding; CP-v2 owns logical-to-physical padding in layers/cp/padding.py.
• Skips the legacy ForwardBatch and FlashAttention padding paths while CP-v2 is active.
• Concatenates K/V into one synchronous CP all-gather and preserves logical output order plus SWA cache locations.
• Uses text warmup for --language-only workers when the model also advertises image capability.
• Adds focused padding/SWA unit coverage and a nightly TP8/CP2 GSM8K E2E test.

This PR does not add multimodal, MTP, interleave, fused-QKV-interleaving, DP-attention, communicator, or decode-CP behavior.

Validation

Validated exact commit 04a86ac7be72d4da1dcb0096ac657293754274b1, based directly on main at b8ec544946f1c5b6e17a919a691b05c5b3e7af84:

• test/registered/cp/test_cp_strategy_unit.py on GB300: 11 passed, 14 subtests passed in 11.44 s.
• Two-node run on baizhou-dev + baizhou-dev-2, with 4 GB300 GPUs per node, --tp 8, --attn-cp-size 2, FA4, zigzag prefill CP, --language-only, and dense MoE TP1.
• GSM8K 5-shot, 200 examples, 8 request threads: 0.990 (198/200; baseline 0.93), latency 1064.273 s, output throughput 40.986 token/s.
• Changed-file pre-commit, compile checks, diff checks, and CI-registry validation passed.
• SHA256 hashes for all 9 changed files matched the exact source on both validation nodes.
• Evaluator and both distributed rank logs were free of tracebacks, runtime/NCCL/watchdog/CUDA errors, and non-2xx responses. Both servers were stopped and all eight GPUs returned to 0 MiB.

The cross-node validation used NCCL_MNNVL_ENABLE=0, NCCL_LAUNCH_ORDER_IMPLICIT=1, and --disable-cuda-graph. These are validation-only settings for the allocated two-node devboxes; this PR adds no CUDA-graph workaround.

---

CI States

Latest PR Test (Base): ❌ Run #29681234251
Latest PR Test (Extra): ❌ Run #29681234158

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Fridge003]

Please consider this comment in this PR. We need to handle this padding somewhere in cp_v2
#28421 (comment)

[Fridge003]

@Fridge003 ready for review at 32b3f419ca8805390268509bf800d1e3ec7cf842.

I rebased and force-pushed on current main at e2d2e8d07ebfa019af5834b81933bd31ef6f4c4f. This revision is zigzag-only for MiMo V2.5 (--cp-strategy zigzag) and does not modify the interleave implementation.

Fresh exact-tree validation on baizhou-dev (4× GB300):

• Registered FA4/zigzag GSM8K E2E: 0.990 (198/200).
• Full three-layer MTP before reload: GSM8K 0.985, avg speculative acceptance 3.5448.
• Full same-path disk reload with CUDA-graph recapture: HTTP 200 / success=true; target + all three distinct MTP stages reloaded on every rank with the expected coordinated graph captures.
• Full three-layer MTP after reload: health 200, math answer 11, GSM8K 0.950, avg speculative acceptance 3.5416.
• Focused CP/MiMo/reload tests and the complete-range pre-commit suite passed; strict error scan was clean; all four GPUs returned to 0 MiB.

The PR title/body now reflect the zigzag implementation and exact validation evidence.

[Fridge003]

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py

[github-actions]

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py:

🚀 4-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

[Fridge003]

Please skip the padding at https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py#L1184 when cp_v2 is enabled.

We wish that CP_V2 can handle the padding on its side. Please open a new file padding.py under layer/cp, which only stores util functions for padding

[Fridge003]

Do we need to consider this place if all the padding logics are managed under cp_v2? If not, we can skip padding max_seq_len_k for cp_v2

sglang/python/sglang/srt/layers/attention/flashattention_backend.py

Lines 632 to 648 in c1067f8

	# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
	# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
	# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
	# FA3's causal reads in-bounds; widened columns index KV slot 0
	# (req_to_token is zero-init) and outputs for padding queries are
	# discarded downstream.
	if (
	self.attn_cp_size > 1
	and forward_batch.global_num_tokens_cpu is not None
	and forward_batch.extend_num_tokens is not None
	and forward_batch.extend_seq_lens_cpu is not None
	):
	padded_extend = int(forward_batch.extend_num_tokens)
	real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
	pad_delta = padded_extend - real_extend
	if pad_delta > 0:
	metadata.max_seq_len_k += pad_delta

[Fridge003]

Add a unit test for checking the correctness of padding logics

[Fridge003]

Remove the async all gather operations as before

[Fridge003]

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_mimo_cp.py

[github-actions]

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_mimo_cp.py:

🚀 4-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

🚀 8-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_mimo_cp.py

[Fridge003]

/rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_mimo_cp.py

[github-actions]

Results for /rerun-test test/registered/cp/test_gqa_prefill_cp.py test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_mimo_cp.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_prefill_cp.py

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

🚀 8-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_mimo_cp.py