Add files using upload-large-folder tool

init/dflash100k_epoch6/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1fc0023cdb9cd6b93336abbb5a36f8f08527368002d68ca9a1a32b72aeb93073
+size 1074860568

runs/dspark_a/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af98cf412711c7d9f7df177dd1bb20b6576bb8233b9d14da5b0b6d3aea24a717
+size 311165128

runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c29ef39cad74a7e958f4c18a31d77e97080db7c1033be6012732e84400b1fd8
+size 7382703887

runs/dspark_joint_6p/export/epoch_6/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e740457069c8cdaf7feebc6f1c9a45af543c8208fa4de0aea29af7a8cf70ddaa
+size 311165128

runs/dspark_joint_6p/export/epoch_6/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcb5d93581ed4fe4cd7405bcebdbf0fe358b8bf0ff2248d2b9e681eba602ddc9
+size 1074860568

runs/dspark_joint_6p_tv/export/epoch_1/config.json

@@ -0,0 +1,55 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [
+      1,
+      9,
+      17,
+      25,
+      33
+    ]
+  },
+  "dtype": "float32",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

runs/dspark_joint_6p_tv/export/epoch_1/dspark_config.json

@@ -0,0 +1,16 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "dflash",
+    "path": "runs/dspark_joint_6p_tv/export/epoch_1"
+  },
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "joint": true,
+  "objective": "dflash",
+  "kd": true,
+  "tv": 0.9,
+  "head_rank": 256,
+  "loss": "DSpark baseline: joint DFlash + Markov head, (1 - tv) CE + tv L1 against Qwen3-4B probabilities, tv = 0.9"
+}

runs/dspark_joint_6p_tv/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebfa3ff872cf4defd1c861d422cb4565c3a288ba4862f7a1c4044463e7282e3d
+size 311165128

runs/dspark_joint_6p_tv/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6e2aa3c88581da6c7c3cacb35f50f678b5fb164f047299deb0f2bc0dbe39589
+size 1074860568

runs/dspark_joint_6p_tv/export/epoch_1/training_recipe.json

@@ -0,0 +1,10 @@
+{
+  "loss": "DSpark baseline: joint DFlash + Markov head, (1 - tv) CE + tv L1 against Qwen3-4B probabilities, tv = 0.9",
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "objective": "dflash",
+  "kd": true,
+  "tv": 0.9,
+  "head_rank": 256
+}

runs/fmlm2_v3/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38594ad798d0869e11e0b7f1aadfeb8fb075f0739a489e129cd28e8ffac4f5f1
+size 1770876772

runs/fmlm2_v3/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:acb621a7e9913a6e9df3e151222a2b0110b2b06cb7b9f6e37c225545b401bd55
+size 311165128

runs/fmlm2_v3/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be6d410548cea0f60bdde21e66fc073ae9d2aa059160c8bfad6a64fd366950cf
+size 1074860568

runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1aaec2649c7dc202b6931d991adafb7d52a7e354bdef79920cca2a06ad8a66da
+size 12695355007

runs/fmlm2_v3_6ep/export/epoch_6/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41f71c01d20ad59568b67a80e24b338f31e3ff71f3c9f16bcc60977b823ee340
+size 1770876772

runs/fmlm2_v3_6ep/export/epoch_6/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52f3ad42a9fb35d38661d0f09a7edc8e30220c8a0435f3ad8ebe70915fa683f5
+size 311165128

runs/fmlm2_v3_6ep/export/epoch_6/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:066d537ccc3da6eb3da7b6855b6d3f226ebd89dbd0280168fab30e3f0b3c3ea7
+size 1074860568

runs/fmlm2_v3_6ep_bv/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fc7410e6b5a259a6006266d871eea6fc7fce72eb9da4f1ebcee53c62eac0aa3
+size 1770876772

runs/fmlm2_v3_6ep_bv/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:252f749694020bf01d7ed7f991776f852a25aea578ad72e431f9b8ea2e3fc329
+size 311165128

runs/fmlm2_v3_6ep_bv/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a3d2040a940e3526e7e9df6abea13efebf51d377c942655545e3e8971249d20
+size 1074860568

runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3fd239bb0f2e19eee270407c0c41fe75012f6e9f9270559d2426e019ce8cec9
+size 1770876772

runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b10c48c786c1ff5aa26624953174f48feba3ad2c10b31143691a4e91fa071806
+size 311165128

runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbc731478b05f229b4e00afed51d29e452f45ced4f588f4e8f85413a05eeeded
+size 1074860568

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [
+      1,
+      9,
+      17,
+      25,
+      33
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dflash.py

@@ -0,0 +1,277 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+from .utils import build_target_layer_ids, extract_context_feature, sample
+
+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
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    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
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    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)
+
+    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)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["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)]
+        )
+        self.target_layer_ids = self.config.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 = self.config.dflash_config.get("mask_token_id", None)
+        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

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/dspark_config.json

@@ -0,0 +1,21 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_bv_lt/export/epoch_1",
+    "mode": "onejump_zero"
+  },
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "joint": true,
+  "objective": "fmlm2",
+  "kd": true,
+  "readout_frac": 0.25,
+  "head_rank": 256,
+  "tv": 0.0,
+  "bv": 0.9,
+  "lafa_time": true,
+  "emb_mse": 0.0,
+  "fix_embed_norm": false
+}

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40b4b27ad2a1a5a4fa1fca89b55b411b39a0beaa4c34d4c25335786ab1256ef9
+size 1770876772

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow_config.json

@@ -0,0 +1,15 @@
+{
+  "learn_embed": true,
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "objective": "fmlm2",
+  "kd": true,
+  "readout_frac": 0.25,
+  "head_rank": 256,
+  "tv": 0.0,
+  "bv": 0.9,
+  "lafa_time": true,
+  "emb_mse": 0.0,
+  "fix_embed_norm": false
+}

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c15930d6cff545a25aad05afe29f81e89b750f1eb4705191473e1d12d8f85e9b
+size 311165128

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9b9a73b54c60f89efab5ce15c70454367ec1b30cc328bbc25cc5f9563d6fcb5
+size 1074860568

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/modeling_dflash.py

@@ -0,0 +1,309 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+
+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 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
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    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
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    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)
+
+    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)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        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
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["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)]
+        )
+        self.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.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,
+        mask_token_id: int,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        self.eval()
+        target.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),
+            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] != 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

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/tokenwise_time.json

