Support Step3.5/3.7 flash mtp3

[forforever73]

Overview

follow-up to #23274.(cc @pwilkin )

📜 Full data-flow trace — couldn't think of a good way to draw this, so I wrote it all down instead. It's long, but every byte is load-bearing.

Notation:

• token@pos / h(pos) — positions are explicit (0-indexed)
• h_tgt(p) — target NextN hidden at p (before the output norm)
• h45(p) / h46(p) — head 45/46 output hidden, chained between heads while drafting
• pending_h = h_tgt(pos of id_last − 1) — always the trunk h, regardless of chaining

Example: a 4-token prompt at positions 0–3.

Time ─────────────────────────────────────────────────────────────────────▶

═══ Round 0: prompt bootstrap ══════════════════════════════════════════════
  prompt: [t0@0, t1@1, t2@2, t3@3]   (need_embd → all logits=true)

  ② target decode → verify_h = [h_tgt(0), h_tgt(1), h_tgt(2), h_tgt(3)]

  ③ process (mirror once per head, all logits=0, h_tgt shifted right by 1):
       for head in {45, 46, 47}:
         seq_rm(ctx_dft, seq, ≥ this ubatch's first pos)   // = 0 here; reset before each head
         token = [t0@0, t1@1,    t2@2,    t3@3]
         embd  = [0,    h_tgt(0), h_tgt(1), h_tgt(2)]   ← row 0 = pending_h = 0 (sentinel)
         → write KV_head@0..3 (teacher-forced, logits=0)
                              h_tgt(3) pushed out → pending_h default = h_tgt(3)
       (no draft → accept not called → pending_h stays = h_tgt(3))

  sample logits at pos 3 → T_first@4
       slot.sampled = T_first@4 ;  pending_h = h_tgt(3)
       invariant: id_last=T_first@4, pending_h=h_tgt(3)=h_tgt(4-1)  ✓

═══ Round 1: generation (id_last=T_first@4, pending_h=h_tgt(3)) ═════════════
  ⑥ draft (accumulated batch; switch head each step, 1 decode/step, batch grows):
       round_start = pos of id_last = 4; seq_rm(≥4) before switching to each head

       set_mtp_layer_offset(ctx_dft, 0)   // head 45
       step0 [h45]  seq_rm(≥4); batch=[T_first@4],            embd=[h_tgt(3)]
                    → decode → write KV45@4; sample draft1@5; output h45(4)

       set_mtp_layer_offset(ctx_dft, 1)   // head 46
       step1 [h46]  seq_rm(≥4); batch=[T_first@4, draft1@5],  embd=[h_tgt(3), h45(4)]
                    → decode → write KV46@4,5; sample draft2@6; output h46(5)

       set_mtp_layer_offset(ctx_dft, 2)   // head 47
       step2 [h47]  seq_rm(≥4); batch=[T_first@4, draft1@5, draft2@6],
                    embd=[h_tgt(3), h45(4), h46(5)]
                    → decode → write KV47@4,5,6; sample draft3@7
                    (draft length capped at n_mtp_layers = 3)

       spec_draft = [draft1, draft2, draft3]

  ── post-draft seq_rm (server, cross-round cleanup; unchanged) ──
       ckpt.pos_max = 3 (ctx_tgt only covers prompt @0..3 here; T_first not decoded yet)
       seq_rm(ctx_dft, slot.id, ≥4)  → every head's KV trimmed back to @0..3

  ① target batch (4 tokens, all logits=true):
       token = [T_first@4, draft1@5, draft2@6, draft3@7]

  ② target decode → verify_h = [h_tgt(4), h_tgt(5), h_tgt(6), h_tgt(7)] + logits[0..3]

  ③ process (once per head, logits=0; embd = h_tgt right-shift, NOT inter-head h):
       for head in {45, 46, 47}:
         seq_rm(ctx_dft, seq, ≥4)   // reset before each head; all heads rewrite from @4
         token = [T_first@4, draft1@5, draft2@6, draft3@7]
         embd  = [h_tgt(3),  h_tgt(4), h_tgt(5), h_tgt(6)]
                  ↑ pending_h           h_tgt(7) pushed out → pending_h default
         → rewrite KV_head@4..7 (teacher-forced, logits=0)

  ④ verify (sample from ctx_tgt logits, identical to the single-block path):
       logits[0] → real tok@5 == draft1  ✓
       logits[1] → real tok@6 == draft2  ✓
       logits[2] → real tok@7 != draft3  ✗ → resample T_new@7
       accepted = [draft1, draft2, T_new] ;  n_accepted = 2

  ⑤ accept (identical to the single-block path):
       pending_h    = verify_h[n_accepted] = verify_h[2] = h_tgt(6)
       slot.sampled = accepted.back()      = T_new (pos 7)
       invariant: id_last=T_new@7, pending_h=h_tgt(6)=h_tgt(7-1)  ✓

═══ Round 2: generation (id_last=T_new@7, pending_h=h_tgt(6)) ═══════════════
  ⑥ step0 [h45]  seq_rm(≥7); batch=[T_new@7],                  embd=[h_tgt(6)]
                 → write KV45@7; draft1'@8, h45(7)
     step1 [h46]  seq_rm(≥7); batch=[T_new@7, draft1'@8],       embd=[h_tgt(6), h45(7)]
                 → write KV46@7,8; draft2'@9, h46(8)
     step2 [h47]  seq_rm(≥7); batch=[T_new@7, draft1'@8, draft2'@9],
                 embd=[h_tgt(6), h45(7), h46(8)]
                 → write KV47@7,8,9; draft3'@10
     ── post-draft seq_rm(≥8); ① target batch; ② target decode; ③ process ×3
        (process resets seq_rm(≥7) per head, then rewrites @7..10)

Core strategy
Each MTP head is its own decoder layer with its own KV, and the driver runs one head per llama_decode. A seq_rm before each head clears the range it re-decodes, so it reuses the same slots (find_slot is deterministic) instead of stacking duplicate positions; find_slot / apply_ubatch are untouched. The two phases differ only in what feeds the heads:

phase	embd fed to each head	purpose
process()	trunk h_tgt, right-shifted by one (not the inter-head hidden)	re-anchor each head's committed-prefix KV to the target's real hidden → next round starts target-aligned
draft()	the previous head's output, chained (slot 0 = trunk pending_h)	generate the draft tokens; each head rebuilds its own layer's KV on the growing prefix

Only the trunk h_tgt crosses rounds, so pending_h / verify_h stay single-layer.

MTP Block Selection Strategy

• cparams.mtp_layer_offset (src/llama-cparams.h) — picks which appended MTP block the DECODER_MTP graph runs: il = n_layer() + offset. Default 0.
• graph_mtp selects the head by offset (il = n_layer() + cparams.mtp_layer_offset, was a hardcoded n_layer()).
• graph_mtp now gathers its output rows via build_inp_out_ids(), like the trunk graph. The fix that makes chaining work: from step 1 on, the output is the last batch row, not row 0, so without it heads 46/47 read the wrong row. Identity gather when n_outputs == n_tokens, so the single-head path is unchanged.
• Loader now requires all n_layer_nextn MTP blocks.
• n_max is clamped to the head count when chaining (each head used once).

Results

./llama-server \
    -m Step-3.7-IQ4_XS.gguf \
    --spec-type draft-mtp \
    --spec-draft-model Step3.7-flash-mtp-Q8_0.gguf \
    -ngl all \
    --spec-draft-ngl all \
    -c 35000 \
    -np 1 \
    -b 2048 \
    -ub 1024 \
    --temp 0 \
    --spec-draft-n-max {n} \
    --spec-draft-p-min 0.0 \
    --host 127.0.0.1 \
    --port 8080

The command is identical on both machines; only --spec-draft-n-max and the build change. Before = single-block MTP on master (one head, looped when n-max > 1); after = the three-layer chain.

DGX Spark GB10

Before (single-block MTP, master)

--spec-draft-n-max 2

  code_python        pred= 192 draft= 164 acc= 108 rate=0.658 tok/s=30.8
  code_cpp           pred= 192 draft= 172 acc= 104 rate=0.605 tok/s=30.5
  explain_concept    pred= 192 draft= 171 acc= 105 rate=0.614 tok/s=30.8
  summarize          pred= 192 draft= 160 acc= 110 rate=0.688 tok/s=33.2
  qa_factual         pred= 192 draft= 146 acc= 117 rate=0.801 tok/s=36.6
  translation        pred= 192 draft= 173 acc= 104 rate=0.601 tok/s=30.8
  creative_short     pred= 192 draft= 189 acc=  95 rate=0.503 tok/s=28.1
  stepwise_math      pred= 192 draft= 156 acc= 113 rate=0.724 tok/s=34.6
  long_code_review   pred= 192 draft= 161 acc= 110 rate=0.683 tok/s=31.7

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1728,
  "total_draft": 1492,
  "total_draft_accepted": 966,
  "aggregate_accept_rate": 0.6475,
  "wall_s_total": 60.19
}

--spec-draft-n-max 3

  code_python        pred= 192 draft= 233 acc= 112 rate=0.481 tok/s=28.8
  code_cpp           pred= 192 draft= 243 acc= 109 rate=0.449 tok/s=27.7
  explain_concept    pred= 192 draft= 252 acc= 106 rate=0.421 tok/s=26.3
  summarize          pred= 192 draft= 233 acc= 112 rate=0.481 tok/s=27.7
  qa_factual         pred= 192 draft= 196 acc= 125 rate=0.638 tok/s=30.5
  translation        pred= 192 draft= 242 acc= 109 rate=0.450 tok/s=25.7
  creative_short     pred= 192 draft= 271 acc=  99 rate=0.365 tok/s=24.9
  stepwise_math      pred= 192 draft= 226 acc= 115 rate=0.509 tok/s=30.0
  long_code_review   pred= 192 draft= 235 acc= 112 rate=0.477 tok/s=27.3

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1728,
  "total_draft": 2131,
  "total_draft_accepted": 999,
  "aggregate_accept_rate": 0.4688,
  "wall_s_total": 70.1
}

