docker-images: glm53-flash-dflash2

docker-images/glm53-flash-dflash2/Dockerfile

@@ -0,0 +1,39 @@
+# syntax=docker/dockerfile:1.7
+# GLM-5.3-Flash DFlash2 draft-model support for DGX Spark (SM121/GB10), TP=2.
+#
+# Stacks on glm53-flash-nvfp4-lab:local (day-0 base + NoPE-MLA + lab
+# checkpoint patches) and adds DFlash2 draft-model execution ported from
+# local-inference-lab/vllm `dev/jovian-judgement` (see UPSTREAM.md):
+#   * qwen3_dflash.py   — refreshed base (fork is a superset of the day-0
+#                         file; identical imports, +DFlash2 wiring incl.
+#                         native MXFP8 serialized projections)
+#   * qwen3_dflash2.py  — the DFlash2DraftModel itself (294 lines)
+#   * dflash.py         — v1 spec-decode proposer (+rope is_neox_style
+#                         handoff from the target, 1 hunk)
+#   * registry.py       — one line: DFlash2DraftModel -> qwen3_dflash2
+#
+# The day-0 base already ships W4A16_NVFP4, MIXED_PRECISION,
+# ModelOptMxFp8LinearMethod, CompressedTensorsW8A8Mxfp8 and the dflash
+# proposer — only the DFlash2 model class + registration were missing.
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+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; \
+    cp /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py "$B/model_executor/models/"; \
+    cp /tmp/dflash.py "$B/v1/spec_decode/dflash.py"; \
+    grep -q '"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),' \
+        "$B/model_executor/models/registry.py"; \
+    sed -i '/"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),/a\    "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"),' \
+        "$B/model_executor/models/registry.py"; \
+    grep -q '"DFlash2DraftModel"' "$B/model_executor/models/registry.py"; \
+    python3 -m py_compile \
+        "$B/model_executor/models/qwen3_dflash.py" \
+        "$B/model_executor/models/qwen3_dflash2.py" \
+        "$B/v1/spec_decode/dflash.py" \
+        "$B/model_executor/models/registry.py"; \
+    python3 /tmp/patch_eagle3.py; \
+    rm -f /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py /tmp/dflash.py /tmp/patch_eagle3.py

docker-images/glm53-flash-dflash2/UPSTREAM.md

@@ -0,0 +1,26 @@
+# Upstream provenance — vendored DFlash2 sources
+
+All three `.py` files are copied verbatim (no local edits) from
+
+- repo:   https://github.com/local-inference-lab/vllm
+- branch: `dev/jovian-judgement`
+- commit: `da4d7be6c97434f6942292ed8abbf4b32dc44355` (2026-08-28,
+  "feat(dflash2): execute serialized MXFP8 projections natively")
+- license: Apache-2.0 (SPDX headers retained)
+
+| local file     | upstream path                                       |
+|----------------|-----------------------------------------------------|
+| qwen3_dflash.py  | vllm/model_executor/models/qwen3_dflash.py        |
+| qwen3_dflash2.py | vllm/model_executor/models/qwen3_dflash2.py       |
+| dflash.py        | vllm/v1/spec_decode/dflash.py                     |
+
+Why vendored instead of "just use the fork image": every published
+`local-inference-lab` / `voipmonitor` vLLM image is **amd64-only** (RTX PRO
+6000 TP=4 world). The day-0 arm64 cu130 base our lab image uses is already a
+superset of the fork for everything EXCEPT the DFlash2 model class — verified
+by import-level diffing: the fork files import only symbols the day-0 image
+ships (`ModelOptMxFp8LinearMethod`, `CompressedTensorsW8A8Mxfp8`, the dflash
+proposer machinery). So the whole port is these 3 files + one registry line.
+
+Refresh: re-fetch the same three paths from the branch tip, `diff` against
+the vendored copies, rebuild.

docker-images/glm53-flash-dflash2/build.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Build the SM121 GLM-5.3-Flash DFlash2 image: lab image + DFlash2 draft-model
+# port from local-inference-lab/vllm dev/jovian-judgement (see UPSTREAM.md).
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a --build-arg
+# to the Dockerfile. TAG is set by the sm121 launcher to the recipe's
+# docker_image (manual build: TAG=glm53-flash-dflash2:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-glm53-flash-dflash2:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    val="${val%\"}"; val="${val#\"}"   # versions.env values may be shell-quoted
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+# Persistent buildkit cache (per node, shared across image dirs): the
+# version-independent stages (apt, torch, CUDA base) are reused across
+# rebuilds — only the pinned source stages (vllm/b12x/…) re-download and
+# re-compile. Override the dir with SM121_BUILDKIT_CACHE if desired.
+CACHE_DIR="${SM121_BUILDKIT_CACHE:-$HOME/.cache/sm121-buildkit}"
+mkdir -p "${CACHE_DIR}"
+docker buildx build "${ARGS[@]}" --cache-from "type=local,src=${CACHE_DIR}" \
+    --cache-to "type=local,dest=${CACHE_DIR},mode=max" \
+    --progress=plain --load -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

