[yanbing-j]

Guard the DSA/MLA prefill context-parallel reduce-scatter checks so they only run for context-parallel extend batches. These predicates are irrelevant for decode, but were still traced during CPU torch.compile decode graph capture and regressed Qwen3.6-35B-A3B decode throughput.

This keeps decode graph capture on the hot path free of prefill-only CP runtime checks while preserving the existing reduce-scatter behavior for context-parallel prefill.

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #35186798049
Latest PR Test (Extra): ❌ Run #35186797885
Latest PR Test (AMD ROCm 10): ❌ Run #35186797991