@@ -0,0 +1 @@
+{"bins": 200, "beta": 0.99, "updates": 29, "t": [0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0], "tau": [[0.0, 0.0396, 0.0807, 0.122, 0.1645, 0.2052, 0.2469, 0.2888, 0.3306, 0.373, 0.4158, 0.4606, 0.5054, 0.552, 0.6004, 0.6523, 0.7076, 0.7689, 0.8392, 0.9115, 1.0], [0.0, 0.0382, 0.0774, 0.1161, 0.1548, 0.1943, 0.2331, 0.2726, 0.3134, 0.3529, 0.3941, 0.4356, 0.4774, 0.5219, 0.5675, 0.6178, 0.6713, 0.732, 0.8069, 0.8926, 1.0], [0.0, 0.0367, 0.0739, 0.1119, 0.1496, 0.1883, 0.2257, 0.2634, 0.302, 0.3401, 0.379, 0.4188, 0.4598, 0.5019, 0.5464, 0.5938, 0.6457, 0.7057, 0.7813, 0.8759, 1.0], [0.0, 0.0361, 0.0727, 0.1095, 0.1465, 0.1838, 0.2211, 0.2588, 0.2964, 0.3341, 0.372, 0.411, 0.4515, 0.4924, 0.5355, 0.5814, 0.632, 0.6912, 0.7651, 0.8664, 1.0], [0.0, 0.0357, 0.0716, 0.1078, 0.1441, 0.1808, 0.2173, 0.2541, 0.2911, 0.3288, 0.3663, 0.4048, 0.4447, 0.4849, 0.5273, 0.5722, 0.6225, 0.6805, 0.7548, 0.8587, 1.0], [0.0, 0.0355, 0.0711, 0.1068, 0.143, 0.1795, 0.2154, 0.2518, 0.2886, 0.3262, 0.3633, 0.4012, 0.4404, 0.4805, 0.5228, 0.5678, 0.6176, 0.675, 0.7473, 0.8523, 1.0], [0.0, 0.0354, 0.0709, 0.1064, 0.1421, 0.1788, 0.2151, 0.2513, 0.2878, 0.325, 0.3619, 0.3996, 0.4385, 0.4778, 0.5196, 0.5638, 0.613, 0.6692, 0.74, 0.8461, 1.0], [0.0, 0.0353, 0.0709, 0.1063, 0.1421, 0.1781, 0.2141, 0.2501, 0.2862, 0.3232, 0.3603, 0.3979, 0.4365, 0.4756, 0.5172, 0.5613, 0.6094, 0.6651, 0.7356, 0.8419, 1.0], [0.0, 0.0353, 0.0707, 0.1059, 0.1415, 0.177, 0.2128, 0.2488, 0.2851, 0.3221, 0.3593, 0.3971, 0.4354, 0.4745, 0.5158, 0.56, 0.6077, 0.6629, 0.7323, 0.8393, 1.0], [0.0, 0.0353, 0.0705, 0.106, 0.1415, 0.1769, 0.2127, 0.2483, 0.2842, 0.3209, 0.3581, 0.3958, 0.4345, 0.473, 0.5142, 0.5578, 0.6056, 0.6596, 0.7293, 0.8359, 1.0], [0.0, 0.0353, 0.0705, 0.106, 0.1412, 0.1768, 0.2125, 0.2482, 0.2843, 0.321, 0.358, 0.3954, 0.4338, 0.4725, 0.513, 0.5565, 0.6041, 0.6578, 0.7265, 0.8336, 1.0], [0.0, 0.0353, 0.0707, 0.1064, 0.142, 0.1774, 0.2131, 0.2491, 0.2851, 0.3219, 0.3586, 0.3961, 0.4343, 0.4729, 0.5134, 0.5565, 0.6038, 0.658, 0.7261, 0.8327, 1.0], [0.0, 0.0357, 0.0713, 0.1068, 0.1427, 0.1782, 0.2141, 0.2501, 0.2864, 0.3231, 0.3605, 0.3978, 0.4359, 0.4746, 0.5154, 0.5591, 0.6066, 0.6602, 0.7286, 0.8344, 1.0], [0.0, 0.0358, 0.0717, 0.1074, 0.1438, 0.18, 0.2165, 0.2529, 0.2897, 0.327, 0.3645, 0.4026, 0.4412, 0.4805, 0.5217, 0.5655, 0.6126, 0.6669, 0.7343, 0.8381, 1.0], [0.0, 0.0372, 0.0744, 0.1116, 0.1488, 0.1862, 0.2237, 0.2615, 0.2991, 0.3375, 0.3761, 0.4151, 0.4546, 0.4952, 0.5374, 0.5824, 0.6307, 0.687, 0.7579, 0.8576, 1.0]], "loss_table": [[0.56088, 0.51558, 0.47125, 0.55303, 0.50507, 0.53824, 0.55543, 0.51585, 0.5447, 0.46988, 0.48713, 0.5043, 0.54745, 0.51792, 0.50553, 0.49522, 0.46846, 0.50874, 0.51544, 0.52513, 0.58333, 0.49041, 0.48576, 0.48477, 0.45076, 0.48379, 0.55411, 0.49368, 0.49259, 0.47939, 0.49061, 0.46604, 0.49517, 0.46604, 0.56361, 0.51253, 0.49062, 0.49547, 0.55671, 0.49089, 0.52201, 0.52091, 0.47561, 0.54581, 0.47084, 0.49916, 0.51017, 0.52189, 0.45984, 0.50827, 0.5119, 0.50995, 0.48959, 0.50931, 0.4923, 0.45744, 0.50718, 0.50554, 0.46289, 0.48852, 0.49966, 0.49434, 0.50693, 0.53308, 0.50433, 0.50986, 0.46111, 0.52461, 0.47881, 0.54803, 0.48718, 0.46517, 0.48453, 0.49171, 0.49162, 0.50611, 0.49551, 0.46613, 0.51336, 0.48207, 0.48288, 0.46742, 0.48628, 0.47021, 0.47866, 0.48778, 0.44934, 0.48692, 0.46111, 0.43962, 0.43403, 0.48642, 0.46903, 0.48561, 0.45442, 0.49051, 0.45499, 0.4871, 0.48842, 0.442, 0.4295, 0.42935, 0.43266, 0.44055, 0.47079, 0.46993, 0.43197, 0.46627, 0.44402, 0.44, 0.44536, 0.44512, 0.42112, 0.40952, 0.44025, 0.4679, 0.43347, 0.41875, 0.41953, 0.42335, 0.37395, 0.44334, 0.4165, 0.38566, 0.39681, 0.37053, 0.39293, 0.41946, 0.41616, 0.41201, 0.39121, 0.39031, 0.34707, 0.3994, 0.38, 0.40136, 0.34779, 0.36477, 0.4093, 0.35792, 0.38638, 0.35186, 0.37903, 0.38355, 0.35282, 0.35586, 0.32737, 0.34689, 0.3105, 0.33526, 0.31963, 0.32623, 0.30799, 0.33363, 0.30344, 0.30949, 0.32834, 0.32899, 0.29014, 0.27143, 0.26848, 0.24515, 0.32626, 0.30014, 0.30361, 0.28792, 0.27329, 0.3174, 0.29754, 0.27771, 0.25178, 0.28983, 0.27826, 0.33397, 0.28043, 0.27025, 0.33233, 0.2596, 0.27997, 0.31038, 0.31204, 0.27314, 0.25264, 0.2712, 0.32573, 0.27814, 0.21979, 0.28643, 0.27291, 0.2474, 0.25133, 0.20983, 0.23553, 0.24367, 0.20727, 0.21999, 0.21533, 0.17414, 0.20101, 0.13969], [1.02762, 0.99538, 0.88432, 1.03223, 0.99607, 0.94777, 1.01154, 0.93535, 0.97038, 0.95691, 0.98914, 0.92718, 0.9606, 0.97058, 0.92196, 0.93416, 0.92885, 0.93895, 0.97835, 0.99102, 1.04719, 0.9829, 0.9134, 0.96962, 0.99004, 0.9783, 1.07615, 0.99307, 0.92447, 0.90467, 0.89409, 0.9108, 0.95659, 0.87792, 0.97548, 0.96517, 0.9452, 0.95927, 1.00171, 0.99361, 0.99984, 0.99025, 0.9442, 0.97116, 0.90401, 0.96519, 0.9426, 0.92727, 0.92647, 0.97093, 0.9458, 0.95741, 0.95577, 0.93808, 0.94552, 0.90474, 0.93727, 0.92243, 0.89924, 0.89257, 0.95599, 0.88696, 0.94607, 0.93743, 0.95191, 0.95026, 0.91248, 0.98515, 0.9636, 0.91119, 0.97326, 0.89274, 0.90524, 0.92688, 0.88133, 0.9491, 0.87298, 0.89009, 0.92506, 0.96889, 0.91295, 0.87118, 0.88589, 0.89584, 0.91134, 0.88502, 0.88745, 0.92258, 0.90019, 0.92157, 0.83013, 0.93339, 0.89736, 0.90471, 0.86812, 0.85899, 0.87822, 0.82619, 0.87228, 0.82634, 0.82514, 0.81643, 0.84316, 0.82564, 0.86575, 0.87875, 0.79413, 0.75266, 0.78757, 0.88704, 0.82372, 0.84408, 0.78424, 0.78401, 0.78807, 0.77592, 0.7541, 0.76547, 0.73941, 0.73318, 0.6793, 0.74895, 0.71661, 0.6989, 0.74665, 0.72632, 0.6692, 0.72292, 0.72928, 0.71417, 0.70138, 0.71887, 0.62195, 0.6543, 0.67166, 0.68749, 0.57142, 0.64435, 0.6259, 0.60847, 0.61048, 0.58654, 0.6167, 0.59466, 0.62625, 0.57759, 0.58388, 0.55195, 0.50945, 0.53128, 0.53833, 0.52966, 0.53456, 0.49407, 0.48102, 0.4915, 0.48779, 0.4682, 0.45283, 0.47576, 0.42874, 0.42646, 0.46813, 0.45699, 0.5021, 0.48529, 0.43728, 0.43921, 0.4256, 0.43359, 0.41289, 0.45334, 0.38461, 0.45049, 0.42304, 0.44604, 0.44099, 0.36123, 0.38382, 0.39788, 0.47385, 0.38239, 0.40329, 0.40414, 0.40447, 0.39779, 0.36871, 0.37758, 0.38914, 0.33856, 0.36904, 0.29247, 0.34438, 0.36225, 0.32639, 0.27743, 0.31428, 0.24372, 0.25904, 0.17693], [1.39247, 1.38406, 1.25961, 1.42022, 1.41178, 1.40185, 1.41545, 1.30526, 1.44834, 1.39428, 1.35023, 1.30685, 1.37198, 1.37245, 1.31612, 1.2943, 1.33304, 1.32829, 1.32196, 1.38355, 1.38361, 1.3113, 1.28669, 1.3425, 1.35155, 1.34296, 1.43906, 1.35421, 1.28963, 1.31364, 1.31461, 1.24478, 1.32654, 1.26231, 1.36446, 1.32051, 1.32026, 1.32412, 1.40552, 1.39609, 1.36348, 1.37519, 1.2652, 1.35077, 1.3234, 1.3659, 1.34366, 1.35858, 1.2994, 1.35514, 1.30413, 1.34826, 1.39355, 1.28966, 1.34892, 1.23785, 1.34728, 1.31551, 1.2906, 1.28194, 1.4105, 1.2745, 1.34553, 1.32836, 1.2922, 1.29844, 1.33686, 1.39398, 1.32354, 1.27727, 1.34379, 1.27164, 1.33153, 1.31322, 1.26039, 1.30327, 1.27951, 1.27229, 1.3401, 1.29432, 1.23657, 1.21544, 1.22778, 1.28945, 1.31309, 1.22509, 1.19337, 1.25486, 1.21533, 1.2803, 1.17303, 1.21369, 1.21296, 1.2955, 1.19461, 1.15021, 1.22812, 1.23281, 1.1915, 1.13653, 1.17004, 1.13911, 1.18272, 1.08145, 1.21147, 1.17842, 1.0878, 1.07052, 1.13294, 1.16885, 1.09225, 1.07897, 1.11282, 1.01316, 1.04632, 1.10286, 1.05417, 1.03913, 0.97601, 1.0391, 0.97449, 1.00173, 0.98357, 0.9542, 0.98566, 1.00895, 0.92706, 0.94812, 0.94763, 0.9082, 0.90004, 0.89454, 0.86265, 0.85017, 0.85458, 0.85139, 0.84949, 0.79366, 0.79639, 0.79003, 0.81175, 0.71356, 0.77763, 0.78755, 0.70645, 0.74831, 0.70985, 0.6955, 0.6052, 0.66888, 0.57857, 0.66656, 0.67551, 0.57231, 0.60089, 0.57981, 0.55275, 0.56903, 0.55284, 0.56615, 0.55746, 0.50724, 0.57554, 0.54949, 0.58663, 0.53995, 0.51674, 0.53292, 0.57764, 0.50524, 0.4991, 0.5182, 0.49202, 0.4968, 0.47355, 0.54508, 0.48482, 0.48534, 0.44467, 0.49241, 0.49128, 0.43389, 0.45196, 0.45235, 0.43132, 0.45529, 0.41751, 0.47448, 0.43961, 0.39298, 0.38583, 0.38189, 0.42117, 0.38821, 0.35213, 0.32267, 0.32789, 0.26984, 0.30796, 0.18058], [1.73223, 1.7755, 1.60913, 1.79926, 1.76108, 1.64749, 1.69944, 1.67286, 1.74003, 1.72781, 1.66074, 1.6537, 1.71424, 1.68032, 1.61124, 1.67927, 1.65061, 1.68088, 1.64426, 1.68928, 1.78297, 1.67711, 1.60839, 1.68218, 1.69622, 1.62467, 1.78005, 1.6707, 1.65837, 1.66272, 1.66166, 1.68785, 1.62045, 1.6244, 1.71604, 1.66162, 1.60908, 1.6545, 1.71802, 1.63318, 1.71253, 1.63038, 1.63466, 1.66389, 1.63579, 1.72739, 1.62524, 1.63875, 1.67306, 1.63177, 1.58273, 1.62124, 1.66628, 1.59659, 1.68948, 1.63038, 1.66174, 1.6504, 1.63493, 1.539, 1.68561, 1.62392, 1.6937, 1.6905, 1.60508, 1.63524, 1.63098, 1.68033, 1.65764, 1.62827, 1.66204, 1.61378, 1.62376, 1.6037, 1.53998, 1.6189, 1.52575, 1.61762, 1.6571, 1.66676, 1.60848, 1.46078, 1.54165, 1.57567, 1.61027, 1.49714, 1.44721, 1.56218, 1.49093, 1.52199, 1.49581, 1.51234, 1.52039, 1.56436, 1.46232, 1.46908, 1.52703, 1.52481, 1.50997, 1.45905, 1.46133, 1.41879, 1.41737, 1.39525, 1.45727, 1.46058, 1.38433, 1.38428, 1.43179, 1.39322, 1.35345, 1.32672, 1.32276, 1.33616, 1.29273, 1.32114, 1.29924, 1.2585, 1.244, 1.25159, 1.23503, 1.23348, 1.24236, 1.12755, 1.21586, 1.17768, 1.14899, 1.20921, 1.13216, 1.0668, 1.00905, 1.03041, 0.99304, 1.00101, 1.01254, 1.02046, 0.98293, 0.98276, 0.9988, 0.9327, 0.91558, 0.94347, 0.91974, 0.91463, 0.80197, 0.85995, 0.78762, 0.81561, 0.70711, 0.74381, 0.72987, 0.74762, 0.78596, 0.61075, 0.68011, 0.68534, 0.65152, 0.70369, 0.64562, 0.6411, 0.62759, 0.56633, 0.66411, 0.58995, 0.60712, 0.6005, 0.57258, 0.54448, 0.64266, 0.55625, 0.54688, 0.5742, 0.49621, 0.62562, 0.5693, 0.54512, 0.51883, 0.52462, 0.49653, 0.51221, 0.56307, 0.49268, 0.50355, 0.54227, 0.514, 0.51658, 0.48565, 0.53674, 0.50523, 0.42903, 0.45687, 0.41596, 0.46201, 0.44731, 0.37892, 0.38808, 0.31884, 0.27808, 0.3284, 0.20143], [1.97126, 2.08037, 1.90608, 2.0225, 2.07594, 1.92145, 1.96177, 1.98986, 2.00115, 2.03013, 2.03682, 1.95524, 1.97639, 1.91878, 1.8989, 1.99178, 1.93718, 1.97153, 1.95258, 1.98456, 2.00565, 1.91504, 1.91812, 1.96398, 1.92278, 1.94249, 2.0507, 1.98532, 1.96421, 1.95409, 1.91143, 1.91287, 1.94543, 1.93829, 1.97525, 1.94248, 1.97382, 1.92359, 1.95789, 1.91097, 2.05972, 1.92409, 1.90063, 1.93666, 1.98724, 1.97588, 1.9345, 1.87683, 1.96674, 1.90331, 1.84661, 1.89825, 1.97702, 1.93353, 1.91475, 1.87314, 1.9655, 1.9246, 1.82814, 1.83979, 1.90342, 1.84831, 1.94839, 1.90732, 1.91583, 1.91716, 1.85779, 1.96533, 1.93294, 1.91351, 1.89157, 1.84167, 1.88035, 1.83268, 1.7964, 1.8815, 1.85157, 1.84454, 1.86801, 1.88055, 1.80864, 1.70796, 1.84623, 1.84818, 1.87898, 1.76972, 1.72057, 1.81081, 1.6984, 1.7711, 1.80347, 1.76713, 1.75191, 1.83051, 1.73397, 1.76624, 1.7259, 1.7206, 1.73823, 1.62979, 1.59856, 1.66426, 1.7394, 1.61395, 1.68658, 1.72749, 1.63855, 1.61201, 1.62635, 1.64742, 1.60493, 1.54626, 1.46957, 1.48811, 1.48467, 1.48309, 1.44593, 1.41797, 1.45297, 1.38391, 1.38196, 1.42553, 1.38703, 1.35507, 1.36824, 1.32788, 1.26502, 1.33334, 1.25181, 1.2521, 1.19892, 1.22811, 1.10653, 1.16323, 1.14006, 1.15436, 1.0692, 1.04438, 1.14141, 1.06509, 1.06131, 0.99856, 0.98601, 0.97514, 0.88663, 0.93474, 0.89377, 0.9366, 0.79659, 0.80256, 0.77371, 0.86765, 0.83331, 0.76765, 0.72854, 0.76583, 0.67592, 0.71234, 0.70147, 0.72735, 0.63839, 0.62299, 0.67196, 0.64772, 0.68219, 0.64364, 0.63281, 0.59526, 0.61076, 0.60775, 0.60326, 0.62464, 0.5648, 0.68379, 0.61686, 0.61613, 0.59626, 0.55689, 0.50304, 0.54579, 0.63629, 0.5264, 0.50529, 0.58868, 0.50082, 0.55016, 0.48655, 0.57274, 0.52673, 0.49394, 0.50557, 0.4614, 0.45599, 0.48904, 0.41707, 0.42265, 0.38373, 0.34169, 0.31872, 0.19825], [2.29948, 2.32162, 2.19231, 2.26537, 2.26241, 2.13559, 2.20833, 2.19048, 2.32071, 2.27736, 2.26993, 2.18542, 2.22269, 2.20385, 2.17519, 2.25906, 2.21508, 2.21766, 2.21907, 2.22475, 2.28629, 2.17539, 2.13424, 2.21827, 2.23114, 2.16394, 2.31205, 2.1835, 2.13948, 2.16864, 2.18447, 2.20678, 2.1782, 2.17311, 2.21006, 2.15668, 2.20512, 2.27438, 2.22203, 2.22484, 2.31204, 2.11205, 2.16434, 2.21357, 2.17164, 2.27059, 2.20433, 2.10947, 2.16184, 2.18248, 2.09832, 2.22668, 2.15529, 2.12148, 2.16426, 2.20675, 2.15353, 2.14081, 2.13692, 2.04048, 2.20354, 2.16712, 2.13363, 2.04236, 2.09574, 2.16475, 2.10062, 2.15669, 2.19656, 2.15811, 2.14745, 2.08735, 2.13835, 2.09538, 2.05957, 2.17074, 2.1248, 2.10821, 2.11247, 2.01868, 2.09146, 1.94866, 2.03071, 2.08639, 2.06046, 1.95902, 1.99248, 2.05275, 1.98023, 2.05391, 1.96141, 1.9765, 1.9662, 2.05133, 1.94789, 1.95636, 1.9139, 1.94715, 1.877, 1.88535, 1.85557, 1.83223, 1.88812, 1.86214, 1.86087, 1.90436, 1.76197, 1.8038, 1.78573, 1.77042, 1.76762, 1.69047, 1.64532, 1.69977, 1.65023, 1.70298, 1.59929, 1.56564, 1.60484, 1.54357, 1.57212, 1.63075, 1.52134, 1.48676, 1.55621, 1.52791, 1.40003, 1.45327, 1.44679, 1.30968, 1.36206, 1.30888, 1.28594, 1.2497, 1.30543, 1.25327, 1.24612, 1.20588, 1.28924, 1.17645, 1.20175, 1.14743, 1.0901, 1.05679, 0.99184, 1.01213, 0.96798, 0.99818, 0.85366, 0.92132, 0.84789, 0.89244, 0.91129, 0.82771, 0.77691, 0.84395, 0.77573, 0.79755, 0.72172, 0.74945, 0.73663, 0.66551, 0.74837, 0.71232, 0.6938, 0.70028, 0.66079, 0.71968, 0.70176, 0.63633, 0.6573, 0.65824, 0.6529, 0.68353, 0.66064, 0.6619, 0.63388, 0.56119, 0.56407, 0.58616, 0.61582, 0.54014, 0.63408, 0.58378, 0.55919, 0.57138, 0.54867, 0.62126, 0.50698, 0.53344, 0.55009, 0.49392, 0.48295, 0.48856, 0.42881, 0.41303, 0.41505, 0.31643, 0.34947, 0.2366], [2.47359, 2.51735, 2.42206, 2.5459, 2.5669, 2.36877, 2.39208, 2.44683, 2.54713, 2.5176, 2.49642, 2.39543, 2.41564, 2.4165, 2.39209, 2.46109, 2.40138, 2.5069, 2.42709, 2.48118, 2.53097, 2.44242, 2.43031, 2.49249, 2.46437, 2.39914, 2.56485, 2.35905, 2.3801, 2.40193, 2.42502, 2.31485, 2.41298, 2.36776, 2.44866, 2.29361, 2.41199, 2.42504, 2.4588, 2.40675, 2.44178, 2.32894, 2.37266, 2.40121, 2.40862, 2.48079, 2.4465, 2.32679, 2.39651, 2.42232, 2.37273, 2.39648, 2.37127, 2.44724, 2.35801, 2.41787, 2.40955, 2.34589, 2.36929, 2.30641, 2.46273, 2.3193, 2.37259, 2.35799, 2.28968, 2.35178, 2.35003, 2.37437, 2.36346, 2.42206, 2.35112, 2.3048, 2.43637, 2.32979, 2.31766, 2.3895, 2.25347, 2.3401, 2.31918, 2.17544, 2.32717, 2.19994, 2.25261, 2.29007, 2.24908, 2.16508, 2.23141, 2.30002, 2.24234, 2.23647, 2.19908, 2.18255, 2.20361, 2.27208, 2.21676, 2.15747, 2.19303, 2.11893, 2.07926, 2.10254, 2.06002, 2.01838, 2.04424, 2.07531, 2.0396, 2.10542, 1.9939, 1.98383, 1.96747, 1.94138, 1.95478, 1.87451, 1.89261, 1.85529, 1.78825, 1.8502, 1.77745, 1.72942, 1.74563, 1.75454, 1.74762, 1.73999, 1.72329, 1.68798, 1.70473, 1.63008, 1.60462, 1.5829, 1.55828, 1.49039, 1.48367, 1.47681, 1.39789, 1.38646, 1.43664, 1.37043, 1.37939, 1.28682, 1.34367, 1.32691, 1.30066, 1.21403, 1.135, 1.15784, 1.07083, 1.1534, 1.05769, 1.01529, 0.9416, 1.02398, 0.96144, 0.88267, 0.91814, 0.92179, 0.88818, 0.91228, 0.8559, 0.80119, 0.75222, 0.78999, 0.7796, 0.79787, 0.82216, 0.73766, 0.73398, 0.68462, 0.70436, 0.71014, 0.70556, 0.65947, 0.70401, 0.67912, 0.65875, 0.68077, 0.69917, 0.68804, 0.69367, 0.59218, 0.61218, 0.66388, 0.66145, 0.56756, 0.64432, 0.58344, 0.63211, 0.55701, 0.57902, 0.61186, 0.51541, 0.58183, 0.5992, 0.55129, 0.53593, 0.50489, 0.44667, 0.42801, 0.37137, 0.31359, 0.36652, 0.22903], [2.62318, 2.68436, 2.63953, 2.72575, 2.72539, 2.60525, 2.54459, 2.60722, 2.69788, 2.66432, 2.64592, 2.60388, 2.63215, 2.53109, 2.62323, 2.63314, 2.65937, 2.65652, 2.61811, 2.65688, 2.68228, 2.59887, 2.59056, 2.69192, 2.60097, 2.60316, 2.67847, 2.56658, 2.5972, 2.54771, 2.69296, 2.55566, 2.67831, 2.57165, 2.58588, 2.53559, 2.57239, 2.61148, 2.66139, 2.59256, 2.63844, 2.55564, 2.57193, 2.58441, 2.62833, 2.66928, 2.56815, 2.58109, 2.60341, 2.56454, 2.63079, 2.59884, 2.66072, 2.57815, 2.57533, 2.56039, 2.52987, 2.56814, 2.57346, 2.49797, 2.61077, 2.45763, 2.52386, 2.49081, 2.53966, 2.44319, 2.51663, 2.57175, 2.56543, 2.47786, 2.52964, 2.51362, 2.59191, 2.51921, 2.45442, 2.50109, 2.49034, 2.4619, 2.47715, 2.36941, 2.48532, 2.41363, 2.43114, 2.51283, 2.42021, 2.32243, 2.40102, 2.45977, 2.38968, 2.41212, 2.46207, 2.37965, 2.37786, 2.38334, 2.33735, 2.31578, 2.26334, 2.27127, 2.27081, 2.27294, 2.1668, 2.18305, 2.2522, 2.24842, 2.22891, 2.20283, 2.12711, 2.06252, 2.10179, 2.10422, 2.03152, 2.06164, 2.05227, 2.00477, 2.01331, 2.01061, 1.94977, 1.87082, 1.95185, 1.87767, 1.87069, 1.85476, 1.78504, 1.79267, 1.78988, 1.74247, 1.73071, 1.68164, 1.6433, 1.64647, 1.53901, 1.57002, 1.539, 1.44472, 1.51117, 1.49991, 1.4721, 1.41414, 1.4426, 1.29382, 1.3529, 1.313, 1.25755, 1.18903, 1.12549, 1.18881, 1.115, 1.11362, 0.98455, 1.062, 0.97614, 1.01485, 0.94196, 0.97848, 0.9171, 0.95446, 0.9208, 0.88427, 0.8355, 0.80805, 0.83618, 0.77923, 0.74079, 0.79172, 0.79593, 0.75683, 0.69412, 0.75011, 0.79956, 0.6692, 0.68808, 0.72722, 0.68257, 0.74523, 0.70298, 0.66785, 0.68852, 0.65846, 0.64833, 0.71774, 0.65158, 0.62951, 0.66502, 0.60977, 0.62296, 0.60892, 0.58539, 0.62181, 0.61424, 0.63007, 0.61194, 0.51127, 0.50514, 0.49971, 0.45259, 0.41851, 0.3856, 0.38595, 0.37841, 0.23612], [2.82419, 2.92267, 2.82569, 2.92023, 2.86845, 2.82141, 2.80209, 2.84906, 2.92991, 2.86659, 2.83081, 2.81449, 2.84483, 2.80286, 2.83685, 2.89596, 2.88936, 2.86864, 2.78336, 2.89779, 2.87942, 2.82319, 2.74958, 2.89023, 2.84567, 2.79206, 2.93384, 2.81109, 2.83433, 2.84752, 2.89032, 2.84853, 2.78525, 2.84456, 2.84237, 2.76241, 2.84098, 2.75858, 2.88887, 2.81955, 2.87501, 2.76508, 2.83625, 2.80329, 2.83168, 2.81657, 2.75809, 2.79208, 2.8253, 2.75134, 2.82072, 2.83687, 2.80838, 2.74163, 2.76983, 2.78096, 2.70803, 2.83452, 2.73566, 2.70242, 2.82788, 2.63962, 2.77372, 2.61207, 2.63566, 2.65495, 2.71234, 2.78421, 2.77395, 2.70352, 2.69608, 2.67341, 2.69233, 2.71602, 2.55893, 2.72036, 2.66004, 2.70568, 2.63769, 2.67469, 2.61028, 2.56249, 2.63501, 2.76148, 2.63692, 2.51733, 2.63762, 2.63352, 2.50485, 2.59365, 2.64709, 2.6075, 2.58196, 2.55877, 2.54121, 2.54955, 2.4715, 2.46349, 2.38599, 2.47717, 2.34273, 2.38537, 2.42968, 2.37872, 2.46197, 2.3476, 2.34792, 2.22425, 2.22211, 2.21652, 2.19202, 2.15866, 2.17138, 2.1962, 2.20693, 2.11529, 2.14265, 2.04314, 2.07428, 2.09459, 2.03244, 2.04183, 1.99073, 1.91163, 2.00481, 1.79645, 1.80112, 1.81571, 1.7579, 1.81156, 1.71715, 1.69223, 1.60108, 1.55026, 1.66163, 1.54717, 1.6153, 1.55495, 1.55324, 1.44619, 1.41721, 1.45099, 1.29747, 1.3363, 1.27268, 1.24945, 1.18074, 1.20979, 1.14556, 1.12436, 0.99397, 1.11128, 0.98087, 1.06874, 0.98153, 1.02682, 0.95536, 0.91077, 0.88113, 0.89216, 0.82765, 0.83009, 0.77034, 0.79597, 0.81433, 0.81407, 0.72453, 0.78853, 0.7526, 0.75931, 0.79678, 0.7755, 0.73067, 0.70257, 0.67445, 0.69331, 0.7876, 0.75945, 0.6748, 0.74804, 0.75049, 0.67241, 0.6849, 0.63639, 0.67249, 0.66684, 0.65885, 0.65661, 0.63446, 0.62242, 0.55868, 0.56933, 0.47697, 0.53679, 0.4607, 0.48541, 0.42857, 0.37027, 0.39667, 0.2613], [2.96078, 3.08528, 2.98202, 3.11185, 3.07657, 2.9585, 2.93471, 3.06781, 3.08004, 2.99539, 3.03066, 2.94857, 3.09964, 2.99414, 2.93686, 3.02473, 2.98893, 3.03495, 2.96796, 3.07862, 2.98506, 2.96166, 2.9367, 3.09086, 3.05165, 2.87443, 3.05996, 3.03893, 3.07283, 3.0431, 3.06567, 2.97463, 2.98853, 3.02234, 2.97629, 2.9424, 2.99425, 2.95168, 2.95752, 3.0136, 3.04441, 2.94657, 2.92354, 3.00422, 3.00517, 3.00526, 3.02892, 2.94892, 3.03256, 2.94899, 3.03226, 2.99525, 2.99134, 2.97932, 2.85408, 2.98672, 2.90411, 2.94759, 2.84602, 2.92846, 3.05468, 2.88836, 2.83443, 2.88146, 2.84313, 2.81507, 2.86537, 2.88481, 2.91202, 2.86346, 2.87113, 2.81804, 2.90101, 2.83194, 2.76212, 2.85033, 2.81501, 2.83308, 2.84575, 2.72701, 2.73601, 2.68478, 2.80374, 2.84256, 2.78915, 2.71373, 2.74311, 2.73093, 2.71021, 2.79797, 2.7897, 2.76234, 2.75704, 2.84562, 2.73079, 2.60782, 2.64465, 2.57399, 2.6244, 2.60075, 2.52205, 2.53391, 2.49652, 2.5121, 2.51621, 2.49046, 2.53721, 2.44207, 2.44807, 2.31918, 2.37956, 2.31683, 2.36704, 2.30048, 2.22249, 2.22051, 2.22153, 2.19143, 2.21278, 2.16742, 2.13992, 2.21071, 2.0473, 2.02574, 2.11375, 1.97923, 1.91556, 1.99928, 1.84571, 1.95333, 1.76351, 1.84303, 1.68061, 1.63882, 1.72266, 1.68518, 1.62932, 1.60295, 1.54238, 1.48728, 1.57523, 1.45951, 1.37566, 1.31152, 1.36451, 1.28431, 1.21921, 1.30121, 1.20177, 1.15255, 1.06057, 1.11965, 1.07252, 1.11186, 1.04809, 0.99592, 1.02337, 0.97258, 0.85537, 0.90561, 0.92244, 0.89452, 0.87062, 0.87724, 0.79011, 0.81471, 0.77344, 0.89019, 0.84675, 0.7863, 0.79655, 0.79136, 0.78903, 0.69109, 0.79933, 0.76122, 0.74675, 0.74444, 0.70679, 0.77344, 0.75115, 0.64208, 0.71032, 0.65917, 0.60391, 0.70882, 0.61943, 0.7233, 0.60813, 0.57173, 0.63867, 0.55988, 0.58275, 0.53571, 0.50102, 0.47117, 0.42009, 0.38298, 0.43448, 0.23533], [3.19289, 3.20332, 3.13118, 3.29195, 3.21025, 3.11974, 3.09509, 3.24822, 3.29626, 3.14663, 3.14776, 3.13062, 3.24695, 3.10309, 3.07283, 3.12355, 3.19411, 3.19611, 3.10147, 3.24791, 3.15023, 3.20357, 3.08551, 3.2674, 3.22793, 3.09495, 3.25732, 3.19573, 3.11871, 3.15339, 3.18287, 3.16111, 3.13278, 3.18914, 3.09748, 3.11531, 3.18987, 3.13633, 3.13197, 3.11869, 3.23783, 3.08722, 3.10859, 3.11983, 3.13705, 3.20026, 3.13958, 3.10017, 3.18473, 3.0822, 3.16972, 3.11158, 3.06028, 3.1369, 3.07121, 3.16635, 3.06199, 3.10235, 3.0049, 3.07904, 3.17164, 2.96408, 3.02321, 3.02179, 3.06421, 2.974, 3.07155, 3.04761, 3.01762, 3.04448, 3.06716, 2.94967, 3.13895, 3.03683, 2.91446, 2.99686, 2.99002, 2.90381, 2.99524, 2.92879, 2.93743, 2.90347, 2.95506, 2.98542, 2.95227, 2.85467, 2.85478, 2.90819, 2.81035, 2.893, 3.00207, 2.86355, 2.92082, 2.92556, 2.80703, 2.80813, 2.81539, 2.85791, 2.81084, 2.85916, 2.66897, 2.62131, 2.66889, 2.60001, 2.72496, 2.65167, 2.69877, 2.55645, 2.51446, 2.41763, 2.46838, 2.43612, 2.48679, 2.4314, 2.40386, 2.38831, 2.31003, 2.29669, 2.24604, 2.33123, 2.24731, 2.26932, 2.20914, 2.15656, 2.27917, 1.99678, 2.00843, 2.07626, 2.01981, 2.02367, 1.93297, 1.86546, 1.8737, 1.74425, 1.77206, 1.72595, 1.6856, 1.68617, 1.61236, 1.59955, 1.65111, 1.5468, 1.49649, 1.39286, 1.40001, 1.38538, 1.27518, 1.29637, 1.24033, 1.28538, 1.13003, 1.0992, 1.08888, 1.1286, 1.13199, 1.04911, 1.00675, 1.01923, 0.93289, 0.91551, 0.91603, 0.92905, 0.87201, 0.85094, 0.8502, 0.86083, 0.81319, 0.8325, 0.85632, 0.78461, 0.80661, 0.78982, 0.76658, 0.74242, 0.77452, 0.86511, 0.7526, 0.75133, 0.68739, 0.75834, 0.83831, 0.7594, 0.75613, 0.75215, 0.68396, 0.75373, 0.66607, 0.66716, 0.64485, 0.57796, 0.61682, 0.60869, 0.64596, 0.59495, 0.51783, 0.48442, 0.41417, 0.36773, 0.44011, 0.23689], [3.37818, 3.4295, 3.32981, 3.39272, 3.38651, 3.268, 3.23625, 3.35117, 3.40409, 3.29405, 3.30436, 3.31998, 3.36086, 3.2574, 3.28336, 3.27311, 3.29237, 3.35928, 3.26916, 3.31437, 3.36866, 3.31251, 3.22481, 3.38328, 3.33019, 3.24876, 3.37842, 3.30472, 3.27889, 3.38972, 3.3586, 3.33149, 3.36475, 3.31092, 3.32532, 3.27716, 3.35026, 3.29164, 3.33912, 3.28857, 3.35237, 3.2818, 3.22231, 3.21942, 3.29937, 3.38346, 3.21985, 3.22702, 3.30349, 3.2728, 3.30489, 3.28141, 3.35561, 3.22667, 3.21544, 3.38781, 3.20666, 3.21179, 3.19872, 3.15896, 3.3226, 3.14365, 3.22858, 3.14938, 3.22504, 3.19481, 3.26758, 3.2346, 3.21266, 3.24453, 3.19085, 3.12807, 3.25126, 3.14237, 3.10679, 3.15485, 3.0806, 3.19298, 3.12594, 3.11598, 3.15629, 3.00586, 3.08751, 3.07022, 3.18309, 3.004, 3.05763, 3.14507, 2.98242, 3.15756, 3.10434, 2.99123, 3.04945, 3.07908, 2.89691, 2.97349, 3.00445, 2.96261, 2.93783, 2.9976, 2.81459, 2.76322, 2.87161, 2.76498, 2.91464, 2.80025, 2.78643, 2.73055, 2.67852, 2.55905, 2.64193, 2.60446, 2.62491, 2.50699, 2.56031, 2.55452, 2.4532, 2.40654, 2.46226, 2.42591, 2.35981, 2.35727, 2.21805, 2.30942, 2.39933, 2.10221, 2.11791, 2.13281, 2.08731, 2.05424, 2.04794, 1.99453, 1.95098, 1.80172, 1.9307, 1.8774, 1.85333, 1.80035, 1.77942, 1.7042, 1.64524, 1.5701, 1.54015, 1.50618, 1.49893, 1.45109, 1.32074, 1.42309, 1.25068, 1.28895, 1.19161, 1.25326, 1.1978, 1.16706, 1.22731, 1.09005, 1.03965, 1.09228, 0.96609, 0.96456, 0.97247, 0.9117, 0.95574, 0.89686, 0.89769, 0.92322, 0.81357, 0.91154, 0.89964, 0.81597, 0.79613, 0.84202, 0.80943, 0.77262, 0.90455, 0.83815, 0.81172, 0.79289, 0.728, 0.76648, 0.89912, 0.78447, 0.76121, 0.77736, 0.77273, 0.6895, 0.67911, 0.74645, 0.68729, 0.64679, 0.65137, 0.60097, 0.60363, 0.58089, 0.56474, 0.51068, 0.47838, 0.40479, 0.41969, 0.27905], [3.45301, 3.54717, 3.4596, 3.47587, 3.54046, 3.40331, 3.38482, 3.51748, 3.54336, 3.46751, 3.45968, 3.4795, 3.50039, 3.41229, 3.39587, 3.42847, 3.43456, 3.57913, 3.52596, 3.48395, 3.57092, 3.43764, 3.40874, 3.53676, 3.40964, 3.39528, 3.54721, 3.4197, 3.42946, 3.47917, 3.52677, 3.46229, 3.51521, 3.4927, 3.46864, 3.45446, 3.5334, 3.43759, 3.46917, 3.42714, 3.44916, 3.39819, 3.42604, 3.38423, 3.45926, 3.49558, 3.4668, 3.36662, 3.4788, 3.42074, 3.42612, 3.36835, 3.48685, 3.41875, 3.31003, 3.46762, 3.39454, 3.48105, 3.34501, 3.32963, 3.43769, 3.2927, 3.38675, 3.37021, 3.31648, 3.35511, 3.31156, 3.33986, 3.38425, 3.27367, 3.28704, 3.24931, 3.36041, 3.30686, 3.28591, 3.33311, 3.37636, 3.32563, 3.28749, 3.27392, 3.31845, 3.18109, 3.29035, 3.2906, 3.24871, 3.19364, 3.19441, 3.25998, 3.25125, 3.26543, 3.1979, 3.13131, 3.18998, 3.21753, 3.0947, 3.04775, 3.07987, 3.1097, 3.07547, 3.09425, 2.96571, 2.90049, 2.99277, 2.88027, 2.94319, 2.93843, 2.92938, 2.74973, 2.85985, 2.75994, 2.7209, 2.72269, 2.79425, 2.69455, 2.65603, 2.61492, 2.6297, 2.55784, 2.53342, 2.46743, 2.47238, 2.51682, 2.39722, 2.39615, 2.483, 2.25951, 2.32966, 2.39053, 2.17893, 2.1878, 2.1899, 2.18898, 2.0122, 1.9554, 2.00041, 2.0044, 1.93296, 1.87236, 1.87714, 1.79056, 1.77574, 1.66991, 1.62939, 1.64213, 1.56681, 1.54654, 1.53242, 1.50248, 1.42946, 1.33256, 1.28468, 1.27894, 1.28488, 1.24691, 1.27329, 1.18784, 1.10191, 1.11608, 1.06125, 1.04563, 1.01995, 0.99509, 1.02159, 0.91851, 0.92137, 0.94517, 0.93674, 0.94291, 0.86937, 0.92664, 0.87062, 0.86865, 0.84939, 0.89536, 0.96442, 0.83353, 0.86622, 0.88632, 0.85741, 0.83835, 0.897, 0.81296, 0.78341, 0.80402, 0.79045, 0.77768, 0.68289, 0.75662, 0.74469, 0.6454, 0.70259, 0.62594, 0.66973, 0.63104, 0.57764, 0.54656, 0.5149, 0.45211, 0.43891, 0.31838], [3.61713, 3.66843, 3.57499, 3.6992, 3.66139, 3.63133, 3.60992, 3.59796, 3.71214, 3.60397, 3.6604, 3.60093, 3.66519, 3.5744, 3.59816, 3.58898, 3.68096, 3.70525, 3.65361, 3.61855, 3.60859, 3.64611, 3.51597, 3.69308, 3.5018, 3.55028, 3.65746, 3.55612, 3.56723, 3.56024, 3.64178, 3.5331, 3.61779, 3.65581, 3.63479, 3.51192, 3.6095, 3.55331, 3.58343, 3.56697, 3.55165, 3.54466, 3.56945, 3.56452, 3.56764, 3.60176, 3.60566, 3.50674, 3.63727, 3.51895, 3.53621, 3.46425, 3.56547, 3.57512, 3.48257, 3.62521, 3.50364, 3.58131, 3.50373, 3.54714, 3.5253, 3.44502, 3.44995, 3.4768, 3.47568, 3.46933, 3.46542, 3.44973, 3.4116, 3.48635, 3.41002, 3.4635, 3.60598, 3.36312, 3.37688, 3.45509, 3.42031, 3.40811, 3.39346, 3.51968, 3.37968, 3.30681, 3.41998, 3.42048, 3.40245, 3.32106, 3.35407, 3.36684, 3.36996, 3.39425, 3.37352, 3.27134, 3.33416, 3.36043, 3.2468, 3.19623, 3.25601, 3.18945, 3.2042, 3.25372, 3.17041, 3.09408, 3.09549, 3.00257, 3.09268, 3.12798, 3.02715, 3.02147, 2.97063, 2.8702, 2.88779, 2.87573, 2.96007, 2.86396, 2.82503, 2.85046, 2.74288, 2.83987, 2.75592, 2.72686, 2.61559, 2.78691, 2.52969, 2.52376, 2.64935, 2.47175, 2.43027, 2.44512, 2.32001, 2.35346, 2.29702, 2.26902, 2.19594, 2.17455, 2.13723, 2.21828, 2.17045, 2.11348, 2.05186, 1.90909, 1.94368, 1.85471, 1.82315, 1.71818, 1.6914, 1.69966, 1.59071, 1.59465, 1.58102, 1.49242, 1.3664, 1.36725, 1.34164, 1.36005, 1.36834, 1.3094, 1.21212, 1.21973, 1.14436, 1.16586, 1.13201, 1.17614, 1.09794, 1.03729, 1.06728, 1.07631, 1.03906, 1.09235, 1.06696, 1.03104, 0.9723, 0.97222, 0.91632, 0.95428, 0.91088, 0.90223, 0.94564, 0.94636, 0.84496, 0.95977, 1.02479, 0.865, 0.86768, 0.84073, 0.80731, 0.91962, 0.72997, 0.85239, 0.75502, 0.67988, 0.75747, 0.68226, 0.69764, 0.68729, 0.6821, 0.5413, 0.57141, 0.46762, 0.45146, 0.30402], [3.7197, 3.75288, 3.73948, 3.77312, 3.81365, 3.78157, 3.66759, 3.72636, 3.74291, 3.74803, 3.73465, 3.70506, 3.83496, 3.72305, 3.72498, 3.69183, 3.75015, 3.8738, 3.70245, 3.7183, 3.7553, 3.78608, 3.64957, 3.79166, 3.81488, 3.67751, 3.81057, 3.73817, 3.75199, 3.69218, 3.80472, 3.73387, 3.7714, 3.73614, 3.80875, 3.61156, 3.72281, 3.63094, 3.72263, 3.82989, 3.72763, 3.64282, 3.72133, 3.65832, 3.73877, 3.75626, 3.69515, 3.6809, 3.69723, 3.66151, 3.66526, 3.66839, 3.82306, 3.70673, 3.65888, 3.7321, 3.69242, 3.72043, 3.64903, 3.65518, 3.65633, 3.54972, 3.68196, 3.59492, 3.65627, 3.55864, 3.69094, 3.66159, 3.60427, 3.68031, 3.61134, 3.59775, 3.69613, 3.59021, 3.53424, 3.57728, 3.51997, 3.60853, 3.54904, 3.6011, 3.57333, 3.52464, 3.60692, 3.57399, 3.49841, 3.54567, 3.56203, 3.59262, 3.45313, 3.57509, 3.49024, 3.53294, 3.47659, 3.52539, 3.4238, 3.31427, 3.36263, 3.33281, 3.45233, 3.47506, 3.33258, 3.35411, 3.28456, 3.25264, 3.32047, 3.25435, 3.23093, 3.2172, 3.16335, 3.16917, 3.02875, 3.05147, 3.12606, 3.14517, 3.0664, 3.0573, 2.93829, 2.95637, 2.94692, 2.89459, 2.88349, 3.00862, 2.86556, 2.82761, 2.85602, 2.74692, 2.57108, 2.72543, 2.58654, 2.51087, 2.47341, 2.46957, 2.44169, 2.38086, 2.33281, 2.42138, 2.37016, 2.34938, 2.22008, 2.12459, 2.21601, 2.1142, 1.96385, 1.97391, 2.01273, 1.95927, 1.84508, 1.88961, 1.77264, 1.81971, 1.62878, 1.63765, 1.63094, 1.61532, 1.61961, 1.56051, 1.54983, 1.44875, 1.41198, 1.42107, 1.46421, 1.45253, 1.27968, 1.23784, 1.2987, 1.28364, 1.34296, 1.2227, 1.26018, 1.20951, 1.24435, 1.22171, 1.18628, 1.19984, 1.15413, 1.21284, 1.14366, 1.1792, 1.10195, 1.0899, 1.1852, 1.07179, 1.09002, 1.13247, 1.00483, 1.05028, 0.98792, 1.08739, 0.93726, 0.98243, 1.05575, 0.93122, 0.95288, 0.83264, 0.82001, 0.68427, 0.69098, 0.56786, 0.55199, 0.42048]]}

runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/utils.py

@@ -0,0 +1,116 @@
+import torch
+from typing import Optional
+from datasets import load_dataset, Features, Sequence, Value
+
+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
+
+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 load_and_process_dataset(data_name: str):
+    # Math datasets
+    if data_name == "gsm8k":
+        dataset = load_dataset("openai/gsm8k", "main", split="test")
+        prompt_fmt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "math500":
+        dataset = load_dataset("HuggingFaceH4/MATH-500", split="test")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "aime24":
+        dataset = load_dataset("HuggingFaceH4/aime_2024", split="train")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+
+    elif data_name == "aime25":
+        dataset = load_dataset("MathArena/aime_2025", split="train")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+
+    # Chat datasets 
+    elif data_name == "alpaca":
+        dataset = load_dataset("tatsu-lab/alpaca", split="train")
+        dataset = dataset.map(lambda x: {"formatted_input": (f"{x['instruction']}\n\nInput:\n{x['input']}" if x['input'] else x['instruction'])})
+        dataset = dataset.map(lambda x: {"turns": [x["formatted_input"]]})
+
+    elif data_name == "mt-bench":
+        dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
+        dataset = dataset.map(lambda x: {"turns": x["prompt"]})
+
+    # Coding datasets
+    elif data_name == "humaneval":
+        dataset = load_dataset("openai/openai_humaneval", split="test")
+        prompt_fmt = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{prompt}\n```"
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "mbpp":
+        dataset = load_dataset("google-research-datasets/mbpp", "sanitized", split="test")
+        dataset = dataset.map(lambda x: {"turns": [x["prompt"]]})
+    
+    elif data_name == "lbpp":
+        LBPP_PY_TEST_URL = "https://huggingface.co/datasets/CohereLabs/lbpp/resolve/main/python/test.parquet"
+        dataset = load_dataset("parquet", data_files={"test": LBPP_PY_TEST_URL})["test"]
+        dataset = dataset.map(lambda x: {"turns": [x["instruction"]]})
+
+    elif data_name == "swe-bench":
+        dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test")
+        prompt_fmt = "Problem Statement:\n{problem_statement}\nPlease fix the issue described above."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "livecodebench":
+        base = "https://huggingface.co/datasets/livecodebench/code_generation_lite/resolve/main/"
+        allowed_files = ["test.jsonl", "test2.jsonl", "test3.jsonl", "test4.jsonl", "test5.jsonl", "test6.jsonl"]
+        urls = [base + fn for fn in allowed_files]
+        dataset = load_dataset("json", data_files={"test": urls})["test"]
+        def format_lcb(doc):
+            system_prompt = (
+                "You are an expert Python programmer. You will be given a question (problem specification) "
+                "and will generate a correct Python program that matches the specification and passes all tests. "
+                "You will NOT return anything except for the program"
+            )
+            question_block = f"### Question:\n{doc['question_content']}"
+            if doc.get("starter_code"):
+                format_message = "### Format: Use the following code structure:"
+                code_block = f"```python\n{doc['starter_code']}\n```"
+            else:
+                format_message = "### Format: Write your code in the following format:"
+                code_block = "```python\n# YOUR CODE HERE\n```"
+            answer_footer = "### Answer: (use the provided format with backticks)"
+            return f"{system_prompt}\n\n{question_block}\n\n{format_message}\n{code_block}\n\n{answer_footer}"
+        target_features = Features({"turns": Sequence(Value("large_string"))})
+        dataset = dataset.map(
+            lambda x: {"turns": [format_lcb(x)]},
+            remove_columns=dataset.column_names,
+            features=target_features
+        )
+    
+    return dataset

runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6d49e24d000733a4b8d9ea5be4a4e660924163512515f45114af9c2551f2946
+size 1797091260

runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39248be7fbbbf3833d300044811fe41539eaa3177644112b719816b519ec70f2
+size 311165128

runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d843fb298102997a57d16a9d0924b5dfbb7058f30a5dc8045d39a3a9f8d72847
+size 1074860568

runs/fmlm2_v3_6ep_ce/export/epoch_1/dflash.py

@@ -0,0 +1,277 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+from .utils import build_target_layer_ids, extract_context_feature, sample
+
+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
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    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
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    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)
+
+    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)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["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)]
+        )
+        self.target_layer_ids = self.config.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 = self.config.dflash_config.get("mask_token_id", None)
+        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

runs/fmlm2_v3_6ep_ce/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:418550eb1cc8e6b3e131beea42e9a7d52540fd63688c3b0da8cca9f8116c35c8
+size 1770876772

runs/fmlm2_v3_6ep_ce/export/epoch_1/flow_config.json

@@ -0,0 +1,10 @@
+{
+  "learn_embed": true,
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "objective": "fmlm2",
+  "kd": true,
+  "readout_frac": 0.25,
+  "head_rank": 256
+}

runs/fmlm2_v3_6ep_ce/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5029aa2fcc79d191a60a735d3c368d6103fcd12e35d7268b6d3bf347fd8f8e3
+size 311165128

runs/fmlm2_v3_6ep_ce/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb3cc23203f226b2b9df422557614d7119f684904244ef7c13b56f5d932f2e19
+size 1074860568

runs/fmlm2_v3_6ep_ce/export/epoch_1/utils.py

@@ -0,0 +1,116 @@
+import torch
+from typing import Optional
+from datasets import load_dataset, Features, Sequence, Value
+
+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
+
+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 load_and_process_dataset(data_name: str):
+    # Math datasets
+    if data_name == "gsm8k":
+        dataset = load_dataset("openai/gsm8k", "main", split="test")
+        prompt_fmt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "math500":
+        dataset = load_dataset("HuggingFaceH4/MATH-500", split="test")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "aime24":
+        dataset = load_dataset("HuggingFaceH4/aime_2024", split="train")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+
+    elif data_name == "aime25":
+        dataset = load_dataset("MathArena/aime_2025", split="train")
+        prompt_fmt = "{problem}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+
+    # Chat datasets 
+    elif data_name == "alpaca":
+        dataset = load_dataset("tatsu-lab/alpaca", split="train")
+        dataset = dataset.map(lambda x: {"formatted_input": (f"{x['instruction']}\n\nInput:\n{x['input']}" if x['input'] else x['instruction'])})
+        dataset = dataset.map(lambda x: {"turns": [x["formatted_input"]]})
+
+    elif data_name == "mt-bench":
+        dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
+        dataset = dataset.map(lambda x: {"turns": x["prompt"]})
+
+    # Coding datasets
+    elif data_name == "humaneval":
+        dataset = load_dataset("openai/openai_humaneval", split="test")
+        prompt_fmt = "Write a solution to the following problem and make sure that it passes the tests:\n```python\n{prompt}\n```"
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "mbpp":
+        dataset = load_dataset("google-research-datasets/mbpp", "sanitized", split="test")
+        dataset = dataset.map(lambda x: {"turns": [x["prompt"]]})
+    
+    elif data_name == "lbpp":
+        LBPP_PY_TEST_URL = "https://huggingface.co/datasets/CohereLabs/lbpp/resolve/main/python/test.parquet"
+        dataset = load_dataset("parquet", data_files={"test": LBPP_PY_TEST_URL})["test"]
+        dataset = dataset.map(lambda x: {"turns": [x["instruction"]]})
+
+    elif data_name == "swe-bench":
+        dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test")
+        prompt_fmt = "Problem Statement:\n{problem_statement}\nPlease fix the issue described above."
+        dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
+    
+    elif data_name == "livecodebench":
+        base = "https://huggingface.co/datasets/livecodebench/code_generation_lite/resolve/main/"
+        allowed_files = ["test.jsonl", "test2.jsonl", "test3.jsonl", "test4.jsonl", "test5.jsonl", "test6.jsonl"]
+        urls = [base + fn for fn in allowed_files]
+        dataset = load_dataset("json", data_files={"test": urls})["test"]
+        def format_lcb(doc):
+            system_prompt = (
+                "You are an expert Python programmer. You will be given a question (problem specification) "
+                "and will generate a correct Python program that matches the specification and passes all tests. "
+                "You will NOT return anything except for the program"
+            )
+            question_block = f"### Question:\n{doc['question_content']}"
+            if doc.get("starter_code"):
+                format_message = "### Format: Use the following code structure:"
+                code_block = f"```python\n{doc['starter_code']}\n```"
+            else:
+                format_message = "### Format: Write your code in the following format:"
+                code_block = "```python\n# YOUR CODE HERE\n```"
+            answer_footer = "### Answer: (use the provided format with backticks)"
+            return f"{system_prompt}\n\n{question_block}\n\n{format_message}\n{code_block}\n\n{answer_footer}"
+        target_features = Features({"turns": Sequence(Value("large_string"))})
+        dataset = dataset.map(
+            lambda x: {"turns": [format_lcb(x)]},
+            remove_columns=dataset.column_names,
+            features=target_features
+        )
+    
+    return dataset