docker-images/glm53-flash-dflash2/dflash.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import replace
+from typing import Any
+
+import torch
+from typing_extensions import override
+
+from vllm.config import VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
+from vllm.v1.spec_decode.utils import (
+    copy_and_expand_dflash_inputs_kernel,
+    next_power_of_2,
+)
+
+logger = init_logger(__name__)
+
+
+class DFlashProposer(SpecDecodeBaseProposer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        device: torch.device,
+        runner=None,
+    ):
+        assert vllm_config.speculative_config is not None
+        assert vllm_config.speculative_config.method == "dflash"
+        super().__init__(
+            vllm_config=vllm_config,
+            device=device,
+            pass_hidden_states_to_model=True,
+            runner=runner,
+        )
+
+        # Only next_token_ids and mask tokens are query tokens, all other context is K/V
+        self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens)
+        self.max_padded_query_tokens = max(
+            self.max_query_tokens,
+            vllm_config.compilation_config.max_cudagraph_capture_size or 0,
+        )
+        # Positions covers both context states + query states
+        self.max_positions = self.max_num_tokens + self.max_padded_query_tokens
+
+        # Separate context buffers to keep query buffer addresses stable for CUDA graphs
+        self._context_slot_mapping_buffer = torch.zeros(
+            self.max_num_tokens,
+            dtype=torch.int64,
+            device=device,
+        )
+        self._slot_mapping_buffer = torch.zeros(
+            self.max_padded_query_tokens,
+            dtype=torch.int64,
+            device=device,
+        )
+        self._context_positions_buffer = torch.zeros(
+            self.max_num_tokens,
+            dtype=torch.int64,
+            device=device,
+        )
+        self.positions = torch.zeros(
+            self.max_padded_query_tokens,
+            dtype=torch.int64,
+            device=device,
+        )
+
+        self.arange = torch.arange(
+            self.max_positions + 1, device=device, dtype=torch.int32
+        )
+
+        # For DFlash we use the input embeddings to embed the mask token
+        self.parallel_drafting_hidden_state_tensor = None
+
+        from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
+
+        self.dflash_causal = not dflash_has_any_non_causal(
+            self.draft_model_config.hf_config
+        )
+
+    @override
+    def load_model(self, target_model: torch.nn.Module) -> None:
+        from vllm.model_executor.models.qwen3_dflash import (
+            dflash_target_rope_is_neox_style,
+        )
+
+        draft_model_config = self.speculative_config.draft_model_config
+        assert draft_model_config is not None
+        # The drafter must rotate Q/K the way its target does. Take that from the
+        # built target before super() constructs the draft.
+        is_neox_style = dflash_target_rope_is_neox_style(target_model)
+        if is_neox_style is not None:
+            draft_model_config.hf_config.is_neox_style = is_neox_style
+        super().load_model(target_model)
+
+    @override
+    def _create_draft_vllm_config(self) -> VllmConfig:
+        base = super()._create_draft_vllm_config()
+        # The draft model is text-only — clear the target's multimodal
+        # flag so flash_attn is not rejected for mm_prefix support.
+        arch = base.model_config.model_arch_config
+        if arch.is_mm_prefix_lm:
+            base.model_config.model_arch_config = replace(arch, is_mm_prefix_lm=False)
+        return replace(
+            base,
+            attention_config=replace(
+                base.attention_config,
+                use_non_causal=not self.dflash_causal,
+            ),
+        )
+
+    @override
+    def _warn_if_multimodal(self):
+        # Override to allow multimodal inputs since DFlash supports Qwen3.5 models
+        pass
+
+    @override
+    def set_inputs_first_pass(
+        self,
+        target_token_ids: torch.Tensor,
+        next_token_ids: torch.Tensor,
+        target_positions: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        token_indices_to_sample: torch.Tensor | None,
+        cad: CommonAttentionMetadata,
+        num_rejected_tokens_gpu: torch.Tensor | None,
+    ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
+        # DFlash cross-attention: context K/V from target hidden states,
+        # Q from query embeddings (bonus + mask tokens).
+        batch_size = cad.batch_size()
+        num_context = target_token_ids.shape[0]
+        num_query_per_req = 1 + self.num_speculative_tokens
+        num_query_total = batch_size * num_query_per_req
+
+        # Store for build_model_inputs_first_pass to use
+        self._dflash_num_context = num_context
+
+        # We don't need to copy into a buffer here since the context preprocessing
+        # does not run in a CUDA graph
+        self._dflash_hidden_states = target_hidden_states
+
+        token_indices_to_sample = torch.empty(
+            batch_size * self.num_speculative_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        # Launch fused triton kernel for input_ids, positions, slot_mapping,
+        # and token_indices_to_sample
+        max_ctx_per_req = cad.max_query_len
+        max_tokens_per_req = max_ctx_per_req + num_query_per_req
+        BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))
+        num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE
+        grid = (batch_size, num_blocks)
+
+        has_num_rejected = num_rejected_tokens_gpu is not None
+        copy_and_expand_dflash_inputs_kernel[grid](
+            # Inputs
+            next_token_ids_ptr=next_token_ids,
+            target_positions_ptr=target_positions,
+            # Outputs
+            out_input_ids_ptr=self.input_ids,
+            out_context_positions_ptr=self._context_positions_buffer,
+            out_query_positions_ptr=self.positions,
+            out_context_slot_mapping_ptr=self._context_slot_mapping_buffer,
+            out_query_slot_mapping_ptr=self._slot_mapping_buffer,
+            out_token_indices_ptr=token_indices_to_sample,
+            # Block table
+            block_table_ptr=cad.block_table_tensor,
+            block_table_stride=cad.block_table_tensor.stride(0),
+            # Metadata
+            query_start_loc_ptr=cad.query_start_loc,
+            num_rejected_tokens_ptr=(
+                num_rejected_tokens_gpu if has_num_rejected else 0
+            ),
+            # Scalars
+            parallel_drafting_token_id=self.parallel_drafting_token_id,
+            block_size=self.block_size,
+            num_query_per_req=num_query_per_req,
+            num_speculative_tokens=self.num_speculative_tokens,
+            total_input_tokens=num_context,
+            BLOCK_SIZE=BLOCK_SIZE,
+            HAS_NUM_REJECTED=has_num_rejected,
+        )
+
+        query_slot_mapping = self._slot_mapping_buffer[:num_query_total]
+        new_query_start_loc = self.arange[: batch_size + 1] * num_query_per_req
+
+        # In padded mode, cad.seq_lens includes rejected tokens. Subtract
+        # them so attention only sees the valid prefix of context states.
+        effective_seq_lens = cad.seq_lens
+        if has_num_rejected:
+            effective_seq_lens = effective_seq_lens - num_rejected_tokens_gpu
+
+        # Skip num_rejected_tokens (GPU-only); overestimating is fine here.
+        new_seq_lens_cpu_upper_bound = (
+            cad.seq_lens_cpu_upper_bound + num_query_per_req
+            if cad.seq_lens_cpu_upper_bound is not None
+            else None
+        )
+        new_cad = CommonAttentionMetadata(
+            query_start_loc=new_query_start_loc,
+            seq_lens=effective_seq_lens + num_query_per_req,
+            query_start_loc_cpu=(
+                torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone()
+                * num_query_per_req
+            ),
+            _seq_lens_cpu=None,
+            _num_computed_tokens_cpu=None,
+            seq_lens_cpu_upper_bound=new_seq_lens_cpu_upper_bound,
+            num_reqs=cad.num_reqs,
+            num_actual_tokens=num_query_total,
+            max_query_len=num_query_per_req,
+            max_seq_len=cad.max_seq_len + num_query_per_req,
+            block_table_tensor=cad.block_table_tensor,
+            slot_mapping=query_slot_mapping,
+            causal=self.dflash_causal,
+        )
+
+        return num_query_total, token_indices_to_sample, new_cad
+
+    @override
+    @torch.inference_mode()
+    def dummy_run(
+        self,
+        num_tokens: int,
+        use_cudagraphs: bool = True,
+        is_graph_capturing: bool = False,
+        slot_mappings: dict[str, torch.Tensor] | None = None,
+    ) -> None:
+        """
+        Key differences to default dummy_run:
+        - Only one forward pass due to parallel drafting
+        - DFlash uses context states as unpadded metadata, so hidden_states will
+        use the unpadded num_tokens instead of num_input_tokens
+        - max_query_tokens is quite small, DFlash only sees spec tokens as queries
+        - Multimodal inputs are not currently supported
+        """
+        num_query_tokens = min(num_tokens, self.max_query_tokens)
+        cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
+            self._determine_batch_execution_and_padding(
+                num_query_tokens, use_cudagraphs=use_cudagraphs
+            )
+        )
+
+        # Slot mapping sized to num_input_tokens (query only), matching
+        # the K/V tensor size from the model forward.  Context KVs are
+        # pre-inserted separately and don't flow through the model.
+        if (
+            self._draft_attn_layer_names
+            and slot_mappings is not None
+            and next(iter(self._draft_attn_layer_names)) in slot_mappings
+        ):
+            slot_mapping_dict = self._get_slot_mapping(num_input_tokens)
+        else:
+            slot_mapping_dict = slot_mappings or {}
+
+        # Context and query positions use separate buffers; no copy needed.
+        context_positions = self._context_positions_buffer[:num_tokens]
+        # Context states will be passed directly to the precomputation without
+        # going through the buffer, since no CUDA graph is used for the precomputation.
+        # For the dummy run, we use the dummy buffer.
+        context_states = self.hidden_states[:num_tokens]
+
+        # Run the KV projection (GEMM + norms + RoPE) for memory profiling,
+        self.model.precompute_and_store_context_kv(context_states, context_positions)
+        with set_forward_context(
+            None,
+            self.vllm_config,
+            num_tokens=num_input_tokens,
+            num_tokens_across_dp=num_tokens_across_dp,
+            cudagraph_runtime_mode=cudagraph_runtime_mode,
+            slot_mapping=slot_mapping_dict,
+        ):
+            self.model(
+                input_ids=self.input_ids[:num_input_tokens],
+                positions=self._get_positions(num_input_tokens),
+                inputs_embeds=None,
+            )
+
+    @override
+    def build_model_inputs_first_pass(
+        self,
+        num_tokens: int,
+        num_input_tokens: int,
+        mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,
+    ) -> tuple[dict[str, Any], int]:
+        # Context and query positions/slots were written to separate
+        # buffers by the kernel — no copy needed.
+        num_context = self._dflash_num_context
+
+        # Pre-insert context KVs directly into cache
+        self.model.precompute_and_store_context_kv(
+            self._dflash_hidden_states,  # Shape is already [num_context, hidden_size]
+            self._context_positions_buffer[:num_context],
+            self._context_slot_mapping_buffer[:num_context],
+        )
+        return (
+            dict(
+                input_ids=self.input_ids[:num_input_tokens],
+                positions=self._get_positions(num_input_tokens),
+                inputs_embeds=None,
+            ),
+            num_input_tokens,
+        )
+
+    @override
+    def build_per_group_and_layer_attn_metadata(
+        self, cad: CommonAttentionMetadata, draft_index: int = 0
+    ) -> tuple[list[object], dict[str, object]]:
+        per_group, per_layer = super().build_per_group_and_layer_attn_metadata(
+            cad, draft_index
+        )
+        if not self.dflash_causal:
+            # Require all layers to support non-causal attention when required by DFlash
+            for layer_name, attn_metadata in per_layer.items():
+                assert getattr(attn_metadata, "causal", None) is False, (
+                    f"Attention metadata for layer {layer_name} does not have"
+                    " non-causal support, which is required for DFlash."
+                    " Consider using a different attention backend, e.g FlashAttention."
+                )
+        return per_group, per_layer
+
+    @override
+    def _get_eagle3_use_aux_hidden_state_from_config(self):
+        return self.dflash_config.get("use_aux_hidden_state", True)
+
+    @property
+    def dflash_config(self):
+        return getattr(self.draft_model_config.hf_config, "dflash_config", None) or {}

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")

