[DeepSeek v3.2][Bugfix] get_index_k_scale_buffer support cp

[xu-yfei]

Motivation

Fix PR #16043 #19319 Prefill CP scenario bug. In the DeepSeek V3.2 prefill with CP (context parallel) scenario, forward_batch.seq_lens and forward_batch.seq_lens_cpu are global information. A request may have fewer than 8 tokens in one chunk, so it will not be assigned to a specific CP rank. In this case, use get_indexer_seq_lens_cpu and get_indexer_seq_len instead.

Modifications

Accuracy Tests

# cp8
nemo-run_1/0 ---------------------------------------- aime25 ----------------------------------------
nemo-run_1/0 evaluation_mode  | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer
nemo-run_1/0 pass@1[avg-of-4] | 30          | 14994      | 3793        | 92.50% ± 1.67%   | 0.83%    
nemo-run_1/0 majority@4       | 30          | 14994      | 3793        | 93.33%           | 0.00%    
nemo-run_1/0 pass@4           | 30          | 14994      | 3793        | 100.00%          | 0.00%    

nemo-run_1/0 ---------------------------------------- aime25 ----------------------------------------
nemo-run_1/0 evaluation_mode  | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer
nemo-run_1/0 pass@1[avg-of-4] | 30          | 14303      | 3464        | 91.67% ± 1.92%   | 0.00%    
nemo-run_1/0 majority@4       | 30          | 14303      | 3464        | 94.17%           | 0.00%    
nemo-run_1/0 pass@4           | 30          | 14303      | 3464        | 96.67%           | 0.00%  

# dp8 updated
nemo-run_1/0 ---------------------------------------- aime25 ----------------------------------------
nemo-run_1/0 evaluation_mode  | num_entries | avg_tokens | gen_seconds | symbolic_correct | no_answer
nemo-run_1/0 pass@1[avg-of-4] | 30          | 14965      | 2013        | 90.00% ± 2.72%   | 1.67%    
nemo-run_1/0 majority@4       | 30          | 14965      | 2013        | 90.00%           | 0.00%    
nemo-run_1/0 pass@4           | 30          | 14965      | 2013        | 96.67%           | 0.00%  

Benchmarking 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 Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

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

[xu-yfei]

@Fridge003 Could you please review this PR?

[Fridge003]

@xu-yfei What's your launching command for accuracy test?

[Fridge003]

/tag-and-rerun-ci

[xu-yfei]

@xu-yfei What's your launching command for accuracy test?
All models are DeepSeek V3.2, and AIME is run following the documentation.

# cp8
MODEL_PATH=/home/models/DeepSeek-V3.2/

python3 -m sglang.launch_server --model-path $MODEL_PATH --trust-remote-code \
--port 8000 --host 0.0.0.0 --attention-backend  nsa \
--enable-metrics --mem-fraction-static 0.85 --max-running-requests 64  --enable-cache-report --page-size 64 \
--tp-size 8 \
--tool-call-parser deepseekv32 \
--reasoning-parser deepseek-v3 \
--chunked-prefill-size 16384 \
--nsa-decode-backend fa3 \
--enable-nsa-prefill-context-parallel \
--nsa-prefill-cp-mode round-robin-split \
--disable-radix-cache 

# dp8
MODEL_PATH=/home/models/DeepSeek-V3.2/

nohup python3 -m sglang.launch_server --model-path $MODEL_PATH --trust-remote-code \
--port 8000 --host 0.0.0.0 --attention-backend  nsa \
--enable-metrics --mem-fraction-static 0.9 --dp-size 8 --enable-dp-attention --enable-cache-report --page-size 64 \
--tp-size 8  \
--tool-call-parser deepseekv32 \
--reasoning-parser deepseek-v3 \
--chunked-prefill-size 16384 \
--nsa-decode-backend fa3 \
--max-running-requests 64 \
--speculative-algorithm EAGLE \
--speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4

# aime
export NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK=1

ns prepare_data aime25

ip=127.0.0.1
PORT=8000
BACKEND=sglang
MODEL="/home/models/DeepSeek-V3.2" # Should be changed to the model name
MODEL_NAME="dsv32-fp8"

echo "Starting AIME25 evaluation with model $MODEL on port $PORT using backend $BACKEND..."
ns eval \
  --benchmarks=aime25:4 \
  --server_type=$BACKEND \
  --model=$MODEL \
  --server_address=http://${ip}:${PORT}/v1 \
  --output_dir=nemo_skills_aime25_${MODEL_NAME}_output_${BACKEND}_$(date +%Y%m%d_%H%M%S) \
  ++chat_template_kwargs.thinking=true \
  ++inference.temperature=1.0 \
  ++inference.top_p=0.95 \
  ++inference.tokens_to_generate=64000
  # ++inference.tokens_to_generate=120000 for Speciale model

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

/rerun-failed-ci

[yuan-luo]

/rerun-failed-ci