Skip to content

Reduce repeated attention setup during speculative decoding - #38162

Closed
ormandj wants to merge 1 commit into
sgl-project:xinyuan/glm-5.3-flash-supportfrom
ormandj:pr/kpool-metadata-fusion
Closed

ormandj wants to merge 1 commit into
sgl-project:xinyuan/glm-5.3-flash-supportfrom
ormandj:pr/kpool-metadata-fusion

Conversation

@ormandj

@ormandj ormandj commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

Superseded by #38213 against main. This PR auto-closed after #36507 merged and its base branch was deleted. The historical description and validation below refer to this original revision.

Motivation

GLM-5.3-Flash's compressed sparse attention repeatedly prepares sequence lengths, page tables and attention schedules during generation. The existing path performs separate GPU operations for this bookkeeping, and draft backends can rebuild metadata already produced for the same step. This PR combines those updates and reuses compatible fresh metadata, reducing repeated setup work around model forwards without changing the attention inputs.

This is an opt-in optimization, not a fix for incorrect model answers. It restores the active metadata-fusion option removed by #38071 as a separate change, stacked on #36507. No isolated end-to-end speedup has been measured.

Modifications

  • Fuse KPool metadata updates for decode, target verification, draft extension and MTP precompute.
  • Preserve live-tail lengths, bound scans, and reuse compatible freshly generated draft metadata without replacing destination buffers.
  • Enable with SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1 only for supported CUDA/page geometry. The option defaults off; removed in-graph and preallocated-plan experiments remain excluded.

Accuracy Tests

The GPU tests feed the same sequence lengths and page mappings to the fused kernels and eager references, then compare their output tensors at pool, page and top-k boundaries, including partially filled live tails. This checks that combining the setup operations preserves the metadata consumed by attention. CPU tests check eligibility and require the original recompute path when source metadata is stale or destination buffers are incompatible.

PYTHONPATH=python compute-sanitizer --tool memcheck --error-exitcode 99 python -m pytest -q test/registered/kernels/ops/attention/test_dsa_metadata.py

Author run at efd2a02d03 on RTX PRO 6000 Blackwell (SM120), Transformers 5.12.1 / Tokenizers 0.22.2: 9 passed, 6 subtests passed; Compute Sanitizer reported ERROR SUMMARY: 0 errors. The CPU contract tests passed three tests and six subtests.

Speed Tests and Profiling

No isolated fusion-on/off benchmark. Full-integration results do not measure this option's contribution.

Checklist

  • Add CPU and GPU regression coverage.
  • Format changed files with upstream tools.
  • Complete final rebased-head tests and review.

Developed with AI assistance.


CI States

Latest PR Test (Base): ❌ Run #34000701731
Latest PR Test (Extra): ❌ Run #34000701594
Latest PR Test (AMD ROCm 7.2): ❌ Run #34000701683

@ormandj ormandj changed the title Fuse repeated metadata updates during speculative decoding Reduce repeated attention setup during speculative decoding Sep 6, 2026
@Fridge003
Fridge003 deleted the branch sgl-project:xinyuan/glm-5.3-flash-support September 6, 2026 09:28
@Fridge003 Fridge003 closed this Sep 6, 2026
@ormandj

ormandj commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

Refiled as #38213 against main after #36507 merged and this PR's base branch was deleted. The patch is unchanged (matching stable Git patch ID); author CPU regression checks on the new head passed 3 tests and 6 subtests. GPU validation attached here remains tied to the original revision and has not been rerun on the new main-based head. Please continue review on #38213.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants