avoid host sync in DSpark prefill slot expansion

[inkcherry]

Motivation

torch.repeat_interleave calculates its output size from GPU data, which introduces a host synchronization during DSpark prefill slot expansion. The required token count is already available from batch.extend_lens. cc @Duyi-Wang @HaiShaw

Modifications

Pass the known token count through output_size for both state_slot and final_pos expansion, avoiding the synchronization without changing the generated values.

Tests

baseline: InferenceX run 35328216206 dsv4 pro agentic bench
baseline+ pr: internal verified @Duyi-Wang

TTFT

Concurrency	Mean	p50	p90	p95
128	4.625→4.194s (-9.32%)	3.385→3.174s (-6.23%)	8.109→6.872s (-15.25%)	12.544→11.189s (-10.81%)
192	7.253→6.633s (-8.54%)	4.594→4.357s (-5.16%)	14.056→12.218s (-13.08%)	21.217→18.691s (-11.91%)
256	11.166→9.323s (-16.50%)	6.568→5.427s (-17.37%)	22.101→18.669s (-15.53%)	33.613→26.908s (-19.95%)

TPOT

Concurrency	Mean	p50	p90	p95
128	13.11→13.09ms (-0.15%)	13.71→13.80ms (+0.66%)	16.30→16.19ms (-0.67%)	17.46→16.74ms (-4.12%)
192	15.42→14.76ms (-4.28%)	16.65→15.93ms (-4.32%)	19.09→18.15ms (-4.92%)	19.78→18.73ms (-5.31%)
256	17.03→17.58ms (+3.23%)	18.30→18.98ms (+3.72%)	22.04→22.34ms (+1.36%)	22.97→23.05ms (+0.35%)

Throughput

Concurrency	Total throughput/GPU	Output throughput/GPU	QPS
128	31,488→31,775 tok/s (+0.91%)	263.14→264.22 tok/s (+0.41%)	4.128→4.147 (+0.46%)
192	43,885→45,056 tok/s (+2.67%)	332.81→338.67 tok/s (+1.76%)	5.305→5.400 (+1.78%)
256	51,446→52,337 tok/s (+1.73%)	379.05→383.22 tok/s (+1.10%)	6.277→6.326 (+0.78%)

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 #35555910621
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ❌ Run #35555910779

[HaiShaw]

/tag-and-rerun-ci

[inkcherry]

ci finished, no nv failure. @HaiShaw @Duyi-Wang

[Duyi-Wang]

@amd-bot ci-status

[bingxche]

@Duyi-Wang

CI Status for PR #40111

Merge verdict: All 8 executed CI failures across AMD/NPU/XPU/Xeon are unrelated to this PR — the changed function is never entered by any of them (two already have in-flight fix PRs). But PR CI never exercises this PR's changed code at all: the is_unified_kv_triton() branch in DSparkWorkerV2._forward_prefill needs ROCm + SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton + --speculative-algorithm DSPARK, and the only test meeting all three is nightly-only. PR CI is also incomplete: PR Test Extra / PR Test Extra (AMD) fast-failed at the label gate, and 6 NPU jobs were fast-fail skipped.

Caution

This PR's changed code is not exercised by any PR-CI test. The edited block at dspark_worker_v2.py:646-661 is guarded by is_unified_kv_triton() (is_hip() and SGLANG_HACK_FLASHMLA_BACKEND == "unified_kv_triton", see python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py:13) inside the DSpark speculative worker. The only test that satisfies both the HIP+env gate and --speculative-algorithm DSPARK is test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py, registered suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro-dspark", nightly=True — that suite appears only in nightly-test-amd.yml, never in pr-test-amd.yml. The green NVIDIA PR Test proves nothing here either: is_hip() is False, so the branch is dead code on CUDA. The one DSpark test that did run, test/registered/unit/spec/test_dspark_target_hidden_projection.py (base-a-test-cpu, ✅), only imports the module — it never calls _forward_prefill.
Before merge, run: nightly-test-amd.yml with job_select=nightly-8-gpu-mi35x-deepseek-v4-pro-dspark (MI35x 8-GPU), or locally SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton python3 test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py.

