Skip to content

refactor: unify decode row construction into decode_input_builder. - #2640

Draft
wai-bao wants to merge 1 commit into
xLLM-AI:mainfrom
wai-bao:refactor/decode-width-unification-v2
Draft

wai-bao wants to merge 1 commit into
xLLM-AI:mainfrom
wai-bao:refactor/decode-width-unification-v2

Conversation

@wai-bao

@wai-bao wai-bao commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Unifies the four per-site decode row-construction copies (MTP validate, MTP draft-extend, DFlash query, fixed/dflash2 verify) into a single decode_input_builder with a declarative plan, and consolidates the shared seq-lens / paged-KV layout knowledge that those copies each reimplemented.

  • Deletes spec_input_builder (852 lines) in favor of framework/batch/decode_input_builder.{h,cpp}: RowSpec rows, DecodeLayout (flat / chunked / block-parallel / chunked-decode), plan-driven commit (update_block_table / commit_decode_input / commit_decode_geometry), and a named expand_idle_rank_block entry for idle DP-rank synthetic blocks.
  • Adds kv_cache/layout/seq_lens_layout.h: the backend seq-lens layout policy (NPU per-seq vs cu_seqlens), single source for seq_len_at, append_seq_len(s), append_qo_indptr, seq_lens_capacity — replaces three get_q_seq_len copies and per-builder build_q_cu_seq_lens_vec copies.
  • Adds kv_cache/layout/paged_kv_layout.h: row-level paged-KV geometry (kv_page_count, kv_last_page_len, append_paged_kv_entries, slot_for_position, compute_slot_mapping) shared by the batch builders, the decode-row builder, the expanded-decode metadata builder, and the DSA metadata builders.
  • Merges the two ~90-line update_decode_step_input copies (MTP / DFlash) behind ExpectedKvPolicy in mtp_runtime_helpers, with the chunked-verify clamp parameters carried inside the policy.
  • Hot-path wins along the way: cached device index tensors (CachedDeviceIndices) for draft-query / draft-extend / validate indices, model_type resolved once per worker instead of re-parsing config.json per predicate call, per-sequence hoists in draft-prefix bounds, and the per-step result.lengths.to(int32) cast removed from the verify commit.

Net −276 lines across 53 files.

Testing

  • decode_input_builder_test (21), task_execution_pipeline_speculative_test (152), mtp_decode_input_test (25), mtp_decode_metadata_test (18), schedule_overlap_test (42), batch_test (134) all pass on NPU.

@wai-bao
wai-bao force-pushed the refactor/decode-width-unification-v2 branch 5 times, most recently from de2fd52 to 0b4514f Compare November 1, 2026 16:30
@wai-bao wai-bao changed the title refactor: unify decode row construction into decode_input_builder with single-source layout, plan-driven commit, and shared paged KV geometry. refactor: unify decode row construction into decode_input_builder. Nov 1, 2026
@wai-bao
wai-bao force-pushed the refactor/decode-width-unification-v2 branch 13 times, most recently from b8f50fd to b5d4a50 Compare November 2, 2026 10:27
@netlify

netlify Bot commented Nov 2, 2026 •

Copy link
Copy Markdown

✅ Netlify Deploy Preview ready

Built without sensitive environment variables

Name Link
🔨 Latest commit f3e7ffb
🔍 Latest deploy log https://app.netlify.com/projects/xllm-docs-preview/deploys/6aec9fb656a2976df8a60454
😎 Deploy Preview https://deploy-preview-2640--xllm-docs-preview.netlify.app
📱 Preview on mobile
Toggle QR Code...

QR Code

Use your smartphone camera to open QR code link.
🤖 Make changes Run an agent on this branch

To edit notification comments on pull requests, go to your Netlify project configuration.

@wai-bao
wai-bao force-pushed the refactor/decode-width-unification-v2 branch 10 times, most recently from 1a5d531 to 7605f8b Compare November 3, 2026 10:18
@wai-bao
wai-bao force-pushed the refactor/decode-width-unification-v2 branch 29 times, most recently from 6b447b9 to 19a13c5 Compare November 4, 2026 18:58

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant