Instructions to use HYHPING2023/checkpoint-draft-dflash2 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use HYHPING2023/checkpoint-draft-dflash2 with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="HYHPING2023/checkpoint-draft-dflash2", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("HYHPING2023/checkpoint-draft-dflash2", trust_remote_code=True)
model = AutoModel.from_pretrained("HYHPING2023/checkpoint-draft-dflash2", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use HYHPING2023/checkpoint-draft-dflash2 with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "HYHPING2023/checkpoint-draft-dflash2"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "HYHPING2023/checkpoint-draft-dflash2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/HYHPING2023/checkpoint-draft-dflash2

SGLang

How to use HYHPING2023/checkpoint-draft-dflash2 with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "HYHPING2023/checkpoint-draft-dflash2" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "HYHPING2023/checkpoint-draft-dflash2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "HYHPING2023/checkpoint-draft-dflash2" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "HYHPING2023/checkpoint-draft-dflash2",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use HYHPING2023/checkpoint-draft-dflash2 with Docker Model Runner:
```
docker model run hf.co/HYHPING2023/checkpoint-draft-dflash2
```

HYHPING2023 commited on 3 days ago

Commit

edfa72c

verified ·

1 Parent(s): 37403bc

DFlash2 draft (epoch_3_step_35300) for Qwen3.5-35B-A3B VCLR3 target: 5-layer draft + two-tap dynamic convs + top-16 candidate selector, block 8

Browse files

Files changed (10) hide show

.gitattributes +1 -0
README.md +60 -0
chat_template.jinja +154 -0
config.json +74 -0
dflash.py +908 -0
model.safetensors +3 -0
preprocessor_config.json +21 -0
processor_config.json +60 -0
tokenizer.json +3 -0
tokenizer_config.json +33 -0

.gitattributes CHANGED Viewed

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md CHANGED Viewed

@@ -1,3 +1,63 @@
 ---
 license: apache-2.0
 ---

 ---
 license: apache-2.0
+pipeline_tag: text-generation
+library_name: transformers
+inference: false
+tags:
+- dflash2
+- speculative-decoding
+- draft-model
+- sglang
 ---
+# checkpoint-draft-dflash2 (Qwen3.5-35B-A3B VCLR3 drafter)
+DFlash2 speculative-decoding **draft model** for the Qwen3.5-35B-A3B (VCLR3 game-video
+SFT) target. Not a standalone language model: it runs inside a speculative decoding
+server and drafts tokens for the target to verify (target weights NOT included;
+point `--model-path` at your own merged Qwen3.5-35B-A3B VCLR3 checkpoint).
+## Architecture
+DFlash2 (Inco AI, [blog](https://inco.ai/blog/dflash2/); z-lab-compatible weight
+layout) with block_size 8:
+- 5-layer Qwen3-style dual-stream draft over the target's captured layer
+  `{1, 10, 19, 28, 37}` hidden states, full 248k vocab (uses the target's
+  embed_tokens / lm_head)
+- Two-tap grouped dynamic causal convolutions (kernel 2, group 16) around every
+  attention/MLP sublayer — fixes block-end (suffix) decay
+- Top-16 candidate path selector (rank 256): `S_t(a,b) = U_t(b) + <A(a) ⊙ H(h_t), B(b)>`
+## Serving (SGLang main, DFlash2-aware DFLASH worker)
+```bash
+python -m sglang.launch_server \
+  --model-path /path/to/Qwen3.5-35B-A3B-vclr3-target \
+  --speculative-algorithm DFLASH \
+  --speculative-draft-model-path HYHPING2023/checkpoint-draft-dflash2 \
+  --speculative-num-draft-tokens 8 \
+  --trust-remote-code --tp 1 --enable-metrics
+```
+Notes: request `chat_template_kwargs: {"enable_thinking": false}` (the draft was
+trained on non-thinking answers); SGLang builds before DFlash2 support silently
+ignore the selector.
+## Evaluation (temperature 0, greedy; block 8)
+| Metric | DFlash1 baseline | **DFlash2 (this)** |
+|---|---|---|
+| Offline accept length / cycle (12.1k blocks, mixed corpus) | 2.72 | **2.97** (+9%) |
+| — video blocks | — | **3.54** |
+| — text blocks | — | 2.89 |
+| SGLang serve accept length (single stream, text) | 1.80 | **2.15** (+19%) |
+Per-position conditional acceptance rises 0.63 → 0.72 toward the block end
+(DFlash1 stays flat ~0.60) — the convolution's suffix-decay fix.
+## Training
+4 epochs on the mixed video+zh+en corpus (70.6k samples), warm-started from the
+DFlash1 b8 mixed2 checkpoint, FSDP + frozen online 35B target. Teacher-forced
+selector CE. Checkpoint: `epoch_3_step_35300`.

chat_template.jinja ADDED Viewed

	@@ -0,0 +1,154 @@

+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

config.json ADDED Viewed

	@@ -0,0 +1,74 @@

+{
+  "transformers_version": "5.8.1",
+  "architectures": [
+    "DFlash2DraftModel"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": "bfloat16",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "vocab_size": 248320,
+  "hidden_size": 2048,
+  "intermediate_size": 8192,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 2,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "max_position_embeddings": 262144,
+  "initializer_range": 0.02,
+  "rms_norm_eps": 1e-06,
+  "use_cache": true,
+  "tie_word_embeddings": false,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "attention_bias": false,
+  "use_sliding_window": false,
+  "sliding_window": null,
+  "max_window_layers": 5,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "attention_dropout": 0.0,
+  "pad_token_id": null,
+  "bos_token_id": null,
+  "eos_token_id": 248044,
+  "_name_or_path": "output/dflash2/draft_config.json",
+  "model_type": "qwen3",
+  "auto_map": {
+    "AutoModel": "dflash.DFlash2DraftModel"
+  },
+  "block_size": 8,
+  "dflash_config": {
+    "mask_token_id": 248043,
+    "target_layer_ids": [
+      1,
+      10,
+      19,
+      28,
+      37
+    ],
+    "conv_kernel_size": 2,
+    "conv_group_size": 16,
+    "selector_rank": 256,
+    "selector_top_k": 16
+  },
+  "num_target_layers": 40,
+  "output_attentions": false
+}

dflash.py ADDED Viewed

	@@ -0,0 +1,908 @@

+# coding=utf-8
+# DFlash2 draft model — self-contained (single file) on purpose.
+#
+# Sources, cross-verified against each other:
+#   1. drafttrain/dflash/_vendored/dflash_model.py — the SpecForge DFlash1 backbone this
+#      repo already trains and serves through SGLang's DFLASH algorithm. The backbone
+#      code below is copied VERBATIM from it (attention / decoder layer / model /
+#      spec_generate) so train<->serve behavior matches the proven DFlash1 path.
+#   2. z-lab/dflash dflash/model.py (Apache-2.0) — the official DFlash2DraftModel
+#      reference: GroupedDynamicCausalConv (conv_kernel_size=2, conv_group_size=16) and
+#      CandidateSelector (selector_rank=256, selector_top_k=16). Parameter names are
+#      kept identical (layers.{i}.attention_conv/mlp_conv.{base_kernel,kernel_projection},
+#      candidate_selector.{predecessor_codebook,successor_codebook,hidden_projection})
+#      so exported weights load under SGLang's DFlash2 support and under the z-lab
+#      reference implementation unchanged.
+#
+# Two deliberate deviations from z-lab's reference:
+#   - The conv is BLOCK-LOCAL during training: when the noise stream is a whole number
+#     of blocks (the training layout: N blocks of [anchor, mask*bs-1] concatenated),
+#     the predecessor tap zero-pads at each block start instead of reading across
+#     blocks — exactly matching inference, where the draft sees one block at a time.
+#     With a single block (<= block_size, the inference shape) the behavior is
+#     bit-identical to z-lab's F.pad reference.
+#   - conv_identity_init (default True): base kernel starts as [1, 0] (identity) and
+#     kernel_projection at zero, so a warm-started DFlash1 checkpoint behaves
+#     identically at step 0 and the convs grow in smoothly during training.
+#
+# This file must stay importable stand-alone: export_draft.py copies it into the served
+# draft directory as ``dflash.py`` (auto_map -> "dflash.DFlash2DraftModel"), where no
+# ``drafttrain`` package exists.
+from typing import Callable, ClassVar, Optional
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import DynamicCache
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen3.modeling_qwen3 import (
+    ALL_ATTENTION_FUNCTIONS,
+    FlashAttentionKwargs,
+    GradientCheckpointingLayer,
+    Qwen3Config,
+    Qwen3MLP,
+    Qwen3PreTrainedModel,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    eager_attention_forward,
+    rotate_half,
+)
+from typing_extensions import Tuple, Unpack
+def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    if temperature < 1e-5:
+        return torch.argmax(logits, dim=-1)
+    bsz, seq_len, vocab_size = logits.shape
+    logits = logits.view(-1, vocab_size)
+    logits = logits / temperature
+    probs = torch.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1).view(bsz, seq_len)
+def _sampling_probs(
+    logits: torch.Tensor,
+    temperature: float,
+    top_p: float = 1.0,
+    top_k: int = 0,
+) -> torch.Tensor:
+    """Softmax over (optionally top-k/top-p filtered) logits, scattered back to full vocab."""
+    scores = logits.float() / temperature
+    vocab_size = scores.shape[-1]
+    if 0 < top_k < vocab_size:
+        scores, indices = torch.topk(scores, top_k, dim=-1)
+    else:
+        indices = None
+    probs = torch.softmax(scores, dim=-1)
+    if top_p < 1.0:
+        sorted_probs, order = probs.sort(dim=-1, descending=True)
+        keep = sorted_probs.cumsum(dim=-1) - sorted_probs < top_p
+        sorted_probs = sorted_probs * keep
+        probs = torch.zeros_like(probs).scatter(-1, order, sorted_probs)
+        probs = probs / probs.sum(dim=-1, keepdim=True)
+    if indices is not None:
+        probs = torch.zeros_like(logits, dtype=probs.dtype).scatter(-1, indices, probs)
+    return probs
+def _sample_probs(probs: torch.Tensor) -> torch.Tensor:
+    shape = probs.shape[:-1]
+    return torch.multinomial(probs.view(-1, probs.shape[-1]), 1).view(shape)
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_len = q.size(-2)
+    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+def _dflash_config(config) -> dict:
+    return getattr(config, "dflash_config", {}) or {}
+def _draft_value(config, name, default=None):
+    return _dflash_config(config).get(name, getattr(config, name, default))
+class Qwen3DFlashAttention(nn.Module):
+    """Dual-stream attention (copied from the vendored DFlash1 model, unchanged).
+    K/V are computed from BOTH the captured target hidden states (context stream) and
+    the draft's own noise-stream hidden states, concatenated.
+    """
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = (
+            config.sliding_window
+            if config.layer_types[layer_idx] == "sliding_attention"
+            else None
+        )
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len = hidden_states.shape[:-1]
+        ctx_len = target_hidden.shape[1]
+        q = self.q_proj(hidden_states)
+        q = q.view(bsz, q_len, -1, self.head_dim)
+        q = self.q_norm(q).transpose(1, 2)
+        k_ctx = self.k_proj(target_hidden)
+        k_noise = self.k_proj(hidden_states)
+        v_ctx = self.v_proj(target_hidden)
+        v_noise = self.v_proj(hidden_states)
+        k = torch.cat([k_ctx, k_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        v = torch.cat([v_ctx, v_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
+        attn_fn: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_output, attn_weights = attn_fn(
+            self,
+            q,
+            k,
+            v,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+# --------------------------------------------------------------------------- #
+# DFlash2 addition 1: two-tap grouped dynamic causal convolution
+# --------------------------------------------------------------------------- #
+def _grouped_dynamic_convolve(hidden, dynamic, base, group_size):
+    """z-lab reference convolve: out[t] = sum_offset (base[offset] + dynamic[t, offset]) * x[t-offset].
+    ``hidden`` (B', L', H), ``dynamic`` (B', L', kernel_size, groups), ``base``
+    (kernel_size, H). Zero left-pad: x[t-offset] = 0 for t < offset.
+    """
+    batch, length, hidden_size = hidden.shape
+    groups = hidden_size // group_size
+    blocks = hidden.view(batch, length, groups, group_size)
+    dynamic = dynamic.view(batch, length, base.shape[0], groups, 1)
+    output = torch.zeros_like(blocks)
+    for offset in range(base.shape[0]):
+        values = blocks if offset == 0 else F.pad(blocks[:, :-offset], (0, 0, 0, 0, offset, 0))
+        kernel = base[offset].view(1, 1, groups, group_size).to(hidden.dtype)
+        output = output + kernel * values
+        output = torch.addcmul(output, dynamic[:, :, offset], values)
+    return output.view_as(hidden)
+class GroupedDynamicCausalConv(nn.Module):
+    """Two-tap grouped dynamic causal convolution (DFlash2).
+    Parameter layout is identical to z-lab/dflash's GroupedDynamicCausalConv:
+      - ``base_kernel``: (2, kernel_size, hidden). Index 0 holds the static taps for
+        the PRE-sublayer conv (``prepare``), index 1 for the POST-sublayer conv
+        (``finish``). The two dynamic kernels are computed once from the pre-sublayer
+        (normed) hidden state and reused by both.
+      - ``kernel_projection``: Linear(hidden, 2 * kernel_size * groups, bias=False) —
+        per-position dynamic tap corrections; every ``group_size`` channels share one.
+    Block-local extension (training): when the input length is a whole number of
+    blocks AND longer than one block, the sequence is processed as concatenated
+    independent blocks — the predecessor tap zero-pads at each block start instead of
+    reading the previous block's tail. This exactly matches inference, where the draft
+    forward sees a single block ([anchor, mask, ...]) at a time. With a single block
+    (length <= block_size) the computation is bit-identical to the z-lab reference.
+    """
+    def __init__(self, hidden_size: int, kernel_size: int, group_size: int, block_size: int):
+        super().__init__()
+        if hidden_size % group_size != 0:
+            raise ValueError(
+                f"GroupedDynamicCausalConv requires group_size to divide hidden_size; "
+                f"got hidden_size={hidden_size}, group_size={group_size}"
+            )
+        if kernel_size < 1:
+            raise ValueError(f"kernel_size must be >= 1, got {kernel_size}")
+        if block_size < 1:
+            raise ValueError(f"block_size must be >= 1, got {block_size}")
+        self.kernel_size = kernel_size
+        self.group_size = group_size
+        self.block_size = block_size
+        self.num_groups = hidden_size // group_size
+        self.base_kernel = nn.Parameter(torch.empty(2, kernel_size, hidden_size))
+        self.kernel_projection = nn.Linear(
+            hidden_size, 2 * kernel_size * self.num_groups, bias=False
+        )
+    def _split_blocks(self, hidden: torch.Tensor):
+        """(B, L, ...) -> ((B*N, bs, ...), bsz, n). Whole-block sequences (L = N*bs > bs,
+        the training noise layout) become N independent blocks; anything else (a single
+        full or partial block, the inference layout) is returned as one block."""
+        bsz, seq_len = hidden.shape[0], hidden.shape[1]
+        if seq_len > self.block_size and seq_len % self.block_size == 0:
+            n = seq_len // self.block_size
+            return hidden.reshape(bsz * n, self.block_size, *hidden.shape[2:]), bsz, n
+        return hidden, bsz, 1
+    def prepare(self, hidden: torch.Tensor):
+        """Pre-sublayer: convolve the (normed) input; stash the finish-step taps.
+        Returns (conv_out (B, L, H), dynamic (B, L, kernel_size, groups)).
+        """
+        bsz, seq_len = hidden.shape[0], hidden.shape[1]
+        blocked, _, _ = self._split_blocks(hidden)
+        dynamic = self.kernel_projection(blocked).view(
+            *blocked.shape[:-1], 2, self.kernel_size, self.num_groups
+        )
+        out = _grouped_dynamic_convolve(
+            blocked, dynamic[..., 0, :, :], self.base_kernel[0], self.group_size
+        )
+        return (
+            out.reshape(bsz, seq_len, hidden.shape[-1]),
+            dynamic[..., 1, :, :].reshape(
+                bsz, seq_len, self.kernel_size, self.num_groups
+            ),
+        )
+    def finish(self, hidden: torch.Tensor, dynamic: torch.Tensor) -> torch.Tensor:
+        """Post-sublayer: convolve the sublayer output with the stashed taps."""
+        bsz, seq_len = hidden.shape[0], hidden.shape[1]
+        blocked, _, _ = self._split_blocks(hidden)
+        dyn_blocked = dynamic.reshape(
+            blocked.shape[0], blocked.shape[1], self.kernel_size, self.num_groups
+        )
+        out = _grouped_dynamic_convolve(
+            blocked, dyn_blocked, self.base_kernel[1], self.group_size
+        )
+        return out.reshape(bsz, seq_len, hidden.shape[-1])
+# --------------------------------------------------------------------------- #
+# Decoder layer (DFlash1 layer + optional DFlash2 conv hooks, z-lab pattern)
+# --------------------------------------------------------------------------- #
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    """DFlash decoder layer with optional DFlash2 conv hooks.
+    With ``attention_conv``/``mlp_conv`` None (DFlash1) the forward is identical to
+    the vendored SpecForge layer. When set (DFlash2), each sublayer is wrapped as:
+    prepare(normed input) -> sublayer -> finish(sublayer output), exactly z-lab's
+    placement. The convs act ONLY on the noise stream (``hidden_states``); the context
+    stream (``target_hidden``) feeds k_ctx/v_ctx directly, unconverted.
+    """
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.attention_conv: Optional[GroupedDynamicCausalConv] = None
+        self.mlp_conv: Optional[GroupedDynamicCausalConv] = None
+    def forward(
+        self,
+        target_hidden: Optional[torch.Tensor] = None,
+        hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attention_kernel = None
+        if self.attention_conv is not None:
+            hidden_states, attention_kernel = self.attention_conv.prepare(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        if attention_kernel is not None:
+            hidden_states = self.attention_conv.finish(hidden_states, attention_kernel)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        mlp_kernel = None
+        if self.mlp_conv is not None:
+            hidden_states, mlp_kernel = self.mlp_conv.prepare(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        if mlp_kernel is not None:
+            hidden_states = self.mlp_conv.finish(hidden_states, mlp_kernel)
+        hidden_states = residual + hidden_states
+        return hidden_states
+def build_target_layer_ids(num_target_layers: int, num_draft_layers: int):
+    if num_draft_layers == 1:
+        return [(num_target_layers // 2)]
+    start = 1
+    end = num_target_layers - 3
+    span = end - start
+    target_layer_ids = [
+        int(round(start + (i * span) / (num_draft_layers - 1)))
+        for i in range(num_draft_layers)
+    ]
+    return target_layer_ids
+def extract_context_feature(
+    hidden_states: list[torch.Tensor],
+    layer_ids: Optional[list[int]],
+) -> torch.Tensor:
+    offset = 1
+    selected_states = []
+    for layer_id in layer_ids:
+        selected_states.append(hidden_states[layer_id + offset])
+    target_hidden = torch.cat(selected_states, dim=-1)
+    return target_hidden
+# --------------------------------------------------------------------------- #
+# DFlash2 addition 2: candidate path selector
+# --------------------------------------------------------------------------- #
+class CandidateSelector(nn.Module):
+    """Top-k candidate path selector (z-lab-compatible parameter layout).
+    Scores adjacent candidate pairs with a gated low-rank bilinear form::
+        S_t(a, b) = U_t(b) + < A(a) ⊙ H(h_t), B(b) >
+    where ``U_t`` is the draft's own logit for candidate ``b`` (how much the drafter
+    liked it on its own), ``A``/``B`` are compact per-token codebooks and ``H(h_t)`` is
+    a context gate projected from the draft hidden state deciding which parts of the
+    predecessor/successor match count.
+    """
+    def __init__(self, config):
+        super().__init__()
+        rank = int(_draft_value(config, "selector_rank", 256))
+        top_k = int(_draft_value(config, "selector_top_k", 16))
+        if rank <= 0:
+            raise ValueError(f"selector_rank must be > 0, got {rank}")
+        if top_k <= 0:
+            raise ValueError(f"selector_top_k must be > 0, got {top_k}")
+        self.rank = rank
+        self.top_k = top_k
+        self.predecessor_codebook = nn.Embedding(config.vocab_size, rank)
+        self.successor_codebook = nn.Embedding(config.vocab_size, rank)
+        self.hidden_projection = nn.Linear(config.hidden_size, rank, bias=False)
+    def pairwise_scores(
+        self,
+        hidden: torch.Tensor,
+        logits: torch.Tensor,
+        prev_ids: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Teacher-forced pairwise scores for TRAINING.
+        hidden (B, L, H), logits (B, L, V), prev_ids (B, L) true predecessor token ids.
+        Returns (candidates (B, L, k), scores (B, L, k)) with
+        ``scores[b, t, j] = logits[b, t, cand_j] + <A(prev) ⊙ H(h_t), B(cand_j)>``.
+        Fully parallel — no sequential walk (teacher forcing supplies the predecessor).
+        """
+        unary, candidates = torch.topk(logits, self.top_k, dim=-1, sorted=False)
+        gate = self.hidden_projection(hidden)                     # (B, L, R)
+        pred = self.predecessor_codebook(prev_ids.long())         # (B, L, R)
+        succ = self.successor_codebook(candidates)                # (B, L, k, R)
+        scores = unary + torch.einsum("blr,blr,blkr->blk", pred, gate, succ)
+        return candidates, scores
+    def select(
+        self,
+        hidden: torch.Tensor,
+        logits: torch.Tensor,
+        anchor_ids: torch.Tensor,
+        temperature: float,
+    ):
+        """Inference-time path walk (z-lak reference): greedy at T=0, else sampling
+        from softmax over the k candidate scores (also returned for lossless
+        rejection sampling). ``anchor_ids`` is the last verified token."""
+        unary, candidates = torch.topk(logits, self.top_k, dim=-1, sorted=False)
+        hidden = self.hidden_projection(hidden)
+        # Accept (B,) or (B, 1) anchor ids (the walk keeps a flat (B,) predecessor).
+        predecessor = anchor_ids.reshape(anchor_ids.shape[0], -1).squeeze(-1)
+        path, q_rows = [], []
+        for position in range(hidden.shape[1]):
+            scores = unary[:, position] + torch.einsum(
+                "br,bkr->bk",
+                self.predecessor_codebook(predecessor) * hidden[:, position],
+                self.successor_codebook(candidates[:, position]),
+            )
+            if temperature > 0:
+                q = _sampling_probs(scores[:, None], temperature)[:, 0]
+                index = _sample_probs(q)
+                q_rows.append(q)
+            else:
+                index = torch.argmax(scores, dim=-1)
+            predecessor = candidates[:, position].gather(-1, index[:, None])[:, 0]
+            path.append(predecessor)
+        return (
+            torch.stack(path, dim=1),
+            candidates,
+            torch.stack(q_rows, dim=1) if q_rows else None,
+        )
+# --------------------------------------------------------------------------- #
+# DFlash (1) backbone — copied verbatim from drafttrain/dflash/_vendored/dflash_model.py
+# --------------------------------------------------------------------------- #
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules: ClassVar[list[str]] = ["Qwen3DFlashDecoderLayer"]
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DFlashDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        dflash_config = getattr(config, "dflash_config", {}) or {}
+        self.target_layer_ids = dflash_config.get(
+            "target_layer_ids",
+            build_target_layer_ids(config.num_target_layers, config.num_hidden_layers),
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config)
+        self.fc = nn.Linear(
+            len(self.target_layer_ids) * config.hidden_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.block_size = config.block_size
+        self.mask_token_id = dflash_config.get("mask_token_id", None)
+        self.projector_type = dflash_config.get("projector_type", None)
+        self.pure_draft_prefix_len = dflash_config.get("pure_draft_prefix_len", 0)
+        self.shift_label = dflash_config.get("shift_label", False)
+        if self.projector_type == "domino":
+            self.emb_dim = dflash_config["emb_dim"]
+            self.gru_hidden_dim = dflash_config["gru_hidden_dim"]
+            self.prefix_gru = nn.GRU(
+                input_size=config.hidden_size,
+                hidden_size=self.gru_hidden_dim,
+                num_layers=1,
+                batch_first=True,
+                bias=False,
+            )
+            in_dim = config.hidden_size + self.gru_hidden_dim
+            self.embed_proj = nn.Sequential(
+                nn.Linear(in_dim, self.emb_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(self.emb_dim, config.vocab_size, bias=False),
+            )
+        elif self.projector_type is not None:
+            raise ValueError(f"Unknown draft projector_type: {self.projector_type}")
+        self.post_init()
+    def forward(
+        self,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        noise_embedding: Optional[torch.Tensor] = None,
+        target_hidden: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = noise_embedding
+        target_hidden = self.hidden_norm(self.fc(target_hidden))
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                target_hidden=target_hidden,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        return self.norm(hidden_states)
+    @torch.inference_mode()
+    def spec_generate(
+        self,
+        target: nn.Module,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        self.eval()
+        num_input_tokens = input_ids.shape[1]
+        max_length = num_input_tokens + max_new_tokens
+        block_size = self.block_size
+        output_ids = torch.full(
+            (1, max_length + block_size),
+            self.mask_token_id,
+            dtype=torch.long,
+            device=target.device,
+        )
+        position_ids = torch.arange(
+            output_ids.shape[1], device=target.device
+        ).unsqueeze(0)
+        past_key_values_target = DynamicCache()
+        past_key_values_draft = DynamicCache()
+        # Prefill stage
+        output = target(
+            input_ids,
+            position_ids=position_ids[:, :num_input_tokens],
+            past_key_values=past_key_values_target,
+            use_cache=True,
+            logits_to_keep=1,
+            output_hidden_states=True,
+        )
+        output_ids[:, :num_input_tokens] = input_ids
+        output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(
+            output.logits, temperature
+        )
+        target_hidden = extract_context_feature(
+            output.hidden_states, self.target_layer_ids
+        )
+        # Decode stage
+        acceptance_lengths = []
+        start = input_ids.shape[1]
+        while start < max_length:
+            block_output_ids = output_ids[:, start : start + block_size].clone()
+            block_position_ids = position_ids[:, start : start + block_size]
+            noise_embedding = target.model.embed_tokens(block_output_ids)
+            draft_logits = target.lm_head(
+                self(
+                    target_hidden=target_hidden,
+                    noise_embedding=noise_embedding,
+                    position_ids=position_ids[
+                        :, past_key_values_draft.get_seq_length() : start + block_size
+                    ],
+                    past_key_values=past_key_values_draft,
+                    use_cache=True,
+                    is_causal=False,
+                )[:, -block_size + 1 :, :]
+            )
+            past_key_values_draft.crop(start)
+            block_output_ids[:, 1:] = sample(draft_logits)
+            output = target(
+                block_output_ids,
+                position_ids=block_position_ids,
+                past_key_values=past_key_values_target,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+            posterior = sample(output.logits, temperature)
+            acceptance_length = (
+                (block_output_ids[:, 1:] == posterior[:, :-1])
+                .cumprod(dim=1)
+                .sum(dim=1)[0]
+                .item()
+            )
+            output_ids[:, start : start + acceptance_length + 1] = block_output_ids[
+                :, : acceptance_length + 1
+            ]
+            output_ids[:, start + acceptance_length + 1] = posterior[
+                :, acceptance_length
+            ]
+            start += acceptance_length + 1
+            past_key_values_target.crop(start)
+            target_hidden = extract_context_feature(
+                output.hidden_states, self.target_layer_ids
+            )[:, : acceptance_length + 1, :]
+            acceptance_lengths.append(acceptance_length + 1)
+            if stop_token_ids is not None and any(
+                stop_token_id in output_ids[:, num_input_tokens:]
+                for stop_token_id in stop_token_ids
+            ):
+                break
+        output_ids = output_ids[:, :max_length]
+        output_ids = output_ids[:, output_ids[0] != self.mask_token_id]
+        if stop_token_ids is not None:
+            stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device)
+            stop_token_indices = torch.isin(
+                output_ids[0][num_input_tokens:], stop_token_ids
+            ).nonzero(as_tuple=True)[0]
+            if stop_token_indices.numel() > 0:
+                output_ids = output_ids[
+                    :, : num_input_tokens + stop_token_indices[0] + 1
+                ]
+        return output_ids
+# --------------------------------------------------------------------------- #
+# DFlash2 draft model
+# --------------------------------------------------------------------------- #
+class DFlash2DraftModel(DFlashDraftModel):
+    """DFlash backbone with the two DFlash2 additions.
+    - Every decoder layer gets ``attention_conv`` + ``mlp_conv``
+      (GroupedDynamicCausalConv, kernel_size=2 / group_size=16 by default).
+    - A ``candidate_selector`` (CandidateSelector, rank=256 / top_k=16 by default).
+    Config keys (in ``dflash_config``): ``conv_kernel_size``, ``conv_group_size``,
+    ``conv_identity_init``, ``selector_rank``, ``selector_top_k`` — mirroring the
+    released z-lab/Qwen3.8-27B-DFlash2 config.json (which carries the first two and
+    the last two; ``conv_identity_init`` is a training-only knob).
+    Weight layout is z-lab-compatible, so exported checkpoints load under SGLang's
+    DFlash2 support and under z-lab/dflash's reference implementation.
+    """
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        # The SERVED checkpoint stores the selector codebooks under bare keys
+        # (no ".weight" — z-lab/SGLang format; see export_draft.py). Map them onto
+        # this model's nn.Embedding parameters when loading via HF transformers.
+        kwargs.setdefault(
+            "key_mapping",
+            {
+                f"candidate_selector.{name}": f"candidate_selector.{name}.weight"
+                for name in ("predecessor_codebook", "successor_codebook")
+            },
+        )
+        return super().from_pretrained(*args, **kwargs)
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        dflash_config = _dflash_config(config)
+        kernel_size = int(dflash_config.get("conv_kernel_size", 2))
+        group_size = int(dflash_config.get("conv_group_size", 16))
+        self.conv_identity_init = bool(dflash_config.get("conv_identity_init", True))
+        for layer in self.layers:
+            layer.attention_conv = GroupedDynamicCausalConv(
+                config.hidden_size, kernel_size, group_size, self.block_size
+            )
+            layer.mlp_conv = GroupedDynamicCausalConv(
+                config.hidden_size, kernel_size, group_size, self.block_size
+            )
+        self.candidate_selector = CandidateSelector(config)
+        # post_init() initializes the newly added Linear/Embedding modules (already
+        # initialized parent modules are skipped); the raw base_kernel Parameters and
+        # the identity pattern are then set explicitly so initialization does not
+        # depend on transformers' double-post_init semantics.
+        self.post_init()
+        self._init_dflash2_weights()
+    def _init_dflash2_weights(self) -> None:
+        std = float(getattr(self.config, "initializer_range", 0.02))
+        with torch.no_grad():
+            for layer in self.layers:
+                for conv in (layer.attention_conv, layer.mlp_conv):
+                    if conv is None:  # pragma: no cover (always set for DFlash2)
+                        continue
+                    # Static taps start as identity: tap-0 (self) = 1, tap-1
+                    # (predecessor) = 0. Anything else destroys the residual stream at
+                    # init (a zero tap-0 zeroes the sublayer input).
+                    conv.base_kernel.zero_()
+                    conv.base_kernel[:, 0, :].fill_(1.0)
+                    if self.conv_identity_init:
+                        conv.kernel_projection.weight.zero_()
+                    else:
+                        conv.kernel_projection.weight.normal_(mean=0.0, std=std)
+            sel = self.candidate_selector
+            sel.predecessor_codebook.weight.normal_(mean=0.0, std=std)
+            sel.successor_codebook.weight.normal_(mean=0.0, std=std)
+            sel.hidden_projection.weight.normal_(mean=0.0, std=std)
+    def propose(
+        self,
+        hidden: torch.Tensor,
+        anchor_ids: torch.Tensor,
+        output_head: nn.Module,
+        temperature: float,
+    ):
+        """Select a draft path through the top-k candidates (SGLang/spec entry point)."""
+        logits = output_head(hidden)
+        return self.candidate_selector.select(hidden, logits, anchor_ids, temperature)
+    @torch.inference_mode()
+    def spec_generate(
+        self,
+        target: nn.Module,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        """DFlash1 spec_generate with the candidate-selector walk replacing the
+        independent per-position argmax/sample (the only behavioral change)."""
+        self.eval()
+        num_input_tokens = input_ids.shape[1]
+        max_length = num_input_tokens + max_new_tokens
+        block_size = self.block_size
+        output_ids = torch.full(
+            (1, max_length + block_size),
+            self.mask_token_id,
+            dtype=torch.long,
+            device=target.device,
+        )
+        position_ids = torch.arange(
+            output_ids.shape[1], device=target.device
+        ).unsqueeze(0)
+        past_key_values_target = DynamicCache()
+        past_key_values_draft = DynamicCache()
+        # Prefill stage
+        output = target(
+            input_ids,
+            position_ids=position_ids[:, :num_input_tokens],
+            past_key_values=past_key_values_target,
+            use_cache=True,
+            logits_to_keep=1,
+            output_hidden_states=True,
+        )
+        output_ids[:, :num_input_tokens] = input_ids
+        output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(
+            output.logits, temperature
+        )
+        target_hidden = extract_context_feature(
+            output.hidden_states, self.target_layer_ids
+        )
+        # Decode stage
+        acceptance_lengths = []
+        start = input_ids.shape[1]
+        while start < max_length:
+            block_output_ids = output_ids[:, start : start + block_size].clone()
+            block_position_ids = position_ids[:, start : start + block_size]
+            noise_embedding = target.model.embed_tokens(block_output_ids)
+            draft_hidden = self(
+                target_hidden=target_hidden,
+                noise_embedding=noise_embedding,
+                position_ids=position_ids[
+                    :, past_key_values_draft.get_seq_length() : start + block_size
+                ],
+                past_key_values=past_key_values_draft,
+                use_cache=True,
+                is_causal=False,
+            )[:, -block_size + 1 :, :]
+            past_key_values_draft.crop(start)
+            draft_logits = target.lm_head(draft_hidden)
+            # DFlash2: walk one coherent path through the top-k candidates instead of
+            # taking each position's top pick independently.
+            draft_tokens, _, _ = self.candidate_selector.select(
+                draft_hidden,
+                draft_logits,
+                block_output_ids[:, 0],
+                temperature,
+            )
+            block_output_ids[:, 1:] = draft_tokens
+            output = target(
+                block_output_ids,
+                position_ids=block_position_ids,
+                past_key_values=past_key_values_target,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+            posterior = sample(output.logits, temperature)
+            acceptance_length = (
+                (block_output_ids[:, 1:] == posterior[:, :-1])
+                .cumprod(dim=1)
+                .sum(dim=1)[0]
+                .item()
+            )
+            output_ids[:, start : start + acceptance_length + 1] = block_output_ids[
+                :, : acceptance_length + 1
+            ]
+            output_ids[:, start + acceptance_length + 1] = posterior[
+                :, acceptance_length
+            ]
+            start += acceptance_length + 1
+            past_key_values_target.crop(start)
+            target_hidden = extract_context_feature(
+                output.hidden_states, self.target_layer_ids
+            )[:, : acceptance_length + 1, :]
+            acceptance_lengths.append(acceptance_length + 1)
+            if stop_token_ids is not None and any(
+                stop_token_id in output_ids[:, num_input_tokens:]
+                for stop_token_id in stop_token_ids
+            ):
+                break
+        output_ids = output_ids[:, :max_length]
+        output_ids = output_ids[:, output_ids[0] != self.mask_token_id]
+        if stop_token_ids is not None:
+            stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device)
+            stop_token_indices = torch.isin(
+                output_ids[0][num_input_tokens:], stop_token_ids
+            ).nonzero(as_tuple=True)[0]
+            if stop_token_indices.numel() > 0:
+                output_ids = output_ids[
+                    :, : num_input_tokens + stop_token_indices[0] + 1
+                ]
+        return output_ids

model.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8114ddb576d991afae44de2f9aebc56921ddada5d698528749bcaf1da81b48ea
+size 1010529880

preprocessor_config.json ADDED Viewed

	@@ -0,0 +1,21 @@

+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

processor_config.json ADDED Viewed

	@@ -0,0 +1,60 @@

+{
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_processor_type": "Qwen2VLImageProcessor",
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "merge_size": 2,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "longest_edge": 16777216,
+      "shortest_edge": 65536
+    },
+    "temporal_patch_size": 2
+  },
+  "processor_class": "Qwen3VLProcessor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 2,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "max_frames": 768,
+    "merge_size": 2,
+    "min_frames": 4,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "size": {
+      "longest_edge": 25165824,
+      "shortest_edge": 4096
+    },
+    "temporal_patch_size": 2,
+    "video_processor_type": "Qwen3VLVideoProcessor"
+  }
+}

tokenizer.json ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
+size 19989325

tokenizer_config.json ADDED Viewed

	@@ -0,0 +1,33 @@

+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "processor_class": "Qwen3VLProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}