Skip to content

[Fix] dspark: Avoid host reads during draft metadata glue capture - #46618

Closed
feng397 wants to merge 3 commits into
sgl-project:mainfrom
feng397:codex/fix-dspark-metadata-glue-capture
Closed

feng397 wants to merge 3 commits into
sgl-project:mainfrom
feng397:codex/fix-dspark-metadata-glue-capture

Conversation

@feng397

@feng397 feng397 commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

DeepSeek V4/4.1 DSpark draft metadata preparation falls back to seq_lens.tolist() when the batch has no CPU sequence-length mirror. This is the normal CUDA DSpark draft path. Since #45450 enables the eligible draft runner's metadata glue graph by default, its capture reaches a device-to-host read and fails:

Metadata glue-graph capture failed for key (1, '5', '5'); falling back to eager metadata prep permanently.
init_forward_metadata_out_graph -> init_forward_metadata_dspark_draft_block
seq_lens_cpu_list = seq_lens.tolist()
RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned.

The existing fallback continues preparing metadata eagerly, but permanently loses the glue optimization for that runner instance.

Modifications

  • Preserve an absent CPU mirror for CUDA sequence lengths. The CUDA causal-expansion kernel already consumes the device tensors directly; retain the existing host-length fallback for CPU inputs.
  • Derive the uniform batch size from seq_lens.shape[0], and allow the extended CPU-length mirror to remain absent through the metadata preparation interfaces.
  • Extend the existing draft metadata tests to reject host reads, verify updated sequence lengths after warm-up for batch sizes 1 and 2, and preserve CPU behavior with and without a host mirror.

Accuracy Tests

  • Full registered test entry point: 7 tests passed in the deployed environment with GPUs masked for this CPU suite. The earlier source-isolated comparison also confirmed the new regression fails on the unmodified base for batch sizes 1 and 2.
  • Real CUDA metadata capture/replay on NVIDIA B200, PyTorch 2.14.1+cu130: the saved original draft-block method reproduces the reported capture failure. The patched backend passes (batch size, block size) cases (1, 3), (2, 3), and (1, 6), with two graph replays per case after warm-up/capture. Updated causal lengths, positions, page tables, SWA indices, and SWA lengths match eager preparation. These checks use the real causal/page/SWA metadata kernels and synthetic KV mappings in a separate test process.
  • PD serving smoke: 4/4 HTTP 200 responses, each with 96 completion tokens (384 total), using DeepSeek-V4.1-Flash with DSpark, TP4, attention DP4, EP4, and flashmla_mega. After restarting with the patch, both workers are healthy and their logs contain zero metadata glue capture failures and zero tracebacks during the check.
  • Passed Ruff checks/formatting, isort, registered-test validation, package test-registration validation, and git diff --check.

Checklist

  • Format code according to the repository's Ruff/isort configuration.
  • Add regression tests in the existing registered unit-test file.
  • Update documentation (no user-facing API or configuration change).
  • Provide full model accuracy and speed benchmark results.
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #40440631650
Latest PR Test (Extra): ❌ Run #40440631535
Latest PR Test (AMD ROCm 10): ❌ Run #40440631478

@hnyls2002

Copy link
Copy Markdown
Collaborator

Thanks @feng397! You tracked down the seq_lens.tolist() host read that broke the DSpark draft metadata glue capture. #46891 landed the same fix (pass None through and take the batch size from seq_lens.shape[0]), and #46897 made the torch reference expansion work from device lengths alone, so non-CUDA devices need no host fallback. #46897 also makes a failed capture of backend-declared glue raise in CI. Closing as superseded; please reopen if I've missed a case.

@hnyls2002 hnyls2002 closed this Oct 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants