Fuse DeepSeek MTP verification into decode dispatch

.agents/skills/profile-dsv4-serving-strace/SKILL.md

@@ -79,6 +79,8 @@ summarizes Effective when reprocessing an older log that contains complete devic
 Serving child processes can write several complete `[STRACE]` records on one physical log
 line. The analyzer splits at every marker before calling Simpler's built-in
 `parse_spans`, `group_invocations`, `to_chrome_trace`, and `_round_metrics` APIs.
+It detects both the legacy four-submission decode layout and the current two-submission
+main-plus-verify / MTP layout from the serving kernel spans and final MTP counters.
 
 ## Validate before reporting success
 
@@ -91,8 +93,8 @@ Require all of the following:
 4. `server.log` contains successful SA profiler start/stop, final request, and MTP
    acceptance lines.
 5. Eight distinct `[chip_process pid=... dev=...] ready` mappings are present.
-6. `serving-trace/trace.json` contains non-empty `traceEvents`, including framework and
-   all four DSV4 prefill/decode kernel spans.
+6. `serving-trace/trace.json` contains non-empty `traceEvents`, including framework spans
+   and all prefill/decode kernel spans required by the detected decode layout.
 7. `server.log` contains host `[STRACE]` records and no `clk=dev` records.
 8. `simpler-swimlane.json`, `strace-8lane.json`, and
    `strace-8lane-host-clock.json` contain non-empty `traceEvents`.

.agents/skills/profile-dsv4-serving-strace/scripts/analyze_profile.py

@@ -87,6 +87,32 @@ def main() -> None:
     server_log = server_log_path.read_text(encoding="utf-8", errors="replace")
     serving_trace = json.loads(serving_trace_path.read_text(encoding="utf-8"))
 
+    serving_events = serving_trace["traceEvents"]
+    kernel_durations_ms: dict[str, list[float]] = defaultdict(list)
+    for event in serving_events:
+        if (
+            event.get("ph") == "X"
+            and event.get("cat") == "kernel"
+            and not event.get("name", "").endswith(".worker_run")
+        ):
+            kernel_durations_ms[event["args"]["kernel"]].append(event["dur"] / 1000.0)
+
+    required_kernels = {"deepseek_v4_prefill", "deepseek_v4_mtp_prefill"}
+    if "deepseek_v4_decode_mtp_fused" in kernel_durations_ms:
+        decode_layout = "two_l2"
+        decode_width = 2
+        required_kernels.add("deepseek_v4_decode_mtp_fused")
+    elif {"deepseek_v4_decode", "deepseek_v4_mtp_decode"}.issubset(kernel_durations_ms):
+        decode_layout = "split"
+        decode_width = 4
+        required_kernels.update({"deepseek_v4_decode", "deepseek_v4_mtp_decode"})
+    else:
+        raise RuntimeError(
+            "serving trace has neither the fused nor split DeepSeek V4 decode kernel layout"
+        )
+    if not required_kernels.issubset(kernel_durations_ms):
+        raise RuntimeError(f"missing serving kernel spans: {required_kernels - kernel_durations_ms.keys()}")
+
     split_log = server_log.replace("[STRACE]", "\n[STRACE]")
     spans = list(parse_spans(split_log.splitlines()))
     invocations = group_invocations(spans)
@@ -130,11 +156,11 @@ def main() -> None:
     invocation_ids = sorted(hid_by_inv)
     if (
         invocation_ids != list(range(1, invocation_ids[-1] + 1))
-        or invocation_ids[-1] < 8
-        or (invocation_ids[-1] - 4) % 4
+        or invocation_ids[-1] < 4 + decode_width
+        or (invocation_ids[-1] - 4) % decode_width
     ):
         raise RuntimeError(f"unexpected invocation ids: {invocation_ids}")
