Преглед на файлове

glm53-flash-dflash2: target-side EAGLE3 port + rigorous KV-page blocker writeup

- patch_eagle3.py: port the fork's SupportsEagle3/EagleModelMixin target
  support (aux hidden states at dflash target_layer_ids 5/14/24/33/42,
  hc_contract-ed for dflash) onto our day-0 model.py surgically; full file
  swap impossible (fork diverges ~150 lines into b12x-only paths).
- Verified via boot: spec JSON parsed, aux layers picked from draft
  config, DFlash2 draft loaded, MXFP8 draft kernel bound.
- KV setup then raises: draft per-token page (bf16 4096 B / fp8 2048 B,
  8 kv heads) is never divisible with the fp8_ds_mla target pages
  (MLA 660 B, kpool indexer 132 B); integer block-scaling cannot bridge
  it and --disable-hybrid-kv-cache-manager only trades the error for
  "cannot unify types" (indexer spec). The fork's sm_120 run escapes via
  b12x stride-indexed pages (indexes_kv_by_block_stride), which do not
  exist on the GB10 base. Recipe documents the math + eliminated knobs
  + three paths forward; args reverted to card-faithful values.
kilork преди 1 месец
родител
ревизия
848772a3e7
променени са 3 файла, в които са добавени 222 реда и са изтрити 9 реда
  1. 3 2
      docker-images/glm53-flash-dflash2/Dockerfile
  2. 182 0
      docker-images/glm53-flash-dflash2/patch_eagle3.py
  3. 37 7
      recipes/glm53-flash-dflash2-mxfp8/recipe.toml

+ 3 - 2
docker-images/glm53-flash-dflash2/Dockerfile

@@ -19,7 +19,7 @@
 ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
-COPY qwen3_dflash.py qwen3_dflash2.py dflash.py /tmp/
+COPY qwen3_dflash.py qwen3_dflash2.py dflash.py patch_eagle3.py /tmp/
 
 RUN set -eux; \
     B=/usr/local/lib/python3.12/dist-packages/vllm; \
@@ -35,4 +35,5 @@ RUN set -eux; \
         "$B/model_executor/models/qwen3_dflash2.py" \
         "$B/v1/spec_decode/dflash.py" \
         "$B/model_executor/models/registry.py"; \
-    rm -f /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py /tmp/dflash.py
+    python3 /tmp/patch_eagle3.py; \
+    rm -f /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py /tmp/dflash.py /tmp/patch_eagle3.py