runs/fmlm2_v3_6ep_tv/export/epoch_1/config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [
+      1,
+      9,
+      17,
+      25,
+      33
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

runs/fmlm2_v3_6ep_tv/export/epoch_1/dflash.py

@@ -0,0 +1,277 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+from .utils import build_target_layer_ids, extract_context_feature, sample
+
+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
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    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
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    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)
+
+    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)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["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)]
+        )
+        self.target_layer_ids = self.config.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 = self.config.dflash_config.get("mask_token_id", None)
+        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

runs/fmlm2_v3_6ep_tv/export/epoch_1/dspark_config.json

@@ -0,0 +1,16 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_tv/export/epoch_1",
+    "mode": "onejump_zero"
+  },
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "joint": true,
+  "objective": "fmlm2",
+  "kd": true,
+  "readout_frac": 0.25,
+  "head_rank": 256
+}

runs/fmlm2_v3_6ep_tv/export/epoch_1/flow.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44c4455cbf38f0838eb86bc64d2cd1f92a41ebde5eb813e4b3aac298abc10ede
+size 1770876772

runs/fmlm2_v3_6ep_tv/export/epoch_1/flow_config.json

@@ -0,0 +1,10 @@
+{
+  "learn_embed": true,
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1",
+  "objective": "fmlm2",
+  "kd": true,
+  "readout_frac": 0.25,
+  "head_rank": 256
+}

runs/fmlm2_v3_6ep_tv/export/epoch_1/head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2a33a09f0e5a5804b0e154d8dab1ee9b25a1d8273bd27609495ce3f8f7ee906
+size 311165128

runs/fmlm2_v3_6ep_tv/export/epoch_1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dba373f994baac3322c5f8f5b771d6aceb2edff9551ddcaa0bbbe40665755f9f
+size 1074860568

runs/fmlm2_v3_6ep_tv/export/epoch_1/modeling_dflash.py

@@ -0,0 +1,309 @@
+from typing import Optional, Callable
+from typing_extensions import Unpack, Tuple
+import torch
+from torch import nn
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    Qwen3Config,
+    Qwen3PreTrainedModel,
+    Qwen3MLP,
+    GradientCheckpointingLayer,
+    FlashAttentionKwargs,
+    rotate_half,
+    eager_attention_forward,
+    ALL_ATTENTION_FUNCTIONS,
+)
+from transformers import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+
+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 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
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    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
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    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)
+
+    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)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        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
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["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)]
+        )
+        self.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.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,
+        mask_token_id: int,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        self.eval()
+        target.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),
+            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] != 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