[PixelFlat]

Description

Capture Python DSpark context KV updates in ACL graphs and reuse the existing task execution pipeline for DSpark decoding.

• Capture context projection, normalization, RoPE, and KV scatter during startup on a private capture stream. Compact execution warms all accepted-row buckets; pipeline execution warms both final Slot bindings with no-write cache slots.

• Replay the pipeline's borrowed target hidden states, positions, and cache-slot views directly. Keep rejected-row masking, context writes, and next-round state handoff on the ordered device stream instead of compacting and uploading accepted rows through the host.

• Add DSpark-specific fixed-width draft bindings and anchor-row sampling while preserving the sequential Markov bias dependency between proposals. Keep target verification, draft forward, and context update graphs distinct.

• Keep the prepared draft's FIA capacity constant while reading live device KV lengths. Missing startup bindings fail explicitly; only declared out-of-capacity queries take the existing eager path, without runtime capture or exception fallback.

• Isolate internal ATB cache symbols in the parent third_party/CMakeLists.txt so KV writes cannot bind to incompatible torch_npu plugin cache definitions. Include the FP16/BF16 regression test; no submodule source or revision changes are needed.

Task pipeline admission is currently limited to DP=1 and single-request capacity (max_seqs_per_batch=1). The existing non-pipeline DSpark DP graph path remains available. This PR does not change target quantization or weight-loading formats.

Related Issues

Follow-up optimization to #2440 and #2473.

Change Type

• Bug fix
• New feature
• Performance improvement
• Test

Validation

On Linux ARM64 / Ascend NPU, the implementation branch passed:

• python setup.py build --device npu --tilelang-jobs 4.
• 43 focused native tests covering algorithm admission, speculative task execution, and draft input layouts.
• Nine relevant Python test modules: atb_cache_isolation, context_kv_acl_graph, model_executor, qwen3_dflash, qwen3_dspark, decode_acl_graph, prepared_acl_graph, eplb_runtime_bridge, and npu_paged_attention.
• Real-NPU eager/graph parity for context writes and prepared DSpark forward with changing device KV lengths, masked/rejected cache slots, and both startup stream types.
• Paired, same-input context update checks across all 16 ranks: 512 checks covering 1-8 accepted rows, zero hidden/KV differences, and unchanged cache guards. This verifies context update parity, not every asynchronous state transition.
• ATB cache regression: both FP16/BF16 cases failed without isolation (RuntimeError: Invalid argument) and passed with the included fix.
• End-to-end GLM-5.2 static W8A8 serving with target, draft, and context graph replay. Capture counters remained unchanged after startup.

These NPU results were collected before rebasing and consolidating the PR onto current main. The validation build included the ATB cache symbol-isolation fix, now included in this PR as a separate bugfix commit. The consolidated PR head passed changed-file pre-commit checks, including clang-format 20.1.6; it has not received a fresh NPU rerun. This is focused validation, not a claim that the complete python setup.py test suite passed.

Performance

All cases use the same GLM-5.2 static W8A8 target, Python model implementation, TP=16 / DP=1 / EP=1, ACL graphs, and asynchronous scheduling. DSpark always uses seven draft steps; MTP uses three steps with task pipeline enabled. The DSpark columns compare recorded configurations at the corresponding implementation stages. One unchanged ShareGPT request with 1,868 input tokens and 128 output tokens, concurrency 1, explicit temperature=0, two warmup requests followed by five non-profiled measurements. Memory settings were unchanged between cases.

Mean metric	DSpark: baseline	DSpark: context graph only	DSpark: context graph + pipeline	MTP3 + pipeline
TTFT	612.00 ms	614.80 ms	619.40 ms	622.80 ms
TPOT	16.10 ms	14.70 ms	13.56 ms	13.46 ms
Total latency	2657.78 ms	2483.22 ms	2339.00 ms	2332.02 ms
Accepted draft tokens / request	84.8	88.2	85.8	80.6
Draft tokens / request	299.6	275.8	294.0	141.6
Draft rounds / request	42.8	39.4	42.0	47.2
Accepted draft tokens / round	1.98	2.24	2.04	1.71
Acceptance rate	28.80%	32.95%	29.22%	57.02%