After (three-layer MTP, this PR)

--spec-draft-n-max 2

  code_python        pred= 192 draft= 154 acc= 113 rate=0.734 tok/s=33.9
  code_cpp           pred= 192 draft= 156 acc= 112 rate=0.718 tok/s=32.6
  explain_concept    pred= 192 draft= 153 acc= 114 rate=0.745 tok/s=34.9
  summarize          pred= 192 draft= 150 acc= 115 rate=0.767 tok/s=35.5
  qa_factual         pred= 192 draft= 146 acc= 117 rate=0.801 tok/s=36.9
  translation        pred= 192 draft= 164 acc= 108 rate=0.658 tok/s=29.8
  creative_short     pred= 192 draft= 171 acc= 104 rate=0.608 tok/s=30.0
  stepwise_math      pred= 192 draft= 139 acc= 121 rate=0.871 tok/s=39.4
  long_code_review   pred= 192 draft= 143 acc= 118 rate=0.825 tok/s=36.2

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1728,
  "total_draft": 1376,
  "total_draft_accepted": 1022,
  "aggregate_accept_rate": 0.7427,
  "wall_s_total": 56.95
}

--spec-draft-n-max 3

  code_python        pred= 192 draft= 188 acc= 128 rate=0.681 tok/s=35.1
  code_cpp           pred= 192 draft= 213 acc= 119 rate=0.559 tok/s=31.4
  explain_concept    pred= 192 draft= 208 acc= 121 rate=0.582 tok/s=32.2
  summarize          pred= 192 draft= 203 acc= 122 rate=0.601 tok/s=33.0
  qa_factual         pred= 192 draft= 181 acc= 130 rate=0.718 tok/s=38.0
  translation        pred= 192 draft= 198 acc= 124 rate=0.626 tok/s=34.7
  creative_short     pred= 192 draft= 244 acc= 108 rate=0.443 tok/s=27.4
  stepwise_math      pred= 192 draft= 170 acc= 134 rate=0.788 tok/s=40.6
  long_code_review   pred= 192 draft= 199 acc= 124 rate=0.623 tok/s=33.1

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1728,
  "total_draft": 1804,
  "total_draft_accepted": 1110,
  "aggregate_accept_rate": 0.6153,
  "wall_s_total": 56.57
}

Mac Studio M4 Max

Before (single-block MTP, master)

--spec-draft-n-max 2

  code_python        pred= 192 draft= 162 acc= 110 rate=0.679 tok/s=42.6
  code_cpp           pred= 192 draft= 179 acc= 101 rate=0.564 tok/s=38.4
  explain_concept    pred= 192 draft= 171 acc= 105 rate=0.614 tok/s=40.2
  summarize          pred= 192 draft= 162 acc= 110 rate=0.679 tok/s=42.5
  qa_factual         pred= 161 draft= 122 acc= 101 rate=0.828 tok/s=47.4
  translation        pred= 192 draft= 171 acc= 104 rate=0.608 tok/s=40.0
  creative_short     pred= 192 draft= 188 acc=  97 rate=0.516 tok/s=36.7
  stepwise_math      pred= 192 draft= 155 acc= 113 rate=0.729 tok/s=44.3
  long_code_review   pred= 192 draft= 165 acc= 108 rate=0.654 tok/s=41.1

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1697,
  "total_draft": 1475,
  "total_draft_accepted": 949,
  "aggregate_accept_rate": 0.6434,
  "wall_s_total": 46.38
}

--spec-draft-n-max 3

  code_python        pred= 192 draft= 233 acc= 113 rate=0.485 tok/s=33.7
  code_cpp           pred= 192 draft= 259 acc= 104 rate=0.402 tok/s=30.4
  explain_concept    pred= 192 draft= 249 acc= 107 rate=0.430 tok/s=31.7
  summarize          pred= 192 draft= 226 acc= 115 rate=0.509 tok/s=34.7
  qa_factual         pred= 161 draft= 165 acc= 105 rate=0.636 tok/s=39.9
  translation        pred= 192 draft= 244 acc= 109 rate=0.447 tok/s=32.3
  creative_short     pred= 192 draft= 278 acc=  98 rate=0.352 tok/s=28.4
  stepwise_math      pred= 192 draft= 220 acc= 117 rate=0.532 tok/s=35.8
  long_code_review   pred= 192 draft= 232 acc= 113 rate=0.487 tok/s=33.0

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1697,
  "total_draft": 2106,
  "total_draft_accepted": 981,
  "aggregate_accept_rate": 0.4658,
  "wall_s_total": 56.7
}

After (three-layer MTP, this PR)

--spec-draft-n-max 2

  code_python        pred= 192 draft= 150 acc= 116 rate=0.773 tok/s=45.3
  code_cpp           pred= 192 draft= 165 acc= 108 rate=0.654 tok/s=41.2
  explain_concept    pred= 192 draft= 155 acc= 113 rate=0.729 tok/s=43.8
  summarize          pred= 192 draft= 150 acc= 116 rate=0.773 tok/s=45.0
  qa_factual         pred= 161 draft= 114 acc= 104 rate=0.912 tok/s=49.7
  translation        pred= 192 draft= 153 acc= 113 rate=0.739 tok/s=43.8
  creative_short     pred= 192 draft= 172 acc= 104 rate=0.605 tok/s=39.1
  stepwise_math      pred= 192 draft= 144 acc= 119 rate=0.826 tok/s=46.8
  long_code_review   pred= 192 draft= 153 acc= 114 rate=0.745 tok/s=43.2

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1697,
  "total_draft": 1356,
  "total_draft_accepted": 1007,
  "aggregate_accept_rate": 0.7426,
  "wall_s_total": 43.83
}

--spec-draft-n-max 3

  code_python        pred= 192 draft= 188 acc= 128 rate=0.681 tok/s=40.8
  code_cpp           pred= 192 draft= 205 acc= 122 rate=0.595 tok/s=37.2
  explain_concept    pred= 192 draft= 207 acc= 121 rate=0.585 tok/s=36.8
  summarize          pred= 192 draft= 194 acc= 126 rate=0.649 tok/s=39.4
  qa_factual         pred= 161 draft= 141 acc= 114 rate=0.808 tok/s=45.6
  translation        pred= 192 draft= 199 acc= 124 rate=0.623 tok/s=38.3
  creative_short     pred= 192 draft= 242 acc= 109 rate=0.450 tok/s=31.5
  stepwise_math      pred= 192 draft= 170 acc= 134 rate=0.788 tok/s=44.9
  long_code_review   pred= 192 draft= 208 acc= 121 rate=0.582 tok/s=36.1

Aggregate: {
  "n_requests": 9,
  "total_predicted": 1697,
  "total_draft": 1754,
  "total_draft_accepted": 1099,
  "aggregate_accept_rate": 0.6266,
  "wall_s_total": 49.38
}

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: CC for code and test, but the core data flow was designed by human