Warning

PR CI is incomplete. PR Test Extra (35555910537) and PR Test Extra (AMD) (35555910536) both failed at call-gate / pr-gate with Missing required label 'run-ci-extra', skipping every downstream job (extra-a-test-*-amd, sgl-kernel-build-wheels, rust-ext-build, simulator-test-cpu). On NPU, multimodal-gen-test-4-npu-a3 fast-failed 6 base-b-test-*-npu-a3 jobs and the 4 base-c-*-npu-a3 jobs were skipped. On XPU, multimodal-gen-test-1-gpu-xpu was skipped after stage-a failed. Those jobs are not tested. If full signal is wanted, add run-ci-extra (via /tag-and-rerun-ci extra) and consider bypass-fastfail for the NPU leg — use it sparingly, it consumes far more CI resources.

Changed files: python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py (+7/-2) · head 32508fb · branch is 75 commits behind main (merge base ab03a8e, 2026-09-21)

Executed CI failure attribution: AMD: 2 failures (0 related) · Others: 6 failures (0 related). Plus 6 fast-fail cascade jobs, 2 label-gate failures, 4 roll-up finish jobs, and 1 matrix-cancelled job — none counted above. No jobs still pending; all 15 workflow runs on this SHA are completed.

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-c-dsv4-flash-fp4-fp8-amd-mi35x (rocm10, linux-mi35x-gpu-8)	test/registered/amd/test_deepseek_v4_flash_fp8_tbo.py	TestDeepseekV4FlashFp8Tbo.test_gsm8k_tbo	AssertionError: np.float64(0.31842304776345715) not greater than 0.91 (07:15:42)	🟢	Server launched with 'speculative_algorithm': None (log line 32098, 'enable_two_batch_overlap': True) → DSparkWorkerV2 is never constructed and _forward_prefill never runs. Known ROCm DSV4 TBO accuracy regression, see in-flight #39804.
stage-b-test-large-8-gpu-mi35x-disaggregation-amd (rocm10, linux-mi35x-gpu-8.fabric)	test/registered/amd/disaggregation/test_nixl_transfer_engine_e2e.py	TestNixlTransferEngineAccuracy.setUpClass	nixl._bindings.nixlBackendError: NIXL_ERR_BACKEND → RuntimeError: Server process exited with code -9 (09:56:25)	🟢	UCX transport selection failure (no active messages transport ... ud_verbs/rdma0:1 - Destination is unreachable) on a Llama-3.1-8B disagg server — no DSV4, no speculative decoding. Exactly the failure in-flight #40641 fixes (red since #37874).

pr-test-amd-finish (106321515161) is the roll-up of the two rows above, not an independent failure. The AMD pipeline itself was complete — stage-a → stage-b → stage-c all ran; the 7 skipped AMD jobs (sgl-kernel/jit-kernel/multimodal-gen) were excluded by check-changes path filtering, not by fast-fail.

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
multimodal-gen-test-4-npu-a3 (0)	sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py	TestDiffusionServerTwoNpu.test_diffusion_generation[qwen_image_t2i_2npu, wan2_2_t2v_14b_w8a8_2npu, minimax_h3_t2va_2npu, mova_360p_ti2va_2npu]	PerformanceValidationError: [performance] Load baseline missing or invalid	🟢	Ascend diffusion perf-baseline file missing on the runner; no LLM decode path, no DSV4/DSpark code.
multimodal-gen-test-4-npu-a3 (1)	same	...[flux_2_image_t2i_2npu, ltx_2_ti2va_2npu]	same	🟢	same
multimodal-gen-test-1-npu-a3 (0)	.../ascend/test_server_1_npu.py	TestDiffusionServerOneNpu.test_diffusion_generation[flux_image_t2i_npu, wan2_1_t2v_1.3b_1_npu, joyai_image_edit_ti2i_1npu]	same	🟢	same
multimodal-gen-test-1-npu-a3 (1)	same	...[ernie_image_t2i_1npu, glm_image_t2i_1npu, flux_2_klein_4b_t2i_1npu, z_image_t2i_1npu]	same	🟢	same
stage-a-test-1-gpu-xpu	test/registered/disaggregation/test_disaggregation_xpu.py	setUpClass	ModuleNotFoundError: No module named 'nixl' → ImportError: Please install NIXL ...	🟢	NIXL not installed in the XPU image; import-time environment gap, marked non-retriable by the CI retry wrapper.
build-test (xeon-spr, spr, ..., spr3, 2, 3)	test/registered/cpu/test_subblock_sparse_attention.py	TestSubBlockSparseAttentionModalities.test_bf16_does_not_require_sparge_attention	AttributeError: <module 'sglang.multimodal_gen.configs.pipeline_configs.minimax_h3'> does not have the attribute 'get_global_forced_attn_backend'	🟢	mock.patch target no longer exists on the installed wheel (test patches a symbol removed from minimax_h3); CPU sparse-attention test, unrelated to speculative decoding.

Cascades / gates not counted as independent failures: NPU base-b-test-{1,2,4×2,8,16}-npu-a3 all exited with Fast-fail: skipping — root cause job(s): multimodal-gen-test-4-npu-a3 (0), multimodal-gen-test-4-npu-a3 (1); Xeon build-test (..., spr1, 0, 3) was cancelled by the matrix after spr3 failed; PR Test Extra / PR Test Extra (AMD) call-gate / pr-gate failed on the missing run-ci-extra label; pr-test-npu-finish, finish (XPU), pr-test-extra-finish, pr-test-amd-extra-finish are roll-ups.

Green: PR Test (NVIDIA), Lint, PR Test (Arm64), PR Test (MLX), PR Test (MUSA), PR Test (sgl-router), PR States, Auto Label PRs. PR Test Mori (AMD) skipped by path filter.

Details / what to do before merge

The coverage gap is the only thing blocking a confident merge.

• Run the DSpark suite before merging — it is the sole test that reaches the modified lines: nightly-test-amd.yml → job_select=nightly-8-gpu-mi35x-deepseek-v4-pro-dspark. A green PR CI here is not evidence the change is correct.
• Diff review supports the claimed no-op: repeats = ctx_lens = torch.tensor(batch.extend_lens), so repeats.sum() == sum(batch.extend_lens) == num_tokens exactly — output_size matches what repeat_interleave would have computed from device data, and the returned values are unchanged. The risk if it were ever wrong is a RuntimeError: Allocated size does not match required size at prefill, not silent corruption.
• Minor, optional: int(sum(batch.extend_lens)) is already computed ~15 lines above for compute_position (dspark_worker_v2.py:631); hoisting it into one local would avoid re-summing the list.
• The branch is 75 commits behind main. Rebasing and rerunning would give a cleaner signal, and would pick up any merged fix for the two known AMD failures.
• For the two AMD reds, no new PR is needed — in-flight fixes already exist: #39804 (open since 2026-09-16, ROCm DSV4 TBO GSM8K regression from #39574, which is in this branch's merge base) and #40641 (open since 2026-09-21, UCX_NET_DEVICES loopback for NIXL_ERR_BACKEND). Chase reviewers on those rather than triaging here.
• Caveat on "pre-existing": stage-c-dsv4-flash-fp4-fp8-amd-mi35x completed in 0 of the 30 other pr-test-amd.yml runs from 2026-09-19..22 that reached it (all cancelled), so there is no directly comparable PR-CI baseline run — the "unrelated" verdict rests on the code-level evidence (speculative_algorithm: None) plus the #39804 description matching the symptom, not on a green sister run.

---

Generated by amd-bot using Claude Code CLI

[kkHuang-amd]

LGTM