-    decode_step_count = (invocation_ids[-1] - 4) // 4
+    decode_step_count = (invocation_ids[-1] - 4) // decode_width
     if len(request_invocations) != len(devices) * invocation_ids[-1]:
         raise RuntimeError(
             f"incomplete rank data: got {len(request_invocations)} invocations, "
@@ -177,10 +203,13 @@ def rank_phase_row(pid: int, main_ids: list[int], mtp_ids: list[int]) -> dict:
     decode_rows = []
     critical_metric = "total_effective_us" if device_effective_available else "total_host_us"
     for step in range(decode_step_count):
-        base = 5 + step * 4
+        base = 5 + step * decode_width
         per_rank = []
         for pid in pids:
-            per_rank.append(rank_phase_row(pid, [base, base + 1], [base + 2, base + 3]))
+            if decode_layout == "two_l2":
+                per_rank.append(rank_phase_row(pid, [base], [base + 1]))
+            else:
+                per_rank.append(rank_phase_row(pid, [base, base + 1], [base + 2, base + 3]))
         critical = max(per_rank, key=lambda row: row[critical_metric])
         decode_row = {
             "step": step + 1,
@@ -200,25 +229,6 @@ def rank_phase_row(pid: int, main_ids: list[int], mtp_ids: list[int]) -> dict:
             )
         decode_rows.append(decode_row)
 
-    serving_events = serving_trace["traceEvents"]
-    kernel_durations_ms: dict[str, list[float]] = defaultdict(list)
-    for event in serving_events:
-        if (
-            event.get("ph") == "X"
-            and event.get("cat") == "kernel"
-            and not event.get("name", "").endswith(".worker_run")
-        ):
-            kernel_durations_ms[event["args"]["kernel"]].append(event["dur"] / 1000.0)
-
-    required_kernels = {
-        "deepseek_v4_prefill",
-        "deepseek_v4_mtp_prefill",
-        "deepseek_v4_decode",
-        "deepseek_v4_mtp_decode",
-    }
-    if not required_kernels.issubset(kernel_durations_ms):
-        raise RuntimeError(f"missing serving kernel spans: {required_kernels - kernel_durations_ms.keys()}")
-
     request_ms = next(
         (
             event["dur"] / 1000.0
@@ -287,6 +297,7 @@ def kernel_summary(name: str) -> dict:
     summary = {
         "run_id": args.run_id,
         "devices": devices,
+        "decode_layout": decode_layout,
         "request": {
             "prompt_tokens": int(completion_match.group(1)),
             "completion_tokens": completion_tokens,
@@ -325,10 +336,29 @@ def kernel_summary(name: str) -> dict:
     (artifact_dir / "profile-summary.json").write_text(json.dumps(summary, indent=2) + "\n")
 
     prefill = summary["simpler"]["prefill_critical"]
-    main_serving = summary["serving_kernel_ms"]["deepseek_v4_decode"]["steady_after_first"]
-    mtp_serving = summary["serving_kernel_ms"]["deepseek_v4_mtp_decode"]["steady_after_first"]
-    if main_serving is None or mtp_serving is None:
-        raise RuntimeError("need more than one decode kernel span for steady statistics")
+    if decode_layout == "two_l2":
+        fused_serving = summary["serving_kernel_ms"]["deepseek_v4_decode_mtp_fused"][
+            "steady_after_first"
+        ]
+        if fused_serving is None:
+            raise RuntimeError("need more than one fused decode kernel span for steady statistics")
+        serving_decode_lines = (
+            f"- Combined decode steady mean (steps 2-{decode_step_count}): "
+            f"{fused_serving['mean']:.3f} ms/iteration"
+        )
+        main_phase_label = "Main+verify"
+    else:
+        main_serving = summary["serving_kernel_ms"]["deepseek_v4_decode"]["steady_after_first"]
+        mtp_serving = summary["serving_kernel_ms"]["deepseek_v4_mtp_decode"]["steady_after_first"]
+        if main_serving is None or mtp_serving is None:
+            raise RuntimeError("need more than one decode kernel span for steady statistics")
+        serving_decode_lines = (
+            f"- Decode main steady mean (steps 2-{decode_step_count}): "
+            f"{main_serving['mean']:.3f} ms/iteration\n"
+            f"- Decode MTP steady mean (steps 2-{decode_step_count}): "
+            f"{mtp_serving['mean']:.3f} ms/iteration"
+        )
+        main_phase_label = "Main"
     critical_rank_counts = Counter(row["critical_device"] for row in decode_rows)
     critical_rank_summary = ", ".join(
         f"device {device}: {count}/{decode_step_count}"
@@ -366,9 +396,9 @@ def kernel_summary(name: str) -> dict:
 Sched windows.
 
 - Prefill critical rank: device {prefill["device"]}, main={prefill["main_effective_us"] / 1000:.3f} ms, MTP={prefill["mtp_effective_us"] / 1000:.3f} ms, total={prefill["total_effective_us"] / 1000:.3f} ms
-- Decode steady critical rank mean: main={steady_effective["critical_main_effective_us"]["mean"] / 1000:.3f} ms, MTP={steady_effective["critical_mtp_effective_us"]["mean"] / 1000:.3f} ms, total={steady_effective["critical_total_effective_us"]["mean"] / 1000:.3f} ms/iteration
+- Decode steady critical rank mean: {main_phase_label}={steady_effective["critical_main_effective_us"]["mean"] / 1000:.3f} ms, MTP={steady_effective["critical_mtp_effective_us"]["mean"] / 1000:.3f} ms, total={steady_effective["critical_total_effective_us"]["mean"] / 1000:.3f} ms/iteration
 
-| Decode iteration | Critical device | Main Effective (ms) | MTP Effective (ms) | Total Effective (ms) |
+| Decode iteration | Critical device | {main_phase_label} Effective (ms) | MTP Effective (ms) | Total Effective (ms) |
 | ---: | ---: | ---: | ---: | ---: |
 {effective_decode_table}
 """
@@ -384,16 +414,15 @@ def kernel_summary(name: str) -> dict:
 
 - Prefill main kernel span: {kernel_durations_ms["deepseek_v4_prefill"][0]:.3f} ms
 - Prefill MTP kernel span: {kernel_durations_ms["deepseek_v4_mtp_prefill"][0]:.3f} ms
-- Decode main steady mean (steps 2-{decode_step_count}): {main_serving["mean"]:.3f} ms/iteration
-- Decode MTP steady mean (steps 2-{decode_step_count}): {mtp_serving["mean"]:.3f} ms/iteration
+{serving_decode_lines}
 
 ## Simpler Host STRACE
 
 - Prefill critical rank: device {prefill["device"]}, main={prefill["main_host_us"] / 1000:.3f} ms, MTP={prefill["mtp_host_us"] / 1000:.3f} ms, total={prefill["total_host_us"] / 1000:.3f} ms
-- Decode steady critical rank mean (steps 2-{decode_step_count}): main={host_stats["critical_main_host_us"]["mean"] / 1000:.3f} ms, MTP={host_stats["critical_mtp_host_us"]["mean"] / 1000:.3f} ms, total={host_stats["critical_total_host_us"]["mean"] / 1000:.3f} ms/iteration
+- Decode steady critical rank mean (steps 2-{decode_step_count}): {main_phase_label}={host_stats["critical_main_host_us"]["mean"] / 1000:.3f} ms, MTP={host_stats["critical_mtp_host_us"]["mean"] / 1000:.3f} ms, total={host_stats["critical_total_host_us"]["mean"] / 1000:.3f} ms/iteration
 - Critical-rank counts across decode: {critical_rank_summary}
 
-| Decode iteration | Critical device | Main host (ms) | MTP host (ms) | Total host (ms) |
+| Decode iteration | Critical device | {main_phase_label} host (ms) | MTP host (ms) | Total host (ms) |
 | ---: | ---: | ---: | ---: | ---: |
 {host_decode_table}
 

.agents/skills/profile-dsv4-serving-strace/scripts/render_8lane.py

@@ -18,9 +18,10 @@
 
 PROCESS_NAME_RE = re.compile(r"inv=(?P<inv>\d+) \(pid=(?P<pid>\d+)\)")
 DEVICE_READY_RE = re.compile(r"\[chip_process pid=(?P<pid>\d+) dev=(?P<device>\d+)\] ready")
+MTP_ACCEPTANCE_RE = re.compile(r"MTP acceptance for .* proposed=(?P<steps>\d+)")
 
 
-def callable_label(invocation: int) -> tuple[str, int | None, str]:
+def callable_label(invocation: int, decode_width: int) -> tuple[str, int | None, str]:
     prefill = {
         1: ("prefill.main", None, "rail_response"),
         2: ("prefill.main.lm_head", None, "rail_animation"),
@@ -29,8 +30,15 @@ def callable_label(invocation: int) -> tuple[str, int | None, str]:
     }
     if invocation in prefill:
         return prefill[invocation]
-    step = (invocation - 5) // 4 + 1
-    phase = (invocation - 5) % 4
+    step = (invocation - 5) // decode_width + 1
+    phase = (invocation - 5) % decode_width
+    if decode_width == 2:
+        decode = {
+            0: ("decode.main+verify", "good"),
+            1: ("decode.mtp", "cq_build_running"),
+        }
+        label, color = decode[phase]
+        return label, step, color
     decode = {
         0: ("decode.main", "good"),
         1: ("decode.main.lm_head", "rail_animation"),
@@ -91,9 +99,10 @@ def main() -> None:
     source = json.loads(args.input.read_text())
     source_events = source["traceEvents"] if isinstance(source, dict) else source
 
+    server_log = args.server_log.read_text(errors="replace")
     pid_to_device = {
         int(match.group("pid")): int(match.group("device"))
-        for match in DEVICE_READY_RE.finditer(args.server_log.read_text(errors="replace"))
+        for match in DEVICE_READY_RE.finditer(server_log)
     }
     devices = sorted(pid_to_device.values())
     if len(devices) != 8 or len(set(devices)) != 8:
@@ -119,6 +128,19 @@ def main() -> None:
         if virtual_pid in virtual_processes and event.get("ph") == "X":
             grouped.setdefault(int(virtual_pid), []).append(event)
 
+    invocation_ids = sorted({invocation for _device, invocation in virtual_processes.values()})
+    if not invocation_ids or invocation_ids != list(range(1, invocation_ids[-1] + 1)):
+        raise ValueError(f"unexpected invocation ids: {invocation_ids}")
+    decode_invocations = invocation_ids[-1] - 4
+    acceptance_matches = list(MTP_ACCEPTANCE_RE.finditer(server_log))
+    proposed_steps = int(acceptance_matches[-1].group("steps")) if acceptance_matches else 0
+    if proposed_steps and decode_invocations == proposed_steps * 2:
+        decode_width = 2
+    elif decode_invocations % 4 == 0:
+        decode_width = 4
+    else:
+        raise ValueError(f"cannot infer decode invocation width from {invocation_ids[-1]} invocations")
+
     roots = [
         event
         for events in grouped.values()
@@ -169,7 +191,7 @@ def main() -> None:
         runner = one_event(events, "simpler_run.runner_run")
         validate = one_event(events, "simpler_run.validate")
         device_wall = one_event(events, "simpler_run.runner_run.device_wall")
-        label, step, color = callable_label(invocation)
+        label, step, color = callable_label(invocation, decode_width)
         event_name = label if step is None else f"D{step:02d} {label}"
         has_device_trace = device_wall is not None
         event_args = {

pypto_serving/model/deepseek/npu_executor.py

@@ -102,12 +102,14 @@
     "decode_attention_hca",
     "decode_attention_swa",
     "decode_fwd",
+    "decode_fwd_mtp",
     "decode_input_pack",
     "decode_indexer",
     "decode_indexer_compressor",
     "decode_layer",
     "decode_metadata_device",
     "decode_mtp",
+    "decode_mtp_verify",
     "lookup_embedding",
     "decode_sparse_attn",
     "decode_sparse_attn_csa",
@@ -438,9 +440,21 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels:
                 self._prefill_dummy_args(model, layout, modules["config"]),
             )
             decode = self._compile_l3_callable(
-                "deepseek_v4_decode",
-                modules["decode_fwd"].l3_decode_fwd,
-                self._decode_dummy_args(model, layout, modules["config"]),
+                "deepseek_v4_decode_mtp_fused" if self._enable_mtp else "deepseek_v4_decode",
+                (
+                    modules["decode_fwd_mtp"].l3_decode_fwd_mtp
+                    if self._enable_mtp
+                    else modules["decode_fwd"].l3_decode_fwd
+                ),
+                (
+                    self._fused_mtp_dummy_args(
+                        modules,
+                        model=model,
+                        layout=layout,
+                    )
+                    if self._enable_mtp
+                    else self._decode_dummy_args(model, layout, modules["config"])
+                ),
             )
             if self._enable_mtp:
                 mtp_prefill = self._compile_l3_callable(
@@ -453,16 +467,6 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels:
                         num_tokens=layout.prefill_seq,
                     ),
                 )
-                mtp_decode = self._compile_l3_callable(
-                    "deepseek_v4_mtp_decode",
-                    modules["decode_mtp"].l3_mtp_decode_layer,
-                    self._mtp_dummy_args(
-                        modules["decode_mtp"],
-                        model=model,
-                        layout=layout,
-                        num_tokens=layout.decode_tokens,
-                    ),
-                )
             freqs_cos, freqs_sin = self._build_rope_tables(modules["rope_tables"], modules["config"])
 
         return DeepSeekV4CompiledKernels(
@@ -488,6 +492,54 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels:
             enable_mtp=self._enable_mtp,
         )
 
+    def _fused_mtp_dummy_args(
+        self,
+        modules: dict[str, object],
+        *,
+        model: RuntimeModel,
+        layout: DeepSeekV4CacheLayout,
+    ) -> tuple[Any, ...]:
+        """Build the combined main-decode and MTP-decode compile signature."""
+        main_args = self._decode_dummy_args(model, layout, modules["config"])
+        mtp_args = self._mtp_dummy_args(
+            modules["decode_mtp"],
+            model=model,
+            layout=layout,
+            num_tokens=layout.decode_tokens,
+        )
+        with _deepseek_v4_import_context(
+            self._kernel_dir,
+            pypto_root=self._kernel_dir.parents[2],
+            ep=len(self._device_ids),
+            lm_head_tp=DEEPSEEK_V4_LM_HEAD_TP_SIZE,
+            moe_shape="decode",
+        ):
+            mtp_specs = modules["decode_mtp"].build_tensor_specs(
+                num_tokens=layout.decode_tokens,
+            )
+        shared_names = {
+            "embed_weight",
+            "main_pre_hc_hidden",
+            "freqs_cos",
+            "freqs_sin",
+            "ori_block_table",
+            "lm_head_weight",
+        }
+        tail_token_ids = torch.empty(
+            (layout.ranks, layout.decode_batch),
+            dtype=torch.int64,
+        )
+        tail_positions = torch.empty(
+            (layout.ranks, layout.decode_batch),
+            dtype=torch.int32,
+        )
+        fused_mtp_args = tuple(
+            arg
+            for spec, arg in zip(mtp_specs, mtp_args, strict=True)
+            if spec.name not in shared_names
+        )
+        return (*main_args, tail_token_ids, tail_positions, *fused_mtp_args)
+
     def _load_kernel_modules(self, layout: DeepSeekV4CacheLayout) -> dict[str, object]:
         """Import DeepSeekV4 pypto-lib modules with EP fixed to the serving world size."""
         pypto_root = self._kernel_dir.parents[2]
@@ -522,10 +574,13 @@ def _load_kernel_modules(self, layout: DeepSeekV4CacheLayout) -> dict[str, objec
             config.DECODE_RECV_MAX = ranks * layout.decode_tokens
             config.RECV_MAX = config.DECODE_RECV_MAX
             modules = {"config": config}
+            decode_module_names = ["decode_layer", "decode_fwd", "lm_head", "rope_tables"]
+            if self._enable_mtp:
+                decode_module_names.extend(("decode_mtp", "decode_fwd_mtp"))
             modules.update(
                 {
                     name: importlib.import_module(name)
-                    for name in ("decode_layer", "decode_fwd", "decode_mtp", "lm_head", "rope_tables")
+                    for name in decode_module_names
                 }
             )
         modules["prefill_layer"] = prefill_layer