Skip to content

perf(mtp): KV-write-only MtpForward variant for the post-verify refresh (#209 item 3) #289

Description

@pekkah

Follow-up to #209 (work item #3). #209's #1 shipped in #287.

Problem

The post-verify refresh in MtpDecoder.DecodeBatched re-runs the full MtpForward once per accepted draft just to rewrite the MTP K/V with the trunk hiddens from BatchVerify. MtpForward includes the ~152K×emb shared-head lm_head matvec and a blocking logits download (CudaHybridGdnForwardPass.MtpForward tail) whose result is then discarded by the refresh.

Goal

A head-block-only MtpForward variant that stops after the MTP attention block writes K/V — skip the shared-head norm + lm_head matvec + the blocking D2H. Estimated ~1–1.6 ms saved per accepted draft.

Notes

  • The refresh only needs the K/V side effect, not the logits.
  • Keep the chained-draft MtpLastHidden capture intact (drafting still needs the self-hidden).
  • Bit-identity of the committed tokens is unaffected (refresh is a K/V rewrite, not a token decision).

Metadata

Metadata

Assignees

No one assigned

    Labels

    perfPerformance optimization opportunity

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions