perf(k3): batched spec-round propose + e2e A/B and round anatomy

docs/models/k3/mtp-dspark.md

@@ -9,7 +9,11 @@ cycle probe and 3.13 on a 4-prompt prose probe vs the same-checkpoint
 sglang reference's 3.0 / ~2.8; the full 896-expert EP16 target commits
 4.0 on code, 2.8–2.9 on English prose, 1.3 on Chinese. Verify is *not* bitwise plain decode
 (chunkwise FlashKDA vs the fused core — near-tie argmaxes flip); the
-`spec_verify` gates certify what is exact instead.
+`spec_verify` gates certify what is exact instead. E2e on the EP16 fleet
+(SPEED-Bench, 2026-08-19): **1.74× on low-entropy c1, 1.8× SLOWER on
+high-entropy c1, −27% throughput at c4 mixed** — the verify step's
+multi-row cost, not propose overhead, is the wall (see the round-anatomy
+section).
 
 Last touched: 2026-08
 
@@ -170,20 +174,78 @@ self-similar text saturate near RadixArk's published band (3.9–5.5,
 measured on chat-template benches); Chinese is the drafter's weak
 spot, not the verify path's.
 
-**Next** (perf follow-up PR — semantics-neutral, kept out of #931 so
-the bring-up PR stays reviewable): measure end-to-end spec-on vs
-spec-off tokens/s on EP16 to size the win, then CUDA-graph the verify
-step (never existed — spec v1 is eager-only; launch geometry varies
-with per-slot pending lengths, needs bucketing by lag-profile) and
-batch the propose pass (one drafter call per slot today). These two go
-together by design: propose is 7 query rows through 5 layers —
-launch-bound, so batching only pays at concurrency — and graph capture
-constrains the same packed geometry, so batching it separately first
-would just get redone when the graph lands.
+## E2e A/B and the round anatomy (2026-08-19)
+
+Full-model EP16 (trays 04–07, `PEGAINFER_K3_MAX_BATCH=16`), vllm-bench
+over SPEED-Bench `throughput_1k` (~1k-token real prompts, seed 42, same
+prompts both legs), main @ 57ff7abe. Artifacts + driver script:
+`~/bench-results/k3-spec-e2e-20260819/`.
+
+| Workload             | spec-on TPOT | spec-off TPOT | verdict |
+|----------------------|--------------|---------------|---------|
+| c1 low_entropy (code)| 44.1 ms      | 76.6 ms       | **1.74× faster** |
+| c1 high_entropy      | 140.3 ms     | 77.0 ms       | **1.82× slower** |
+| c4 mixed (median)    | 51.4 ms      | 78.0 ms       | median hides the tail |
+| c4 mixed (mean)      | 102.6 ms     | 78.9 ms       | 1.3× slower |
+| c4 mixed steady tok/s| 35.7         | 49.1          | **−27% throughput** |
+
+Spec currently wins only on low-entropy single-stream traffic. Plain
+TPOT is category-independent (~77 ms), so the whole spread is
+`round_time / committed_per_round`.
+
+**Round anatomy** (nsys on pruned EP4, tray08 — per-rank layout is
+isomorphic to full EP16; trace `~/k3-prof/specround.nsys-rep`): round
+p50 64.6 ms = verify step 61.4 ms + propose ~3 ms (draft fwd 1.9,
+Markov 0.09, readback+feed 1.1). Inside the median verify step (5039
+kernels): small cuBLAS GEMMs 22.3 ms, MLA attention 15.0 ms (23 calls ×
+0.62 ms), MegaMoE only 9.6 ms, per-row `k3_*_b1` elementwise 7.5 ms,
+FlashKDA 2.8 ms, launch gaps 5.5 ms. Two corrections to the old plan
+fall out of this:
+
+* **The verify step is not launch-bound** — gaps are 9% of the step. A
+  verify CUDA graph buys ≤ ~5 ms; the real cost is small-batch math
+  (per-group GEMMs re-reading weights, per-row elementwise). Packing
+  those across rows is the lever, and it also shrinks what a graph
+  would capture.
+* **Propose is ~5% of the round at c1** — batching/graphing it cannot
+  move single-stream latency. It *does* matter at concurrency, where
+  per-slot propose calls and their per-call pipeline drains stack.
+
+On EP16 the multi-row verify step inflates far worse than EP4 (round
+must be ≥ 140 ms on high-entropy for the TPOT to come out at 140 with
+≥1 commit/round, vs a 77 ms plain step) — the fabric all-to-all pays
+per row, so low-acceptance rounds are double-taxed: fewer commits AND
+a pricier step.
+
+**Landed from this round**: batched propose
+(`feat/k3-batched-propose`) — `decode_spec` now makes one drafter call
+per round over every active slot (sorted `split_at_mut` hands out the
+disjoint `&mut` slot states; drafts un-permute back to batch order for
+the verify split loop). One Markov readback/pipeline drain per round
+instead of one per slot. Semantics-neutral, certified at the right
+granularity: **c1 sequential serve A/B vs main is byte-identical**
+(active=1 hits identical GEMM shapes); at c8 outputs fork in the known
+near-tie noise class — batching the dense draft rows retiles GEMMs, a
+one-draft acceptance shift repacks the verify step, and this
+architecture flips near-tie argmaxes on any geometry change (same class
+#931 certified) — while per-slot acceptance stays in-band (dev 3.59 vs
+main 3.45 mean tokens/round over 8 slots, no collapse → no anchor/state
+cross-wiring). `dspark_reference_cross_check` and the 6 `spec_verify`
+gates are green on the branch build.
+
+**Next, in value order**: (1) pack the verify step's per-row/per-group
+math across rows — biggest single lever on both TPOT columns; (2) size
+the EP16 fabric row-scaling directly (verify-step wall vs rows on the
+fleet) — if confirmed it caps spec at concurrency regardless of (1);
+(3) propose CUDA graph only after (1) fixes the packed geometry it
+would capture.
 
 **Not planned** (decided 2026-08-19): adaptive block length via the
 confidence head. The checkpoint ships a confidence head we don't load;
 using it to truncate low-confidence draft blocks would cut wasted
 verify rows where acceptance is weak (e.g. Chinese at 1.33/round).