The DSpark baseline already enables target/draft graphs and asynchronous scheduling, but uses eager context KV updates and no task pipeline. The middle DSpark column enables context KV graphs without task pipeline; the final DSpark column enables both optimizations.

Acceptance rate is the mean of the five per-request rates. Accepted draft tokens per round is total accepted draft tokens divided by total draft rounds; it excludes the target correction/bonus token.

DSpark TPOT changed from 16.10 ms to 13.56 ms in these five-run means (15.8% lower). Acceptance and draft-round counts also vary, so these separately collected results do not isolate the causal contribution of each optimization or establish a universal speedup. Full-model dynamic W8A8 validation and DSpark task pipeline with DP>1 are outside this validation scope.

Timeline Evidence

Each trace is a separate single HTTP request after warmup, captured on rank 0/device 0. Steady decode analysis excludes the initial boundary and final incomplete round. The three DSpark traces contain 44, 43, and 39 complete steady rounds respectively. These profiled intervals are diagnostic data, separate from the five-request non-profiled performance table.

Mean observed interval / count	Eager context, pipeline off	Context graph, pipeline off	Context graph + pipeline
Context graph executions / round	0	1	1
Context KV scatter kernels / round	5	5	5
Context projection to last KV scatter	1.010 ms	0.610 ms	0.411 ms
Draft execution work envelope	2.506 ms	1.802 ms	1.435 ms
Target execution work envelope	39.290 ms	38.563 ms	38.381 ms
Draft graph launch to target graph launch	4.560 ms	4.577 ms	2.064 ms
Prefix verification start to next draft launch	3.513 ms	3.750 ms	0.520 ms
Full draft-launch-to-next-draft interval	47.625 ms	47.155 ms	41.191 ms
Post-verify interval outside recorded compute/copy union	2.725 ms	2.963 ms	0.031 ms

Context is inside the graph: the optimized device trace repeatedly shows draft -> target -> context graph launches. Between greedy_prefix_verify_kernel and the next draft, one independent MODEL_EXECUTE is followed by context MatMul, RMSNorm and RoPE operations, then five reshape_and_cache_200000000 kernels. The latest trace records 41 draft, 41 target, and 41 context graph executions. All 16 ranks warm two borrowed context bindings; startup capture counters remain unchanged during serving.

For a concrete optimized interval, relative to the first timestamped X event, verification starts at 821.8435 ms, the context MODEL_EXECUTE starts at 821.88375 ms, context computation runs from 821.88825 to 822.29875 ms, and the next draft starts at 822.3685 ms. The control event's 3.46 us duration is not the context graph's execution time.

Reduced state-handoff bubble: versus context-graph-only execution, the post-verify interval decreases from 3.750 to 0.520 ms. The recorded compute/copy union occupies 0.787 and 0.488 ms respectively; subtracting that union leaves approximately 2.963 versus 0.031 ms uncovered within the selected window. Recorded host API starts in this window change from about 28.0 Node@launch, 8.0 synchronous copies, 19.9 asynchronous copies, and 10.0 stream synchronizations per round to zero for those categories. Pipeline host preparation can occur earlier: this does not mean all host work or every device bubble has disappeared.

Intervals include dependent device work, not just idle time. Work envelopes must not be summed across overlapping streams; the uncovered interval is a recorded-task gap proxy, not a measurement of every rank's utilization. Phases are identified by actual attention/context kernels, not Model Id. Analysis used JSON events; Perfetto UI inspection was not performed.

The latest MTP3 trace has about 1.626 ms summed draft work across three forwards versus DSpark's 1.435 ms single draft forward. Its target work envelope is 33.883 ms versus DSpark's 38.381 ms: the remaining difference in this observation is predominantly wider target verification, not a slower DSpark draft.

Pull Request Checklist

• PR title and commit message follow the xLLM commit format.
• Pre-commit hooks are installed; changed-file checks passed.
• Self-reviewed against the project-specific code style.
• Based on latest main at submission time.
• Focused tests added or updated.
• Complete NPU python setup.py build test rerun on the consolidated PR head.

Reviewer Notes

Please focus on final Slot binding lifetimes and stream ordering, no-write startup context warmup, DSpark's N-row draft versus N+1-row verification contract, and preservation of Markov proposal sampling. The prerequisite ATB cache isolation is included as a separate commit. It changes parent-repository target visibility only; the torch_npu_ops submodule remains unchanged.

Summary by CodeRabbit

• New Features
  • Added DSpark speculative decoding support for eligible task-pipeline configurations.
  • Added support for accelerated graph execution during DSpark decoding on Ascend NPU.
• Performance
  • Added graph-based context KV updates, with replay for supported workloads and eager execution when requests exceed warmed capacity.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-12T11:16:05.024343Z	b2f3cd5	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 45289748-6b33-4bca-9240-a045f1d7ce9a

---

📥 Commits

Reviewing files that changed from the base of the PR and between ce249b4 and 6edae15.

---

📒 Files selected for processing (23)

• tests/core/runtime/acl_graph_executor_test.cpp
• tests/core/runtime/mtp_decode_input_test.cpp
• tests/core/runtime/task_execution_pipeline_speculative_test.cpp
• tests/python/test_atb_cache_isolation.py
• tests/python/test_context_kv_acl_graph.py
• tests/python/test_eplb_runtime_bridge.py
• tests/python/test_model_executor.py
• tests/python/test_prepared_acl_graph.py
• third_party/CMakeLists.txt
• xllm/core/distributed_runtime/llm_master.cpp
• xllm/core/framework/config/speculative_config.h
• xllm/core/runtime/dflash_worker_impl.cpp
• xllm/core/runtime/py_executor_impl.cpp
• xllm/core/runtime/slot_buffer.cpp
• xllm/core/runtime/slot_buffer.h
• xllm/core/runtime/task_execution_pipeline.h
• xllm/core/runtime/task_execution_pipeline_speculative.cpp
• xllm/core/runtime/worker.cpp
• xllm/models/llm/py_causal_lm.cpp
• xllm/models/llm/py_causal_lm.h
• xllm/python/model_executor/executor.py
• xllm/python/model_executor/runners/context_kv_acl_graph.py
• xllm/python/model_executor/runners/prepared_acl_graph.py

---

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

---

📝 Walkthrough

📝 Walkthrough

📝 Walkthrough

Walkthrough

The pull request adds DSpark as a speculative task-pipeline algorithm. It implements DSpark block-draft sampling and adds context-KV and prepared ACL graph execution paths, with tests for graph capture, replay, cache updates, and proposal handling.

Changes

DSpark speculative decoding and graph execution

