[Fix] Keep diffusion encoder TP context bindings consistent

[ch-wan]

Motivation

The LTX-2 two-GPU encoder accuracy test fails while entering use_tensor_parallel_group: diffusion initialization stamps a TP=1 attention group, but encoder folding only overrides the TP group and widths. The context then reports attn_tp_size=2 alongside an attn_tp_group of size 1 and rejects the scope before loading the encoder.

This is present on main (00986c81be687d82e66651fb3ce9a096070b06e6) and was exposed by this component accuracy CI failure.

Modifications

Override attn_tp_group, attn_tp_rank, and moe_tp_rank alongside the folded TP group. The existing context manager restores these bindings when the scope exits.

Strengthen the existing folding tests to start from diffusion's stamped TP=1 context, cover both ranks of TP=2, and check restoration after an exception in a nested folding scope.

Accuracy Tests

• On a clean checkout of current main 00986c81be687d82e66651fb3ce9a096070b06e6, the strengthened folding regressions fail in all three cases before applying the fix.
• After the fix: PYTHONPATH=python python3 -m pytest python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py -q — 9 passed.
• Pre-commit checks passed.

The full LTX-2 model-weight accuracy test has not been rerun.

The manual rerun stopped at the first Wan VAE test because the generic rerun workflow started only one process for a two-GPU test. That launcher issue is fixed separately in #40652; this run did not reach the LTX-2 encoder case.

Speed Tests and Profiling

Not applicable: this changes the encoder-loading context, with no kernel or forward-path changes.

---

CI States

Latest PR Test (Base): ❌ Run #35668042419
Latest PR Test (Extra): ❌ Run #35668042190
Latest PR Test (AMD ROCm 10): ❌ Run #35668042267

[ch-wan]

/rerun-test test_component_accuracy_2_gpu.py test_component_accuracy_parallel_runtime.py test_qwen_image21_distributed.py

[github-actions]

Results for /rerun-test test_component_accuracy_2_gpu.py test_component_accuracy_parallel_runtime.py test_qwen_image21_distributed.py:

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

python3 -m pytest python/sglang/multimodal_gen/test/single_test_file/component_accuracy/test_component_accuracy_2_gpu.py -x

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

python3 -m pytest python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py -x
python3 -m pytest python/sglang/multimodal_gen/test/unit/test_qwen_image21_distributed.py -x