docker-images/glm53-flash-dflash2/qwen3_dflash.py

@@ -0,0 +1,974 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import io
+from collections.abc import Iterable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import Qwen3Config
+
+from vllm import _custom_ops as ops
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.multimodal.inputs import NestedTensors
+from vllm.transformers_utils.config import set_default_rope_theta
+from vllm.transformers_utils.repo_utils import get_hf_file_bytes
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
+    get_eagle3_aux_layers_from_config,
+)
+
+from .qwen2 import Qwen2MLP as Qwen3MLP
+from .qwen3 import Qwen3ForCausalLM
+from .utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    get_draft_quant_config,
+    maybe_prefix,
+    process_eagle_weight,
+)
+
+logger = init_logger(__name__)
+
+
+_SLIDING_ATTENTION = "sliding_attention"
+
+
+def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
+    """Resolve explicit causality before falling back to legacy layer defaults."""
+    is_causal = getattr(config, "is_causal", None)
+    if is_causal is not None:
+        return bool(is_causal)
+    override = (getattr(config, "dflash_config", None) or {}).get("causal")
+    if override is not None:
+        return bool(override)
+    layer_types = getattr(config, "layer_types", None)
+    return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION
+
+
+def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
+    """Whether the draft needs a non-causal-capable backend, resolved from config
+    (config mirror of the model's ``get_draft_attn_causal``, usable pre-build)."""
+    return not all(
+        _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers)
+    )
+
+
+def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None:
+    """The target's RoPE layout, from its first attention layer.
+
+    A DFlash head must rotate Q/K the way the target it was distilled against
+    does, and a mismatch is silent — acceptance collapses but nothing errors and
+    the output stays correct. Draft checkpoints do not carry this, so take it
+    from the target. None if the target uses no RoPE.
+    """
+    language_model = (
+        target_model.get_language_model()
+        if hasattr(target_model, "get_language_model")
+        else target_model
+    )
+    for module in language_model.modules():
+        style = getattr(module, "is_neox_style", None)
+        if isinstance(style, bool):
+            return style
+    return None
+
+
+def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
+    spec_config = vllm_config.speculative_config
+    config = spec_config.draft_model_config.hf_config
+    aux_layers = get_eagle3_aux_layers_from_config(spec_config)
+    num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers
+    target_hidden_size = (
+        getattr(config, "target_hidden_size", None) or config.hidden_size
+    )
+    return target_hidden_size * num_features_to_use
+
+
+def _resolve_layer_attention(
+    config: Qwen3Config, layer_idx: int
+) -> tuple[int | None, bool]:
+    """Resolve ``(sliding_window, causal)`` for one DFlash draft layer.
+
+    +----------------------+-------------------------+--------------------------------+
+    | Config               | ``layer_type``          | *``causal``                    |
+    +======================+=========================+================================+
+    | ``layer_types``      | SWA if ``use_swa``      | True if ``layer_types[i]=SWA`` |
+    |                      | else ``layer_types[i]`` | else False                     |
+    +----------------------+-------------------------+--------------------------------+
+    | ``layer_types=None`` | SWA                     | False                          |
+    | + ``use_swa=True``   |                         |                                |
+    +----------------------+-------------------------+--------------------------------+
+    | ``layer_types=None`` | Full                    | False                          |
+    | + ``use_swa=False``  |                         |                                |
+    +----------------------+-------------------------+--------------------------------+
+    * If ``dflash_config.causal`` is set, its value overrides ``causal`` for all layers.
+
+    This is to support a varied ecosystem of checkpoints, including:
+    - XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash (sets "use_swa", assumes non-causal)
+    - z-lab/gemma-4-31B-it-DFlash (has mixed layer types, assumes causal only for SWA)
+    - z-lab/Qwen3.5-9B-DFlash ("standard" DFlash, all full attn, assumes non-causal)
+    """
+    dflash_config = getattr(config, "dflash_config", None) or {}
+    layer_types = getattr(config, "layer_types", None)
+    use_swa = dflash_config.get("use_swa", False)
+
+    any_sliding = False
+    if layer_types is not None:
+        num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types)
+        any_sliding = num_sliding > 0
+        # Mixed sliding/full attention needs multiple KV groups (V2 runner only).
+        if (
+            0 < num_sliding < len(layer_types)
+            and not get_current_vllm_config().use_v2_model_runner
+        ):
+            raise NotImplementedError(
+                "DFlash drafters with mixed sliding/full attention require "
+                "the V2 model runner; relaunch with "
+                "VLLM_USE_V2_MODEL_RUNNER=1."
+            )
+
+    # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``.
+    if layer_types is None or (use_swa and not any_sliding):
+        is_sliding = use_swa
+    else:
+        is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION
+
+    sliding_window = None
+    if is_sliding:
+        sliding_window = dflash_config.get(
+            "swa_window_size", getattr(config, "sliding_window", None)
+        )
+        if sliding_window is None:
+            raise ValueError(
+                "DFlash sliding attention requires a window size configured in "
+                "dflash_config.swa_window_size or the top-level sliding_window."
+            )
+
+    return sliding_window, _dflash_layer_causal(config, layer_idx)
+
+
+class DFlashQwen3Attention(nn.Module):
+    """Attention for DFlash speculative decoding.
+
+    Context KVs are pre-inserted into the KV cache before the forward pass.
+    This layer handles only query tokens via standard attention.
+    Adapted from Qwen3Attention."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_parameters: dict,
+        max_position: int = 4096 * 32,
+        head_dim: int | None = None,
+        rms_norm_eps: float = 1e-06,
+        attention_bias: bool = False,
+        add_swa_attention_sink_bias: bool = False,
+        sliding_window: int | None = None,
+        causal: bool = False,
+        is_neox_style: bool = True,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
+        super().__init__()
+        self.layer_name = prefix
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=attention_bias,  # DFlash has o_proj bias when using attention bias
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position,
+            is_neox_style=is_neox_style,
+            rope_parameters=rope_parameters,
+        )
+
+        self.attention_sink_bias = (
+            torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
+            if add_swa_attention_sink_bias
+            else None
+        )
+
+        self.sliding_window = sliding_window
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            per_layer_sliding_window=sliding_window,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+            sinks=self.attention_sink_bias,
+        )
+        self.causal = causal
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        """DFlash attention assumes that the KV cache is already populated
+        with the context K/V from the target model's hidden states. This forward op
+        computes attention for the query tokens only.
+        See also: precompute_and_store_context_kv"""
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        # Per-head RMSNorm
+        q_shape, k_shape = q.shape, k.shape
+        q = self.q_norm(
+            q.view(*q_shape[:-1], q_shape[-1] // self.head_dim, self.head_dim)
+        ).view(q_shape)
+        k = self.k_norm(
+            k.view(*k_shape[:-1], k_shape[-1] // self.head_dim, self.head_dim)
+        ).view(k_shape)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class DFlashQwen3DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        config: Qwen3Config,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        set_default_rope_theta(config, default_theta=1000000)
+        attn_type = AttentionType.DECODER
+
+        # DFlash drafts store the sink-bias flag inside dflash_config; fall back
+        # to the top-level attribute used by other (e.g. MiMo) configs.
+        dflash_config = getattr(config, "dflash_config", None) or {}
+        add_swa_attention_sink_bias = dflash_config.get(
+            "attention_sink_bias",
+            getattr(config, "add_swa_attention_sink_bias", False),
+        )
+
+        # Resolve this layer's attention mode (full vs sliding window, causal vs
+        # non-causal) from the draft config.
+        sliding_window, causal = _resolve_layer_attention(config, layer_idx)
+
+        # RoPE layout, copied off the target at load time by the draft loader
+        # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it:
+        # a head distilled from an interleaved-RoPE target must rotate the way
+        # that target does, or every drafted Q/K is wrong and acceptance
+        # collapses with no error raised.
+        is_neox_style = getattr(config, "is_neox_style", True)
+
+        self.self_attn = DFlashQwen3Attention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            max_position=config.max_position_embeddings,
+            num_kv_heads=config.num_key_value_heads,
+            rms_norm_eps=config.rms_norm_eps,
+            attention_bias=getattr(config, "attention_bias", False),
+            add_swa_attention_sink_bias=add_swa_attention_sink_bias,
+            sliding_window=sliding_window,
+            causal=causal,
+            is_neox_style=is_neox_style,
+            head_dim=getattr(config, "head_dim", None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            rope_parameters=config.rope_parameters,
+            prefix=f"{prefix}.self_attn",
+            attn_type=attn_type,
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        else:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+@support_torch_compile
+class DFlashQwen3Model(nn.Module):
+    decoder_layer_cls = DFlashQwen3DecoderLayer
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "midlayer.": "layers.0.",
+            # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder
+            # `encoder.fc` / `encoder.output_norm_enc`; this head calls them
+            # `fc` / `hidden_norm`. Same tensors and shapes, different names.
+            "encoder.output_norm_enc.": "hidden_norm.",
+            "encoder.fc.": "fc.",
+        },
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        },
+    )
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.vocab_size = self.config.vocab_size
+        self.quant_config = get_draft_quant_config(vllm_config)
+
+        drafter_config = getattr(self.config, "eagle_config", {})
+        drafter_config.update(getattr(self.config, "dflash_config", {}))
+
+        self.use_aux_hidden_state = drafter_config.get(
+            "use_aux_hidden_state",
+            getattr(self.config, "use_aux_hidden_state", True),
+        )
+
+        current_vllm_config = get_current_vllm_config()
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash
+        # checkpoints will have the mask embedding in the vocabulary embedding table
+        # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship
+        # with a separate mask embedding tensor to use instead. When present, we load it
+        # and substitute it for embed_tokens[mask_token_id] when computing embeddings.
+        self.mask_token_id = drafter_config.get(
+            "mask_token_id", getattr(self.config, "mask_token_id", None)
+        )
+        self.mask_embedding = nn.Parameter(
+            torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype),
+            requires_grad=False,
+        )
+        self.has_separate_mask_embedding = False
+
+        self.layers = nn.ModuleList(
+            [
+                self.decoder_layer_cls(
+                    current_vllm_config,
+                    config=self.config,
+                    layer_idx=layer_idx,
+                    cache_config=current_vllm_config.cache_config,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
+                )
+                for layer_idx in range(self.config.num_hidden_layers)
+            ]
+        )
+        if self.use_aux_hidden_state:
+            self.fc = ReplicatedLinear(
+                input_size=_get_dflash_fc_input_size(
+                    vllm_config,
+                ),
+                output_size=self.config.hidden_size,
+                bias=False,
+                params_dtype=vllm_config.model_config.dtype,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "fc"),
+                return_bias=False,
+            )
+        self.hidden_norm = RMSNorm(
+            self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+        self.norm = RMSNorm(
+            self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+        # The context projection concatenates K/V weights from every draft
+        # layer. It is not a LinearBase module because its output layout is a
+        # DFlash-specific fusion. Serialized MXFP8 checkpoints retain an
+        # independently packed copy in this parameter container.
+        self._fused_kv_linear = nn.Module()
+        self._fused_kv_quant_method = None
+        self._fused_kv_weight: torch.Tensor | None = None
+        self._fused_kv_weight_scale: torch.Tensor | None = None
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        embeds = self.embed_tokens(input_ids)
+        if self.has_separate_mask_embedding and self.mask_token_id is not None:
+            # Replace masked slots with the dedicated mask embedding.
+            is_mask = (input_ids == self.mask_token_id).unsqueeze(-1)
+            embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds)
+        return embeds
+
+    def _build_context_kv_buffers(
+        self,
+        layers_attn: list[nn.Module],
+        has_bias: bool,
+    ) -> None:
+        from vllm.model_executor.layers.quantization.modelopt import (
+            ModelOptMxFp8LinearMethod,
+        )
+
+        quant_methods = [a.qkv_proj.quant_method for a in layers_attn]
+        uses_mxfp8 = [
+            isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods
+        ]
+        if any(uses_mxfp8) and not all(uses_mxfp8):
+            raise ValueError(
+                "Every DFlash attention layer must use the same MXFP8 "
+                "linear format for the fused context K/V projection."
+            )
+
+        self._hidden_norm_weight = self.hidden_norm.weight.data
+
+        # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
+        kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
+        self._fused_kv_weight = torch.cat(kv_weights, dim=0)
+
+        if all(uses_mxfp8):
+            kv_scales = [a.qkv_proj.weight_scale[a.q_size :] for a in layers_attn]
+            self._fused_kv_weight_scale = torch.cat(kv_scales, dim=0)
+        else:
+            self._fused_kv_weight_scale = None
+        if has_bias:
+            kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]
+            self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)
+        else:
+            self._fused_kv_bias = None
+
+        # K-norm weights stacked into one contiguous [num_layers, head_dim]
+        # tensor so the per-layer K-norm runs as a single grouped kernel.
+        self._k_norm_weights = torch.stack(
+            [a.k_norm.weight.data for a in layers_attn], dim=0
+        ).contiguous()
+
+    def _build_fused_kv_buffers(self) -> None:
+        """Build fused weight buffers for precompute_and_store_context_kv.
+
+        Must be called after weights are loaded. Stacks the KV-projection
+        weights, K-norm weights, and RoPE parameters from every attention
+        layer so that precompute_and_store_context_kv can run one fused
+        GEMM for all layers at once. Also aliases the weight of the hidden_norm.
+        """
+        layers_attn = [layer.self_attn for layer in self.layers]
+        attn0 = layers_attn[0]
+        has_bias = attn0.qkv_proj.bias is not None
+
+        self._build_context_kv_buffers(layers_attn, has_bias)
+
+        # RoPE parameters
+        self._rope_head_size = attn0.rotary_emb.head_size
+        self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache
+        self._rope_is_neox = attn0.rotary_emb.is_neox_style
+        # Validation that RoPE params are the same across all layers
+        for attn in layers_attn[1:]:
+            assert (
+                attn.rotary_emb.head_size == self._rope_head_size
+                and attn.rotary_emb.is_neox_style == self._rope_is_neox
+            ), "All layers must have the same RoPE parameters for DFlash precomputation"
+
+        # Layer metadata
+        self._num_attn_layers = len(layers_attn)
+        self._kv_size = attn0.kv_size
+        self._head_dim = attn0.head_dim
+        self._num_kv_heads = attn0.num_kv_heads
+        self._rms_norm_eps = attn0.q_norm.variance_epsilon
+        # Validation that all layers have the same attention config
+        for attn in layers_attn[1:]:
+            assert (
+                attn.kv_size == self._kv_size
+                and attn.head_dim == self._head_dim
+                and attn.num_kv_heads == self._num_kv_heads
+                and attn.q_norm.variance_epsilon == self._rms_norm_eps
+            ), "All layers must have the same attn config for DFlash precomputation"
+
+        # References to inner Attention layers for direct cache writes
+        self._attn_layers = [layer.self_attn.attn for layer in self.layers]
+
+    def process_weights_after_loading(self) -> None:
+        """Pack the serialized MXFP8 context projection for its GEMM backend."""
+        from vllm.model_executor.layers.quantization.modelopt import (
+            ModelOptMxFp8LinearMethod,
+        )
+
+        quant_method = self.layers[0].self_attn.qkv_proj.quant_method
+        if not isinstance(quant_method, ModelOptMxFp8LinearMethod):
+            return
+        if self._fused_kv_weight is None or self._fused_kv_weight_scale is None:
+            raise RuntimeError(
+                "The DFlash MXFP8 context projection requires serialized K/V "
+                "weights and block scales to be fused during checkpoint loading."
+            )
+
+        output_size, input_size = self._fused_kv_weight.shape
+        self._fused_kv_linear.input_size_per_partition = input_size
+        self._fused_kv_linear.output_size_per_partition = output_size
+        self._fused_kv_linear.logical_widths = [output_size]
+        self._fused_kv_linear.register_parameter(
+            "weight", nn.Parameter(self._fused_kv_weight, requires_grad=False)
+        )
+        self._fused_kv_linear.register_parameter(
+            "weight_scale",
+            nn.Parameter(self._fused_kv_weight_scale, requires_grad=False),
+        )
+        quant_method.process_weights_after_loading(self._fused_kv_linear)
+        self._fused_kv_quant_method = quant_method
+        self._fused_kv_weight = None
+        self._fused_kv_weight_scale = None
+        logger.info_once(
+            "Using %s for the fused DFlash context K/V projection.",
+            type(quant_method.kernel).__name__,
+        )
+
+    def _project_context_kv(
+        self,
+        context_states: torch.Tensor,
+        num_ctx: int,
+        num_layers: int,
+        num_kv_heads: int,
+        head_dim: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # --- Fused KV projection (one GEMM for all layers) ---
+        normed_context_states = torch.empty_like(context_states)
+        ops.rms_norm(
+            normed_context_states,
+            context_states,
+            self._hidden_norm_weight,
+            self._rms_norm_eps,
+        )
+        if self._fused_kv_quant_method is None:
+            if self._fused_kv_weight is None:
+                raise RuntimeError("DFlash context K/V projection is not initialized.")
+            all_kv_flat = F.linear(
+                normed_context_states, self._fused_kv_weight, self._fused_kv_bias
+            )
+        else:
+            all_kv_flat = self._fused_kv_quant_method.apply(
+                self._fused_kv_linear,
+                normed_context_states,
+                self._fused_kv_bias,
+            )
+        # Single contiguous copy that separates K/V and transposes to
+        # layer-major layout.  Result: [2, L, num_ctx, nkv, hd] contiguous.
+        # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
+        all_kv = (
+            all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)
+            .permute(2, 1, 0, 3, 4)
+            .contiguous()
+        )
+        all_k = all_kv[0]  # [L, num_ctx, nkv, hd], contiguous
+        all_v = all_kv[1]  # [L, num_ctx, nkv, hd], contiguous
+        return all_k, all_v
+
+    def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
+        # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
+        # The weight is selected per layer by the outermost (layer) index.
+        all_k_normed = torch.empty_like(all_k)
+        ops.rms_norm(
+            all_k_normed,
+            all_k,
+            self._k_norm_weights,
+            self._rms_norm_eps,
+        )
+        return all_k_normed
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """Precompute K/V for context states write them into each layer's KV cache.
+
+        Input context states are projected to K/V, normed, and have RoPE applied.
+        Since the context shape is different than the query shape, we can't rely on the
+        regular forward pass to apply torch.compile and CUDA graphs to this section.
+        As such, this function is optimized to minimize the number of torch ops present:
+        we use fused vLLM kernels for RMSNorm and RoPE, fuse the GEMM into one
+        large projection, and avoid cloning buffers (with .contiguous()) where possible.
+
+        When context_slot_mapping is None (e.g. during dummy_run) only
+        the computation runs, and no K/V is written to cache.
+        """
+        if not hasattr(self, "_num_attn_layers"):
+            logger.warning_once(
+                "DFlash buffer initialization was skipped. If dummy weights are not "
+                "in use, this may indicate an error in weight loading."
+            )
+            self._build_fused_kv_buffers()
+
+        num_ctx = context_states.shape[0]
+        L = self._num_attn_layers
+        kv = self._kv_size
+        hd = self._head_dim
+        nkv = self._num_kv_heads
+
+        all_k, all_v = self._project_context_kv(context_states, num_ctx, L, nkv, hd)
+        all_k_normed = self._normalize_context_k(all_k)
+
+        # --- Fused RoPE across all layers ---
+        # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).
+        # In-place RoPE: pass K as the "query" arg with key=None.
+        all_k_flat = all_k_normed.view(L * num_ctx, kv)
+        positions_repeated = context_positions.repeat(L)
+        cos_sin_cache = self._rope_cos_sin_cache
+        if cos_sin_cache.dtype != all_k_flat.dtype:
+            cos_sin_cache = cos_sin_cache.to(dtype=all_k_flat.dtype)
+        ops.rotary_embedding(
+            positions_repeated,
+            all_k_flat,
+            None,
+            self._rope_head_size,
+            cos_sin_cache,
+            self._rope_is_neox,
+        )
+
+        if context_slot_mapping is None:
+            return
+
+        # --- Per-layer cache insert ---
+        all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
+        per_layer = isinstance(context_slot_mapping, (list, tuple))
+        for i in range(L):
+            slot_mapping = (
+                context_slot_mapping[i] if per_layer else context_slot_mapping
+            )
+            if slot_mapping is None:
+                continue  # dummy run: skip cache ops
+            attn = self._attn_layers[i]
+            kv_cache = attn.kv_cache
+            attn.impl.do_kv_cache_update(
+                attn,
+                all_k_final[i],
+                all_v[i],
+                kv_cache,
+                slot_mapping,
+            )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        input_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if input_embeds is None:
+            input_embeds = self.embed_input_ids(input_ids)
+
+        hidden_states = input_embeds
+
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                residual=residual,
+            )
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def _preprocess(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        for name, loaded_weight in weights:
+            if "attention_sink_bias" in name:
+                # Sink bias is per-head; shard it across TP ranks like the
+                # attention heads themselves.
+                heads_per_rank = loaded_weight.shape[0] // tp_size
+                loaded_weight = loaded_weight.narrow(
+                    0, tp_rank * heads_per_rank, heads_per_rank
+                )
+            yield name, loaded_weight
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            self._preprocess(weights), mapper=self.hf_to_vllm_mapper
+        )
+
+
+class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
+    model_cls = DFlashQwen3Model
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.draft_model_config = vllm_config.speculative_config.draft_model_config
+        self.config = self.draft_model_config.hf_config
+        if getattr(self.config, "draft_vocab_size", None) is None:
+            self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+        self.model = self.model_cls(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            start_layer_id=target_layer_num,
+        )
+
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.lm_head = ParallelLMHead(
+            self.config.draft_vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(
+            self.config.draft_vocab_size, scale=logit_scale
+        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: NestedTensors | None = None,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions, inputs_embeds)
+
+    def get_draft_kv_cache_layer_names(self) -> list[str]:
+        return [layer.self_attn.attn.layer_name for layer in self.model.layers]
+
+    def get_draft_attn_causal(self) -> list[bool]:
+        """Per-layer attention causality, aligned with
+        get_draft_kv_cache_layer_names."""
+        return [layer.self_attn.causal for layer in self.model.layers]
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if self.draft_id_to_target_id is None:
+            return logits
+
+        base = torch.arange(self.config.draft_vocab_size, device=logits.device)
+        targets = base + self.draft_id_to_target_id
+        logits_new = logits.new_full(
+            (logits.shape[0], self.config.vocab_size),
+            float("-inf"),
+        )
+        logits_new[:, targets] = logits
+        return logits_new
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """Precompute projected + RoPE'd K/V and write to cache."""
+        self.model.precompute_and_store_context_kv(
+            context_states, context_positions, context_slot_mapping
+        )
+
+    def process_weights_after_loading(self) -> None:
+        """Finalize the DFlash fused context projection after linear packing."""
+        self.model.process_weights_after_loading()
+
+    def combine_hidden_states(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.model.use_aux_hidden_state:
+            return hidden_states
+        needs_squeeze = hidden_states.dim() == 1
+        if needs_squeeze:
+            hidden_states = hidden_states.unsqueeze(0)
+        expected = self.model.fc.input_size
+        if hidden_states.shape[-1] != expected:
+            raise ValueError(
+                f"DFlash drafter expects {expected} concatenated aux hidden "
+                f"features but received {hidden_states.shape[-1]}. This usually "
+                "means the draft model's target_layer_ids reference layers that "
+                "do not exist in the target model (incompatible draft/target pair)."
+            )
+        result = self.model.fc(hidden_states)
+        if needs_squeeze:
+            result = result.squeeze(0)
+        return result
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        model_weights = {}
+        includes_draft_id_mapping = False
+        includes_embed_tokens = False
+        for name, loaded_weight in weights:
+            assert "mask_hidden" not in name, (
+                "DFlash embeds masked slots via mask_token_id (optionally "
+                "overridden by a mask_embedding.pt file); it should not ship a "
+                "mask_hidden weight."
+            )
+            if "t2d" in name:
+                continue
+            if "d2t" in name:
+                name = name.replace("d2t", "draft_id_to_target_id")
+                includes_draft_id_mapping = True
+            elif "lm_head" not in name:
+                name = "model." + name
+            if "embed_tokens" in name:
+                includes_embed_tokens = True
+            model_weights[name] = loaded_weight
+            process_eagle_weight(self, name)
+
+        # Route the separately-trained mask embedding (if shipped) through the
+        # standard weight loader alongside the rest of the draft weights.
+        mask_embedding = self._read_mask_embedding()
+        if mask_embedding is not None:
+            model_weights["model.mask_embedding"] = mask_embedding
+            self.model.has_separate_mask_embedding = True
+
+        orig_to_new_substr = {}
+        if not includes_draft_id_mapping:
+            orig_to_new_substr["draft_id_to_target_id"] = None
+        if not includes_embed_tokens:
+            orig_to_new_substr["embed_tokens"] = None
+        if not self.model.use_aux_hidden_state:
+            orig_to_new_substr["fc."] = None
+        if not self.model.has_separate_mask_embedding:
+            orig_to_new_substr["mask_embedding"] = None
+        mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
+        loader = AutoWeightsLoader(self)
+        loader.load_weights(model_weights.items(), mapper=mapper)
+        self.model._build_fused_kv_buffers()
+
+    def _read_mask_embedding(self) -> torch.Tensor | None:
+        """Checks for an override mask embedding in `mask_embedding.pt` and returns it.
+
+        Some checkpoints ship a separately-trained mask embedding for the mask token,
+        which we use to overwrite the embedding for `mask_token_id`. This helper
+        checks for the file, loads the pytorch tensor, and returns the embedding to use.
+
+        Returns None if the override file is not present.
+        """
+        mask_token_id = self.model.mask_token_id
+        if mask_token_id is None:
+            return None
+
+        MASK_EMBEDDING_FILENAME = "mask_embedding.pt"
+        data = get_hf_file_bytes(
+            MASK_EMBEDDING_FILENAME,
+            self.draft_model_config.model,
+            self.draft_model_config.revision,
+        )
+        if data is None:
+            return None
+
+        state = torch.load(io.BytesIO(data), weights_only=True)
+        if isinstance(state, dict):
+            if state.get("mask_token_id", mask_token_id) != mask_token_id:
+                raise ValueError(
+                    f"{MASK_EMBEDDING_FILENAME} mask_token_id does not match "
+                    f"dflash_config.mask_token_id ({mask_token_id}). "
+                    f"Got {state.get('mask_token_id')}."
+                )
+            state = state["embedding"]
+
+        logger.info(
+            "Loaded DFlash mask embedding for mask_token_id %s from %s",
+            mask_token_id,
+            MASK_EMBEDDING_FILENAME,
+        )
+        return state.reshape(-1)

docker-images/glm53-flash-dflash2/qwen3_dflash2.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.compilation.backends import set_model_tag
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+from .qwen3_dflash import (
+    DFlashQwen3DecoderLayer,
+    DFlashQwen3ForCausalLM,
+    DFlashQwen3Model,
+)
+from .utils import maybe_prefix
+
+
+def _grouped_conv(
+    hidden_states: torch.Tensor,
+    delta: torch.Tensor,
+    base: torch.Tensor,
+    block_size: int,
+    num_groups: int,
+    group_size: int,
+    taps: int,
+) -> torch.Tensor:
+    blocks = hidden_states.unflatten(-1, (num_groups, group_size))
+    coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1)
+    output = coefficients[:, 0] * blocks
+    position = torch.arange(hidden_states.shape[0], device=hidden_states.device)
+    if block_size & (block_size - 1) == 0:
+        position = position & (block_size - 1)
+    else:
+        position = position % block_size
+    for tap in range(1, taps):
+        shifted = F.pad(blocks[:-tap], (0, 0, 0, 0, tap, 0))
+        output += coefficients[:, tap] * shifted * (position >= tap).view(-1, 1, 1)
+    return output.flatten(-2)
+
+
+class DFlashGroupedConv(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        taps: int,
+        group_size: int,
+        block_size: int,
+        params_dtype: torch.dtype,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        if hidden_size % group_size:
+            raise ValueError(
+                f"conv_group_size={group_size} must divide hidden_size={hidden_size}."
+            )
+        self.block_size = block_size
+        self.taps = taps
+        self.group_size = group_size
+        self.num_groups = hidden_size // group_size
+        self.base_kernel = nn.Parameter(
+            torch.empty(2, taps, hidden_size, dtype=params_dtype),
+            requires_grad=False,
+        )
+        self.kernel_projection = ReplicatedLinear(
+            hidden_size,
+            2 * taps * self.num_groups,
+            bias=False,
+            params_dtype=params_dtype,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "kernel_projection"),
+            return_bias=False,
+        )
+
+    def _convolve(
+        self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int
+    ) -> torch.Tensor:
+        return _grouped_conv(
+            hidden_states,
+            delta,
+            self.base_kernel[side],
+            self.block_size,
+            self.num_groups,
+            self.group_size,
+            self.taps,
+        )
+
+    def prepare(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        coefficients = self.kernel_projection(hidden_states).reshape(
+            hidden_states.shape[0], 2, self.taps, self.num_groups
+        )
+        return self._convolve(hidden_states, coefficients[:, 0], 0), coefficients[:, 1]
+
+    def finish(
+        self, hidden_states: torch.Tensor, coefficients: torch.Tensor
+    ) -> torch.Tensor:
+        return self._convolve(hidden_states, coefficients, 1)
+
+
+class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        config,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            vllm_config,
+            config=config,
+            layer_idx=layer_idx,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+        draft_config = config.dflash_config
+        speculative_config = vllm_config.speculative_config
+        assert speculative_config is not None
+        conv_args = dict(
+            hidden_size=config.hidden_size,
+            taps=int(draft_config["conv_kernel_size"]),
+            group_size=int(draft_config["conv_group_size"]),
+            # Query tokens per request: the bonus token plus the mask tokens.
+            block_size=1 + speculative_config.num_speculative_tokens,
+            params_dtype=vllm_config.model_config.dtype,
+            quant_config=quant_config,
+        )
+        self.attention_conv = DFlashGroupedConv(
+            **conv_args, prefix=maybe_prefix(prefix, "attention_conv")
+        )
+        self.mlp_conv = DFlashGroupedConv(
+            **conv_args, prefix=maybe_prefix(prefix, "mlp_conv")
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states, coefficients = self.attention_conv.prepare(hidden_states)
+        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+        hidden_states = self.attention_conv.finish(hidden_states, coefficients)
+
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states, coefficients = self.mlp_conv.prepare(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.mlp_conv.finish(hidden_states, coefficients)
+        return hidden_states, residual
+
+
+def _score_edges(
+    predecessor_table: torch.Tensor,
+    successor_table: torch.Tensor,
+    candidate_ids: torch.Tensor,
+    unary_logits: torch.Tensor,
+    hidden: torch.Tensor,
+    anchor_token_ids: torch.Tensor,
+    top_k: int,
+) -> torch.Tensor:
+    successors = successor_table[candidate_ids]
+    predecessor_ids = torch.cat(
+        (
+            anchor_token_ids[:, None, None].expand(-1, 1, top_k),
+            candidate_ids[:, :-1],
+        ),
+        dim=1,
+    )
+    predecessors = predecessor_table[predecessor_ids]
+    return unary_logits[:, :, None] + torch.einsum(
+        "blpr,blcr->blpc", predecessors * hidden[:, :, None], successors
+    )
+
+
+@support_torch_compile
+class CandidateSelector(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        vocab_size: int,
+        rank: int,
+        top_k: int,
+        params_dtype: torch.dtype,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> None:
+        super().__init__()
+        self.top_k = top_k
+        self.predecessor_codebook = nn.Parameter(
+            torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False
+        )
+        self.successor_codebook = nn.Parameter(
+            torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False
+        )
+        self.hidden_projection = ReplicatedLinear(
+            hidden_size,
+            rank,
+            bias=False,
+            params_dtype=params_dtype,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "hidden_projection"),
+            return_bias=False,
+        )
+
+    def forward(
+        self,
+        candidate_ids: torch.Tensor,
+        unary_logits: torch.Tensor,
+        hidden_states: torch.Tensor,
+        anchor_token_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden = self.hidden_projection(hidden_states)
+        return _score_edges(
+            self.predecessor_codebook,
+            self.successor_codebook,
+            candidate_ids,
+            unary_logits,
+            hidden,
+            anchor_token_ids,
+            self.top_k,
+        )
+
+
+class DFlash2Qwen3Model(DFlashQwen3Model):
+    decoder_layer_cls = DFlash2Qwen3DecoderLayer
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            vllm_config=vllm_config,
+            start_layer_id=start_layer_id,
+            prefix=prefix,
+        )
+        draft_config = self.config.dflash_config
+        self.input_embedding_scale = float(
+            draft_config.get("input_embedding_scale", 1.0)
+        )
+        # Without its own tag the selector shares the draft head's compile cache.
+        with set_model_tag("dflash2_candidate_selector"):
+            self.candidate_selector = CandidateSelector(
+                hidden_size=self.config.hidden_size,
+                vocab_size=self.config.vocab_size,
+                rank=int(draft_config["selector_rank"]),
+                top_k=int(draft_config["selector_top_k"]),
+                params_dtype=vllm_config.model_config.dtype,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "candidate_selector"),
+            )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return super().embed_input_ids(input_ids) * self.input_embedding_scale
+
+
+class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM):
+    model_cls = DFlash2Qwen3Model
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        draft_config = self.config.dflash_config
+        softcap = float(draft_config.get("final_logit_softcapping") or 0.0)
+        self.candidate_logits_processor = LogitsProcessor(
+            vllm_config.model_config.get_vocab_size(),
+            scale=float(draft_config.get("output_multiplier", 1.0)),
+            soft_cap=softcap if softcap > 0 else None,
+        )
+
+    def compute_candidates(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.candidate_logits_processor.get_top_k_tokens(
+            self.lm_head, hidden_states, self.model.candidate_selector.top_k
+        )
+
+
+EntryClass = DFlash2Qwen3ForCausalLM

docker-images/glm53-flash-dflash2/versions.env

@@ -0,0 +1,6 @@
+# Authoritative build inputs for the GLM-5.3-Flash DFlash2 image overlay.
+# BASE_IMAGE is itself digest-pinned at its root (see
+# ../glm53-flash-nvfp4-lab/versions.env); the chain starts there.
+# Vendored fork sources: local-inference-lab/vllm dev/jovian-judgement
+# @ da4d7be6c97434f6942292ed8abbf4b32dc44355 (UPSTREAM.md).
+BASE_IMAGE="glm53-flash-nvfp4-lab:local"