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).
Follow-up to #209 (work item #3). #209's #1 shipped in #287.
Problem
The post-verify refresh in
MtpDecoder.DecodeBatchedre-runs the fullMtpForwardonce per accepted draft just to rewrite the MTP K/V with the trunk hiddens fromBatchVerify.MtpForwardincludes the ~152K×emb shared-head lm_head matvec and a blocking logits download (CudaHybridGdnForwardPass.MtpForwardtail) whose result is then discarded by the refresh.Goal
A head-block-only
MtpForwardvariant 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
MtpLastHiddencapture intact (drafting still needs the self-hidden).