-Recorded as an idea only — revisit if low-acceptance traffic shows up
-in real serving profiles.
+Recorded as an idea only — **the revisit condition has now fired**:
+the e2e A/B above shows exactly this failure (high-entropy c1 1.8×
+slower; every wasted draft row is a fabric-taxed verify row). Worth a
+fresh look, but it changes serve behavior, so it stays a decision, not
+a default.

pegainfer-k3/src/executor/mod.rs

@@ -1390,33 +1390,64 @@ impl K3Executor {
             return self.verify_inner(&[]);
         }
         self.enter_step()?;
-        // Propose per slot. The draft is rank-local and collective-free, so
-        // this adds no cross-rank coupling; one call per slot for now (the
-        // batched form needs disjoint `&mut` slot states).
-        let mut verify_batch = Vec::with_capacity(batch.len());
-        for entry in batch {
-            let anchor_pos = self.decode_state.positions[entry.slot];
-            let dspark = self.dspark.as_mut().expect("armed above");
-            let drafts = dspark.model.propose(
-                &self.ctx,
-                &self.model.embed,
-                &self.model.w_lm,
-                &mut [&mut dspark.slots[entry.slot]],
-                &[(entry.last_token, anchor_pos)],
-                &mut dspark.scratch,
-            )?;
-            // Admission reserves `prompt + max_tokens` context, not the
-            // draft span: near the ceiling the verify appends at
-            // `anchor_pos + 1 ..= anchor_pos + drafts` must shed drafts
-            // instead of tripping the verify guard (fatal under EP). A
-            // 0-draft verify is a legal one-token deferred-commit step.
-            let headroom = (self.max_ctx - 1).saturating_sub(anchor_pos);
-            let keep = drafts[0].len().min(headroom);
-            verify_batch.push(K3VerifySlot {
+        let max_ctx = self.max_ctx;
+        let positions: Vec<usize> = batch
+            .iter()
+            .map(|entry| self.decode_state.positions[entry.slot])
+            .collect();
+        // One batched propose for the whole round: the draft is rank-local
+        // and collective-free, the dense draft rows batch across slots, and
+        // the Markov readback becomes a single round trip instead of one per
+        // slot. `propose` wants disjoint `&mut` slot states — a sorted
+        // `split_at_mut` walk over the slot array hands them out.
+        let mut order: Vec<usize> = (0..batch.len()).collect();
+        order.sort_unstable_by_key(|&index| batch[index].slot);
+        let dspark = self.dspark.as_mut().expect("armed above");
+        let mut states: Vec<&mut K3DsparkSlotState> = Vec::with_capacity(batch.len());
+        let mut anchors = Vec::with_capacity(batch.len());
+        let mut rest: &mut [K3DsparkSlotState] = &mut dspark.slots;
+        let mut consumed = 0usize;
+        for &index in &order {
+            let entry = &batch[index];
+            ensure!(
+                entry.slot >= consumed,
+                "K3 decode-spec batch repeats slot {}",
+                entry.slot
+            );
+            let (_, tail) = rest.split_at_mut(entry.slot - consumed);
+            let (state, tail) = tail
+                .split_first_mut()
+                .with_context(|| format!("K3 decode-spec slot {} is out of range", entry.slot))?;
+            states.push(state);
+            anchors.push((entry.last_token, positions[index]));
+            consumed = entry.slot + 1;
+            rest = tail;
+        }
+        let proposed = dspark.model.propose(
+            &self.ctx,
+            &self.model.embed,
+            &self.model.w_lm,
+            &mut states,
+            &anchors,
+            &mut dspark.scratch,
+        )?;
+        // Admission reserves `prompt + max_tokens` context, not the
+        // draft span: near the ceiling the verify appends at
+        // `anchor_pos + 1 ..= anchor_pos + drafts` must shed drafts
+        // instead of tripping the verify guard (fatal under EP). A
+        // 0-draft verify is a legal one-token deferred-commit step.
+        let mut verify_batch: Vec<K3VerifySlot> = batch
+            .iter()
+            .map(|entry| K3VerifySlot {
                 slot: entry.slot,
                 anchor: entry.last_token,
-                drafts: drafts[0][..keep].to_vec(),
-            });
+                drafts: Vec::new(),
+            })
+            .collect();
+        for (&index, drafts) in order.iter().zip(&proposed) {
+            let headroom = (max_ctx - 1).saturating_sub(positions[index]);
+            let keep = drafts.len().min(headroom);
+            verify_batch[index].drafts = drafts[..keep].to_vec();
         }
         // A slot's packed rows are its deferred-commit replay plus the
         // speculative span — up to `2 * K3_DSPARK_BLOCK` — so a full batch can