dspark perf optimization

[ZhangLirong-amd]

What

Remove two DSpark-only per-decode-step host syncs from the spec-decode hot path (both gated on use_dspark / the ragged path; plain V4 never hits them).

before

after

1. eagle.py prepare_inputs — int(cu_seqlens_q[-1]) read a GPU scalar back to the host every decode step (aten::item → _local_scalar_dense → hipMemcpyWithStream). The copy is stream-ordered after the whole decode kernel, so the CPU blocked until the GPU queue drained, collapsing CPU/GPU overlap into a ping-pong. Replaced with an on-device clamp (torch.minimum against a 0-dim device scalar); clamp_(min=0) on the bound preserves the old total_tokens > 0 guard.

2. model_runner _ragged_fill_deferred_all_same — assembled ragged decode input_ids via 2× D2H (prev/draft ids) + a Python per-seq scatter loop + a full H2D each step (the two D2H also forced a sync on the sampler/draft kernels). Rebuilt entirely on-device (gather anchors/drafts + masked select), mirroring the rectangular all-same path.

Impact

Restores CPU/GPU overlap on the DSpark ragged + PIECEWISE + DP-attention decode path. The large aten::item span disappears from the profiler trace and decode throughput improves.

Validation

• Losslessness: ragged on/off produce identical tokens under --ignore-eos; MTP acceptance rate unchanged. Change Update logo #2 alters input_ids assembly — please confirm before merge.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1698 --add-label <label>

[Copilot]

Pull request overview

This PR optimizes the DSpark speculative-decode hot path by removing per-decode-step GPU→CPU synchronization points and rebuilding parts of the ragged-input assembly to avoid host round-trips.

Changes:

• Replaces a DSpark-only aten::item scalar read in EagleProposer.prepare_inputs() with an on-device clamp.
• Reworks _ragged_fill_deferred_all_same() to assemble ragged deferred input_ids directly on the GPU buffer and zero-fill graph tail padding on-device.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
atom/spec_decode/eagle.py	Removes a per-step GPU scalar read by clamping token indices fully on-device for DSpark.
atom/model_engine/model_runner.py	Replaces CPU scatter + D2H/H2D rebuild of ragged deferred input_ids with a GPU-side gather/select write into the flat buffer, plus GPU tail zeroing.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.