Fix reduce_scatterv producer contract for SUM_LEN

[YAMY1234]

Motivation

Related to #23554.

_scatter_hidden_states() uses DP reduce_scatterv whenever should_use_dp_reduce_scatterv() is true. However, LayerCommunicator.should_use_reduce_scatter() only reported true for _scatter_hidden_states in MAX_LEN mode. In SUM_LEN prefill, producers such as dense MLP layers could still perform their TP all-reduce before the communicator performed reduce_scatterv, which double-reduced the hidden states and caused long-prompt accuracy regressions.

Kimi K2.6 exposes this because it has an early dense MLP layer (first_k_dense_replace=1), so the dense producer path must follow the same reduce-scatter contract as the communicator postprocess path.

Modifications

Make should_use_reduce_scatter() return true for _scatter_hidden_states whenever should_use_dp_reduce_scatterv() is true. This aligns the producer contract with the communicator path: if the communicator will run DP reduce_scatterv, the producer skips its internal TP all-reduce.

Accuracy Tests

GSM8K on Kimi K2.6, GB300, TP8/DP8/EP8, DP attention, num_examples=1319, num_threads=1319, num_shots=20, max_tokens=16384.

Bad baseline before this fix:

• score 0.5642802155504234, invalid extracted answers 422, empty responses 410.

This PR:

• score 0.9753656658968437, correct 1267, wrong 32, invalid extracted answers 0, empty responses 0.

Full final GSM8K log:

GSM8K Config: endpoint=http://localhost:30580; num_examples=1319; max_tokens=16384; num_threads=1319; num_shots=20; temperature=default; top_p=default; top_k=default
Running GSM8K evaluation...
/usr/local/lib/python3.12/dist-packages/torchao/quantization/quant_api.py:1731: SyntaxWarning: invalid escape sequence '\.'
  """Configuration class for applying different quantization configs to modules or parameters based on their fully qualified names (FQNs).
Downloading from https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl to /tmp/test.jsonl
/tmp/test.jsonl:   0%|          | 0.00/237k [00:00<?, ?B/s]
/tmp/test.jsonl: 732kB [00:00, 38.7MB/s]
ChatCompletionSampler initialized with self.system_message=None self.temperature=0.0 self.max_tokens=16384 self.reasoning_effort=None self.extra_body=None
  0%|          | 0/1299 [00:00<?, ?it/s]
  0%|          | 1/1299 [01:39<35:58:03, 99.76s/it]
  1%|▏         | 18/1299 [01:46<1:31:51,  4.30s/it]
  2%|▏         | 20/1299 [02:23<2:13:48,  6.28s/it]
  5%|▌         | 68/1299 [02:41<28:40,  1.40s/it]
  6%|▌         | 74/1299 [04:22<1:09:25,  3.40s/it]
  8%|▊         | 100/1299 [05:42<1:05:07,  3.26s/it]
100%|██████████| 1299/1299 [05:42<00:00,  3.79it/s]
Total latency: 343.077 s
Score: 0.975
Output throughput: 3616.264 token/s
[METRIC] gsm8k_score=0.9753656658968437 labels={"model": "moonshotai/Kimi-K2.6", "eval": "gsm8k"}
[METRIC] gsm8k_latency=343.0773218280001 labels={"model": "moonshotai/Kimi-K2.6", "eval": "gsm8k"}
Writing report to /tmp/gsm8k_moonshotai_Kimi-K2.6.html
{'score:std': np.float64(0.15500801168472017), 'score': np.float64(0.9753656658968437), 'latency': 343.0773218280001, 'output_throughput': 3616.263509897623}
Writing results to /tmp/gsm8k_moonshotai_Kimi-K2.6.json
Results saved to: /logs/accuracy/gsm8k_moonshotai_Kimi-K2.6.json
GSM8K evaluation complete

Speed Tests and Profiling

This is a correctness fix for the producer/communicator reduce-scatter contract, not a performance-oriented change. The GSM8K validation run above reported total latency 343.077 s and output throughput 3616.264 token/s.

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.

[gemini-code-assist]

Code Review

This pull request modifies the should_use_reduce_scatter method in communicator.py to permit reduce-scatter operations when should_use_dp_reduce_scatterv() is true, expanding the conditions beyond the previous requirement of maximum length padding. I have no feedback to provide.

[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!

[YAMY1234]

/tag-and-rerun-ci

[mmangkad]

Looks good

ChatCompletionSampler initialized with self.system_message=None self.temperature=1.0 self.max_tokens=512 self.reasoning_effort=None self.extra_body={'chat_template_kwargs': {'thinking': False}}
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1299/1299 [01:08<00:00, 19.00it/s]
Total latency: 69.474 s
Score: 0.965
Output throughput: 2176.574 token/s
[METRIC] gsm8k_score=0.9645881447267128 labels={"model": "moonshotai/Kimi-K2.6", "eval": "gsm8k"}
[METRIC] gsm8k_latency=69.47386905899998 labels={"model": "moonshotai/Kimi-K2.6", "eval": "gsm8k"}
Writing report to /tmp/gsm8k_moonshotai_Kimi-K2.6.html
{'score:std': np.float64(0.18481844004154702), 'score': np.float64(0.9645881447267128), 'latency': 69.47386905899998, 'output_throughput': 2176.573754249705}
Writing results to /tmp/gsm8k_moonshotai_Kimi-K2.6.json

[b8zhong]

I can also confirm this fix

python3 benchmark/gsm8k/bench_sglang.py --num-shots 20 --num-questions 1209 --parallel 1209 --platinum
Loading GSM8K Platinum dataset from HuggingFace...
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1209/1209 [00:50<00:00, 23.89it/s]
Accuracy: 0.979
Invalid: 0.000
Latency: 50.623 s
Output throughput: 2378.002 token/s