Layer / File(s)	Summary
Task-pipeline admission and anchor input xllm/core/framework/config/speculative_config.h, xllm/core/distributed_runtime/llm_master.cpp, xllm/core/runtime/worker.cpp, xllm/core/runtime/dflash_worker_impl.cpp, xllm/core/runtime/task_execution_pipeline.h, xllm/core/runtime/slot_buffer.*, tests/core/runtime/acl_graph_executor_test.cpp, tests/core/runtime/mtp_decode_input_test.cpp	Task-pipeline support now includes DSpark. DSpark receives a task kind, and block-draft input can sample from the anchor without an extra row. Tests cover algorithm admission, input tensors, and rejected invocation types.
DSpark draft execution xllm/core/runtime/task_execution_pipeline_speculative.cpp, tests/core/runtime/task_execution_pipeline_speculative_test.cpp	DSpark setup uses per-sample draft indices and DSpark-specific graph warmup and decode settings. Proposal sampling uses a sequential Markov path. Tests cover warmup, overlapping decodes, rejected proposals, and unsupported topology.
Context-KV graph capture and dispatch xllm/python/model_executor/runners/context_kv_acl_graph.py, xllm/python/model_executor/executor.py, xllm/core/runtime/py_executor_impl.cpp, xllm/models/llm/py_causal_lm.*, third_party/CMakeLists.txt, tests/python/test_context_kv_acl_graph.py, tests/python/test_atb_cache_isolation.py, tests/python/test_eplb_runtime_bridge.py, tests/python/test_model_executor.py	The context-KV runner binds caches and captures and replays graph entries. The executor and PyCausalLM pass cache bindings and eligible context writes through the runner. Tests cover capture, replay, stream ordering, cache writes, and startup warmup.
Prepared block-draft graph execution xllm/python/model_executor/runners/prepared_acl_graph.py, xllm/python/model_executor/executor.py, tests/python/test_prepared_acl_graph.py, tests/python/test_model_executor.py	Prepared ACL graph handling adds block-draft validation, warmup metadata, replay, and an admission check. DSpark requests that do not fit the prepared graph use eager execution. Tests cover graph reuse, capacity limits, and fallback.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant TaskExecutionPipeline
  participant DraftModel
  participant Sampler
  TaskExecutionPipeline->>DraftModel: Request Markov bias for the preceding token
  DraftModel-->>TaskExecutionPipeline: Return vocabulary bias
  TaskExecutionPipeline->>Sampler: Sample proposal
  Sampler-->>TaskExecutionPipeline: Return proposal token
  TaskExecutionPipeline->>DraftModel: Use proposal as the next-step input

Loading

sequenceDiagram
  participant ModelExecutor
  participant ContextKvAclGraphRunner
  participant PyCausalLM
  participant KVCaches
  ModelExecutor->>ContextKvAclGraphRunner: Bind caches and warm graph buckets
  PyCausalLM->>ContextKvAclGraphRunner: Execute context-KV update
  ContextKvAclGraphRunner->>KVCaches: Replay cache update
  ContextKvAclGraphRunner-->>PyCausalLM: Return updated hidden states

Loading

Suggested reviewers: oleanderrr

---

Merge Risk: ⚪ Minimal · up to 6edae

No confirmed merge-blocking defect remains. Lowercase DSpark configuration support has not been established, so it is not treated as a demonstrated regression.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 6edae

The change alters how request state is updated. Enforced capacity limits, stable cache ownership, and ordered updates bound the exposure. Recovery after interrupted execution remains incompletely established.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The reviewed state-mutation path is bounded to a configured worker's target and draft caches, with one sequence in flight across up to two pipeline Slots. A fatal execution failure can affect that worker's availability. Fleet-wide, tenant, and deployment recovery scope cannot be established from the supplied topology evidence.

Trust Boundaries and Controls

• observed — In the inspected serving path, the Python executor is constructed and invoked through the internal C++ bridge. Runtime code supplies cache bindings and graph metadata. Request-dependent data influences model computation, but the new calls do not themselves give requests cache-binding or capture authority.

Resilience and Maintainability Implications

• inferred — The launch loop runs directly in a std::thread without an exception handler. An exception escaping replay therefore terminates the process rather than permitting that loop to continue with partially published state. This structure predates the PR. It does not establish safe in-process retry or deployment restart behavior after capture or device failure.

---

Pre-merge checks | 4 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 10.47% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 86 functions across 22 files. (1 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Title check		The title clearly summarizes the main changes: capturing DSpark context KV in ACL graphs and enabling the task pipeline. It is concise and uses the required perf prefix.
Description check		The description is comprehensive. It explains the implementation, scope, validation results, performance data, limitations, checklist status, and reviewer focus areas. It does not reproduce every temp…

---

Full details: Docstring Coverage

Explanation

Docstring coverage is 10.47% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 86 functions across 22 files. (1 skipped: 1 unsupported.)

---

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

---

🧪 Generate unit tests (beta)

• Create a new PR

---

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}