+ 182 - 0
docker-images/glm53-flash-dflash2/patch_eagle3.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Target-side EAGLE3/aux-hidden-state support for Glm5Next (DFlash2 needs it).
+
+The DFlash2 drafter consumes the TARGET's hidden states at layers
+5/14/24/33/42 (draft config.json dflash_config.target_layer_ids). The vLLM
+spec-decode runner gates that on supports_eagle3(target) and on the model
+forward emitting aux states (EagleModelMixin). The day-0 nvidia Glm5Next
+model implements neither; the fork (local-inference-lab/vllm
+dev/jovian-judgement @ da4d7be6, Apache-2.0, see UPSTREAM.md) does. A full
+file swap is NOT possible: the fork diverges from our day-0 base in ~150
+lines (b12x backends, weight remaps) and our base carries fixes the fork
+lacks. So this is a surgical port of exactly the fork's:
+
+  * EagleModelMixin on Glm5NextModel + SupportsEagle3 on both roots
+  * dflash_capture flag (fork: speculative_config.use_dflash(), absent on
+    our base -> method == "dflash" equivalent)
+  * _prepare_aux_hidden_state (mHC: hc_post then hc_contract for dflash)
+  * forward-loop aux collection + tuple return (after self.norm, fork parity)
+
+Idempotent; every anchor asserted exactly-once; py_compile-checked.
+"""
+
+import py_compile
+
+B = "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia"
+PATH = f"{B}/model.py"
+
+EDITS = [
+    # 1. imports
+    (
+        """from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    SupportsPP,
+)""",
+        """from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    SupportsEagle3,
+    SupportsPP,
+)""",
+    ),
+    # 2. model gets the mixin (runner sets aux layers through it)
+    (
+        "class Glm5NextModel(nn.Module):\n",
+        "class Glm5NextModel(nn.Module, EagleModelMixin):\n",
+    ),
+    # 3. dflash capture flag in Glm5NextModel.__init__
+    (
+        """        self.is_sequence_parallel = (
+            vllm_config.parallel_config.use_sequence_parallel_moe
+        )
+""",
+        """        self.is_sequence_parallel = (
+            vllm_config.parallel_config.use_sequence_parallel_moe
+        )
+        _spec = vllm_config.speculative_config
+        # fork parity: dflash drafts consume hc_contract-ed aux states
+        # (SpeculativeConfig.use_dflash() only exists on the fork).
+        self.dflash_capture = (
+            _spec is not None and _spec.method == "dflash"
+        )
+""",
+    ),
+    # 4. aux extraction method (fork body, verbatim apart from type-name quote)
+    (
+        """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+""",
+        """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def _prepare_aux_hidden_state(
+        self,
+        layer: "Glm5NextDecoderLayer",
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        post: torch.Tensor | None,
+        comb: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if not layer.mhc or residual is None:
+            return hidden_states
+        assert post is not None and comb is not None
+        aux_hidden_state = layer.hc_post(hidden_states, residual, post, comb)
+        if self.dflash_capture:
+            return hc_contract(aux_hidden_state, layer.n)
+        return aux_hidden_state.flatten(1)
+""",
+    ),
+    # 5. forward return type
+    (
+        """    ) -> torch.Tensor:
+        if get_pp_group().is_first_rank:
+""",
+        """    ) -> (
+        torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]
+    ):
+        if get_pp_group().is_first_rank:
+""",
+    ),
+    # 6. aux collection in the layer loop (fork semantics: capture the state
+    # ENTERING layer_idx+1, i.e. the OUTPUT of layer_idx)
+    (
+        """        for layer in self._active_layers:
+            hidden_states, residual, post, comb = layer(
+                positions, hidden_states, residual, post, comb
+            )
+""",
+        """        aux_hidden_states: list[torch.Tensor] = []
+        if self.start_layer in self.aux_hidden_state_layers:
+            aux_hidden_states.append(hidden_states)
+
+        for layer_idx, layer in enumerate(
+            self._active_layers, start=self.start_layer
+        ):
+            if layer_idx + 1 in self.aux_hidden_state_layers:
+                aux_hidden_state = self._prepare_aux_hidden_state(
+                    layer, hidden_states, residual, post, comb
+                )
+                if self.is_sequence_parallel:
+                    aux_hidden_state = sp_all_gather(aux_hidden_state)[
+                        :full_num_tokens
+                    ]
+                aux_hidden_states.append(aux_hidden_state)
+            hidden_states, residual, post, comb = layer(
+                positions, hidden_states, residual, post, comb
+            )
+""",
+    ),
+    # 7. tuple return after norm (fork parity — logits stay well-formed)
+    (
+        """        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+""",
+        """        hidden_states = self.norm(hidden_states)
+        if aux_hidden_states:
+            return hidden_states, aux_hidden_states
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+""",
+    ),
+    # 8. both registered roots advertise eagle3 (the runner checks the root)
+    (
+        """    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
+):
+""",
+        """    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid,
+    SupportsEagle3
+):
+""",
+    ),
+    (
+        """    Glm4vForConditionalGeneration, HasInnerState, IsHybrid
+):
+""",
+        """    Glm4vForConditionalGeneration, HasInnerState, IsHybrid, SupportsEagle3
+):
+""",
+    ),
+]
+
+src = open(PATH).read()
+if "dflash_capture" in src:
+    print(f"already patched: {PATH}")
+else:
+    for i, (old, new) in enumerate(EDITS, 1):
+        n = src.count(old)
+        assert n == 1, f"edit {i}: anchor {n}x in {PATH} — review me"
+        src = src.replace(old, new, 1)
+    open(PATH, "w").write(src)
+    py_compile.compile(PATH, doraise=True)
+    print(f"patched: {PATH}")
+
+print("patch_eagle3 OK")

+ 37 - 7
recipes/glm53-flash-dflash2-mxfp8/recipe.toml

@@ -32,13 +32,42 @@
 # Adopted verbatim from the card: --speculative-config {method:dflash,
 # model:...DFlash2-MXFP8, num_speculative_tokens:7, kv_cache_dtype:auto}.
 #
-# Boot markers: target loads as in glm53-flash-nvfp4-lab (92.72 GiB/node) →
-# a "DFlash2" speculator/drafter line in the spec-decode init → "Application
-# startup complete". A "DFlash2DraftModel not in registry" error means the
-# image overlay didn't apply (rebuild glm53-flash-dflash2:local). Verify the
-# DFlash proposal acceptance rate is > the lab MTP-3 baseline (that's the
-# whole point of DFlash2); if accuracy looks off, this recipe is the eager
-# control.
+# =====================================================================
+# STATUS 2026-08-29: BLOCKED on GB10 — boots to weight-load + draft load
+# (both fine, MXFP8 draft kernel binds), then dies in KV-cache setup:
+#
+#   NotImplementedError: layers.3.self_attn.indexer.k_cache: page size is
+#   not divisible by the maximum page size and cannot be padded
+#   (indexes_kv_by_block_stride is False).
+#
+# MATH (block size cancels, so no --block-size / draft kv dtype fixes it):
+# per-token pages on this base are indexer 132 B, MLA fp8_ds_mla 660 B,
+# DFlash2 draft (8 kv heads x 128) bf16 4096 B / fp8 2048 B. vLLM unifies
+# target+draft pages by integer block-size scaling, which needs divisio-
+# nility: 4096%660 != 0, 2048%660 != 0. Tried and eliminated: draft
+# kv_cache_dtype fp8, target kv-cache-dtype variants, --disable-hybrid-
+# kv-cache-manager (then: "failed to convert the KV cache specs to one
+# unified type" — the kpool indexer spec is not coercible either).
+#
+# WHY THE FORK DOES NOT HIT THIS: its qualified run (voipmonitor image,
+# sm_120) uses B12X attention/indexer backends which opt into
+# indexes_kv_by_block_stride, so unify PADS the strided pages instead.
+# The b12x kernels are not part of the arm64 GB10 base and are compiled
+# for sm_120 consumer Blackwell.
+#
+# PATHS FORWARD (pick before spending boot time on this again):
+#   a) build local-inference-lab/vllm dev/jovian-judgement for arm64/cu130
+#      incl. b12x kernels (hours; sm_120->sm_121 compat unknown);
+#   b) wait for a GB10 indexer/MLA backend that sets
+#      indexes_kv_by_block_stride (upstream);
+#   c) give up — acceptance-rate comparison falls back to MTP (already
+#      characterized in specs/tech-architecture/glm53-recipe-state.md).
+#
+# What IS done and committed regardless: glm53-flash-dflash2:local image
+# (DFlash2DraftModel port + registry), patch_eagle3.py (target-side
+# EAGLE3/aux-hidden-state support for Glm5Next, fork parity), verified
+# config plumbing (spec JSON parsed, aux layers 5/14/24/33/42 extracted
+# from the draft config, FlashInferCutlassMxfp8LinearKernel bound).
 #
 # PREREQUISITES: BOTH models cached on both nodes:
 #   hf download local-inference-lab/GLM-5.3-Flash-NVFP4
@@ -79,6 +108,7 @@ args = [
     "--disable-custom-all-reduce",
     # DFlash2 draft from local-inference-lab/GLM-5.3-Flash-DFlash2-MXFP8
     # (block-diffusion drafter, proposal length 7 — verbatim from the card)
+    # Card: kv_cache_dtype auto (bf16 draft pages) — kept faithful.
     "--speculative-config", """
 '{"method":"dflash","model":"local-inference-lab/GLM-5.3-Flash-DFlash2-MXFP8","num_speculative_tokens":7,"kv_cache_dtype":"auto"}'""",
     "--language-model-only",