Add files using upload-large-folder tool

DATA_SHA256SUMS

@@ -0,0 +1,3 @@
+e3be6bda242c860dd6f2e15d086b4dcd08b5ea4826ddee9455ad5c9462be4d8f  data/tokens.pt
+5d4f4fe7735468a7cb4715b6c6711cc76543686038dd82a9d890d54427549489  data/data.jsonl
+1fc0023cdb9cd6b93336abbb5a36f8f08527368002d68ca9a1a32b72aeb93073  init/dflash100k_epoch6/model.safetensors

README.md

@@ -0,0 +1,76 @@
+---
+license: mit
+base_model: Qwen/Qwen3-4B
+tags:
+- speculative-decoding
+- dflash
+- flow-matching
+- drafter
+---
+
+# FlowDraft checkpoints
+
+Drafter weights for **FlowDraft**: flow-map drafters for speculative decoding with **Qwen3-4B**.
+Code, reports and full instructions: https://github.com/ProtonPratt/FlowDraft-Main
+
+A drafter guesses the next 15 tokens of Qwen3-4B in one forward pass. Qwen3-4B checks all 15 guesses in one pass
+and keeps the longest correct prefix. The output is exactly what Qwen3-4B writes alone. The score is **τ**: tokens
+written per guess-and-check round.
+
+## Layout
+
+This repo has the same layout as the GitHub repo. Download into the clone and every file is at the path the code
+reads:
+
+```bash
+git clone https://github.com/ProtonPratt/FlowDraft-Main.git && cd FlowDraft-Main
+python scripts/hf_checkpoints.py fetch                 # DFlash-100K + the best drafter (4.1 GB)
+python scripts/hf_checkpoints.py fetch --set all       # everything here (~62 GB)
+sha256sum -c WEIGHTS_SHA256SUMS --ignore-missing
+```
+
+Run the code from the repo root. The export configs (`dspark_config.json`) hold repo-relative paths.
+
+## Exports (drafters)
+
+Results: 480 benchmark answers (80 per set: GSM8K, MATH-500, HumanEval, MBPP, MT-Bench), one H100, batch size 1,
+overhead-free benchmark loop.
+
+| Path | What it is | greedy τ | t1 τ | greedy speed | t1 speed |
+|---|---|---:|---:|---:|---:|
+| `init/dflash100k_epoch6` | DFlash-100K (baseline; also the start of every run) | 4.78 | 4.08 | 3.59x | 3.05x |
+| `runs/fmlm2_v3_6ep_bv/export/epoch_1` | **best**: fmlm2 V3, 6 passes + BV pass | **5.31** | **4.70** | 3.50x | **3.11x** |
+| `runs/fmlm2_v3_6ep_bv_lt/export/epoch_1` | BV pass + Lafa learned time per slot | 5.30 | 4.71 | | |
+| `runs/fmlm2_v3_6ep_bv_emb/export/epoch_1` | BV pass + Lafa embedding loss | 5.27 | 4.65 | | |
+| `runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1` | Lafa A + self-conditioning, two-step drafting | | | | |
+| `runs/fmlm2_v3_6ep_tv/export/epoch_1` | 6 passes + TV pass | 5.26 | 4.70 | | |
+| `runs/fmlm2_v3_6ep_ce/export/epoch_1` | 6 passes + CE pass (control) | 5.17 | 4.37 | | |
+| `runs/fmlm2_v3_6ep/export/epoch_6` | fmlm2 V3 after 6 passes (start of every extra pass) | | | | |
+| `runs/fmlm2_v3/export/epoch_1` | fmlm2 V3 at 1 pass | | | | |
+| `runs/dspark_joint_6p/export/epoch_6` | DSpark (no flow) after 6 passes | | | | |
+| `runs/dspark_joint_6p_tv/export/epoch_1` | DSpark 6 passes + TV pass | 5.20 | 4.64 | 3.48x | 3.09x |
+| `runs/dspark_a/export/epoch_1` | Markov head on frozen DFlash-100K (head only) | 5.04 | 4.36 | 3.37x | 2.90x |
+| `runs/full_scratch/export/epoch_6` | plain FMLM+ from scratch, 6 passes | | | | |
+
+Empty cells: see the reports in the GitHub repo (`reports/README.md`). An fmlm2 export has `model.safetensors`
+(DFlash layout), `flow.safetensors` (flow parts) and `head.safetensors` (joint head). A DSpark export has
+`model.safetensors` + `head.safetensors`.
+
+## Resume checkpoints
+
+Full training state (weights, optimizer, schedule, data cursor). Use them only to continue training. Resume needs
+8 GPUs and the same training settings; `src/train.py` checks both.
+
+| Path | Size | What it is |
+|---|---:|---|
+| `runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt` | 12 GB | fmlm2 V3 at the end of pass 6 |
+| `runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt` | 6.8 GB | DSpark at the end of pass 6 |
+| `runs/fmlm2_v3_tv1p/ckpt/step_1250.pt` | 12 GB | 1-pass matched TV run, paused at step 1,250 |
+
+## Limits
+
+- Trained on 99,999 **greedy** Qwen3-4B answers, one seed per run.
+- The target model, Qwen3-4B, is not in this repo. Download it from `Qwen/Qwen3-4B` (Apache-2.0).
+- The model code is upstream DFlash (z-lab/dflash, MIT).
+
+Team TokenYahu: Vipen Loka, Dheeraj Hegde, Pratyush Jena.

WEIGHTS_SHA256SUMS

@@ -0,0 +1,88 @@
+af98cf412711c7d9f7df177dd1bb20b6576bb8233b9d14da5b0b6d3aea24a717  runs/dspark_a/export/epoch_1/head.safetensors
+c13ca1bdb106deb0d2976764978b31edd7348249b1b45bb4c2d3cc5421ad581b  runs/dspark_b/export/epoch_1/head.safetensors
+e6a43cbc9fd15b3e1a45e5ec67585ef00a484df4ac8d32a5756c9dfc6fad989f  runs/dspark_joint_1p/export/epoch_1/head.safetensors
+746f093afcc86db5157a5f3a3ebf4d19f0bd4c5dc7656dc47453450b5cdea263  runs/dspark_joint_1p/export/epoch_1/model.safetensors
+4c29ef39cad74a7e958f4c18a31d77e97080db7c1033be6012732e84400b1fd8  runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt
+fb37b0fb99a46d77fcc3f400d9025bbdb7ad9df30598b7b8cd4b518aac6e0adf  runs/dspark_joint_6p/export/epoch_1/head.safetensors
+efdcaa86ddd35f91ed2d00359c8db68606168f947a2a9826acf1b87e73a1379d  runs/dspark_joint_6p/export/epoch_1/model.safetensors
+dc9877a12fdf3968deb6cd9280c0893a15adeaa80fe9019cf6982cb75870f228  runs/dspark_joint_6p/export/epoch_2/head.safetensors
+166d31c81d0211258d931800af1503fa2e7f6d53847f5dfec6e6aeebb98973b7  runs/dspark_joint_6p/export/epoch_2/model.safetensors
+9f4d661ebc20971a8582269ae753ee307ed01613b974b1904b8fbf03c3553397  runs/dspark_joint_6p/export/epoch_3/head.safetensors
+22dc256624d34fef7cb19a9489bd94d0c50f717f195a95556c991d196841f625  runs/dspark_joint_6p/export/epoch_3/model.safetensors
+283277fb2a9e86d0c8c7c6869a9fc8aa8e795eedc846422e430b2cf0b08531b5  runs/dspark_joint_6p/export/epoch_4/head.safetensors
+be252f7217184734db733e5d6ba9278fe4fc72027879b9de26cc0e5d4057d771  runs/dspark_joint_6p/export/epoch_4/model.safetensors
+4bb3c51602d505d8e756e5958f673045fd64c167766bf40bcd7941494d4edfcf  runs/dspark_joint_6p/export/epoch_5/head.safetensors
+ad09d74cedc943b623b0c16c7bf64ab8223f971360c794012051d3c4f6e7cff7  runs/dspark_joint_6p/export/epoch_5/model.safetensors
+e740457069c8cdaf7feebc6f1c9a45af543c8208fa4de0aea29af7a8cf70ddaa  runs/dspark_joint_6p/export/epoch_6/head.safetensors
+fcb5d93581ed4fe4cd7405bcebdbf0fe358b8bf0ff2248d2b9e681eba602ddc9  runs/dspark_joint_6p/export/epoch_6/model.safetensors
+ebfa3ff872cf4defd1c861d422cb4565c3a288ba4862f7a1c4044463e7282e3d  runs/dspark_joint_6p_tv/export/epoch_1/head.safetensors
+a6e2aa3c88581da6c7c3cacb35f50f678b5fb164f047299deb0f2bc0dbe39589  runs/dspark_joint_6p_tv/export/epoch_1/model.safetensors
+0b59f8bccfb56dcfd168945aed82e84e726e64543f22e8c80e1723350cfd99d4  runs/fmlm2_v0/export/epoch_1/flow.safetensors
+36dd7f37a5f73723366f5cf091ca5ab3a52a8c1bb5192615cd27598ac66aa88f  runs/fmlm2_v0/export/epoch_1/model.safetensors
+4cd79d686c4ae11f63084a4526f75afacedc7b78936e4d76c627ff8e15273153  runs/fmlm2_v1/export/epoch_1/flow.safetensors
+cbb1ccdee9181b14d8b4a302bb040bb4ba49641180f4d8832920b01a696a9ba3  runs/fmlm2_v1/export/epoch_1/model.safetensors
+64d5c2f1084327c572c28aae885227f623bfc5681fc8061b6ca13a2c025ea3c4  runs/fmlm2_v2/export/epoch_1/flow.safetensors
+96ac9001b6d4bf869fd8f7e0ea9d0d6e02e5eb1f5caf0d45e6dc000e6e9ff761  runs/fmlm2_v2/export/epoch_1/model.safetensors
+38594ad798d0869e11e0b7f1aadfeb8fb075f0739a489e129cd28e8ffac4f5f1  runs/fmlm2_v3/export/epoch_1/flow.safetensors
+acb621a7e9913a6e9df3e151222a2b0110b2b06cb7b9f6e37c225545b401bd55  runs/fmlm2_v3/export/epoch_1/head.safetensors
+be6d410548cea0f60bdde21e66fc073ae9d2aa059160c8bfad6a64fd366950cf  runs/fmlm2_v3/export/epoch_1/model.safetensors
+1aaec2649c7dc202b6931d991adafb7d52a7e354bdef79920cca2a06ad8a66da  runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt
+8b888f1c23cd327ec52c7f00c6ceacf0ca2d054bfb11a17dc2edc68f83beaa70  runs/fmlm2_v3_6ep/export/epoch_1/flow.safetensors
+48e79cb8bd7a69b588f4c8aa8dbf2d08fed0a368a6da8023b030bbd01eccaddd  runs/fmlm2_v3_6ep/export/epoch_1/head.safetensors
+9ac25c511b304356b7e16b9d23d3ea168e828516bc3c842ab4eee4fa2e41835f  runs/fmlm2_v3_6ep/export/epoch_1/model.safetensors
+08437009317e083e1d604584e6ee8491ff8ec641daebda5813e6ecdb5d7e0b69  runs/fmlm2_v3_6ep/export/epoch_2/flow.safetensors
+8a5d57b02569c76ac520827bb987965ed077af187b761770519256da36d9bb72  runs/fmlm2_v3_6ep/export/epoch_2/head.safetensors
+5bd14b4605c260f202515942b0cc5f39e9a65733ed1aea1b37975344dd12603b  runs/fmlm2_v3_6ep/export/epoch_2/model.safetensors
+2d189b359ce54312bb43c4bf6f71149a5df16c5d22db3b8de06f0b59a0aea090  runs/fmlm2_v3_6ep/export/epoch_3/flow.safetensors
+82217651ae08854e9fa842c31944c1f0d62331d3230326ac295088aafcc7fd79  runs/fmlm2_v3_6ep/export/epoch_3/head.safetensors
+e3837711b71074e6663a63355c30fbdcc9486054d0eda4f9fd7d44836ca3d30a  runs/fmlm2_v3_6ep/export/epoch_3/model.safetensors
+dc821b2d76134a5fbfff294355eb1e3fac89a4f41ae27d785073b8da732d2bad  runs/fmlm2_v3_6ep/export/epoch_4/flow.safetensors
+fcd281ae6434e60d361fcddd636a6b9c0a7be7c4ac99ac7985af8c3f63ece647  runs/fmlm2_v3_6ep/export/epoch_4/head.safetensors
+fd3588fed866220ef216f3b6257bb700b89ed639db9af41d49c6433f344b8fae  runs/fmlm2_v3_6ep/export/epoch_4/model.safetensors
+de82bb61766d86edea41d007b122a337d3d4fc7e23c4c605ed8084ec2284e51b  runs/fmlm2_v3_6ep/export/epoch_5/flow.safetensors
+19107b05598f495290a4f6cdb43c70d5b000a3484f6e67313249f963145f4d3b  runs/fmlm2_v3_6ep/export/epoch_5/head.safetensors
+3bb926d172224981ab4f80a5e0f018dd8cd10afbb2bdc38671b9a82c7c931f55  runs/fmlm2_v3_6ep/export/epoch_5/model.safetensors
+41f71c01d20ad59568b67a80e24b338f31e3ff71f3c9f16bcc60977b823ee340  runs/fmlm2_v3_6ep/export/epoch_6/flow.safetensors
+52f3ad42a9fb35d38661d0f09a7edc8e30220c8a0435f3ad8ebe70915fa683f5  runs/fmlm2_v3_6ep/export/epoch_6/head.safetensors
+066d537ccc3da6eb3da7b6855b6d3f226ebd89dbd0280168fab30e3f0b3c3ea7  runs/fmlm2_v3_6ep/export/epoch_6/model.safetensors
+4fc7410e6b5a259a6006266d871eea6fc7fce72eb9da4f1ebcee53c62eac0aa3  runs/fmlm2_v3_6ep_bv/export/epoch_1/flow.safetensors
+252f749694020bf01d7ed7f991776f852a25aea578ad72e431f9b8ea2e3fc329  runs/fmlm2_v3_6ep_bv/export/epoch_1/head.safetensors
+6a3d2040a940e3526e7e9df6abea13efebf51d377c942655545e3e8971249d20  runs/fmlm2_v3_6ep_bv/export/epoch_1/model.safetensors
+418550eb1cc8e6b3e131beea42e9a7d52540fd63688c3b0da8cca9f8116c35c8  runs/fmlm2_v3_6ep_ce/export/epoch_1/flow.safetensors
+f5029aa2fcc79d191a60a735d3c368d6103fcd12e35d7268b6d3bf347fd8f8e3  runs/fmlm2_v3_6ep_ce/export/epoch_1/head.safetensors
+bb3cc23203f226b2b9df422557614d7119f684904244ef7c13b56f5d932f2e19  runs/fmlm2_v3_6ep_ce/export/epoch_1/model.safetensors
+44c4455cbf38f0838eb86bc64d2cd1f92a41ebde5eb813e4b3aac298abc10ede  runs/fmlm2_v3_6ep_tv/export/epoch_1/flow.safetensors
+b2a33a09f0e5a5804b0e154d8dab1ee9b25a1d8273bd27609495ce3f8f7ee906  runs/fmlm2_v3_6ep_tv/export/epoch_1/head.safetensors
+dba373f994baac3322c5f8f5b771d6aceb2edff9551ddcaa0bbbe40665755f9f  runs/fmlm2_v3_6ep_tv/export/epoch_1/model.safetensors
+2f1a7eb42eaafb9dd24c780d46047658e94de9d2c89feb883c5e3908b1c8d123  runs/fmlm2_v3_sd3/export/epoch_1/flow.safetensors
+f3bea77b73c2148134aa6c9e87785231d4bd9c577d2426d3ac1ffc6bbc16d51b  runs/fmlm2_v3_sd3/export/epoch_1/head.safetensors
+e7e240532886e144bc0476df2ccb4ec006bd5b4447093247760ebefd94fd0dfe  runs/fmlm2_v3_sd3/export/epoch_1/model.safetensors
+946c8bac947dc3a2d6d7e74ba04b8b6496e7d3574d680258bf8351c3ce37183f  runs/fmlm2_v3_sd4/export/epoch_1/flow.safetensors
+18d4d696ba39c04623ec010097734bc3193bd03291510946ef7a0fe88034609c  runs/fmlm2_v3_sd4/export/epoch_1/head.safetensors
+e8bff64928ce119eeffcd909acb9c2d4a2cc7b5dcf326fe1923fd530869c2fea  runs/fmlm2_v3_sd4/export/epoch_1/model.safetensors
+9071b4ab56db18de4722350a15f045102032589e7c1793030b517db7f6fd5dd1  runs/fmlm2_v3_tv1p/ckpt/step_1250.pt
+54426ca7edad92181f401bb96776669636ea1d925d6f7c7a871fca42badb8574  runs/full_scratch/export/epoch_1/flow.safetensors
+f56dbb526f9ba97531027595e70e347589452cc90096069e0a1cd370d3f5febf  runs/full_scratch/export/epoch_1/model.safetensors
+7ff997f4cc74b45d2dc36ff5e41a80625303cb746d08a9cbc5b1f8c914006e2f  runs/full_scratch/export/epoch_2/flow.safetensors
+4ad056cc527c7c38d7a9298232de34e8c72da6b9043a276fefecc2a9c3b54eb5  runs/full_scratch/export/epoch_2/model.safetensors
+d11ccb31d0f629067b0d8b1365a4332cbf22421dd020c2f2fbd9c94dfb3c3558  runs/full_scratch/export/epoch_3/flow.safetensors
+2b5b7a4a41ed54d24e54046f232d3e8b01997e508de9972be7ecf25c2fd66479  runs/full_scratch/export/epoch_3/model.safetensors
+def9bfe4b15b8d43f7203fa10eed8dca39b17d04c44ef52aa83ded19ce70cf96  runs/full_scratch/export/epoch_4/flow.safetensors
+f78dd65f29cd666eac2a18924c77d9696d3924d21ee8962403d340acebb67bff  runs/full_scratch/export/epoch_4/model.safetensors
+b674b4fca0a931181afaa51bcb0dd530a39ca8f57bc1c1f39714d684350b8bfe  runs/full_scratch/export/epoch_5/flow.safetensors
+b3974dfea27535508796b4cff20dc0269ce8e67244138e3ffa66f8a600c68d61  runs/full_scratch/export/epoch_5/model.safetensors
+d9e232bedcd5475f9de6bb35fe7ef49f6fb04bf487bc8d1861a56b84a9ef3d30  runs/full_scratch/export/epoch_6/flow.safetensors
+698112c5e83c2e4bcad11b0d3328cba7576cc146c332e6fb8d58e1ff31f0a9c7  runs/full_scratch/export/epoch_6/model.safetensors
+a1a8149b77e16ff388576ef0322e159eeb13ce594af0ef08856e230d8aadf83a  runs/gsf_1ep/export/epoch_1/flow.safetensors
+b38c96dd5e904447782a08826e82d6bd67c6a5701c0cc590a7a7225ca231e1ea  runs/gsf_1ep/export/epoch_1/model.safetensors
+d1fcb89e8078ebae0bc4096139d84ce6dbb6b5743eec330012eaaf4122fe222a  runs/gsf_1ep_linear/export/epoch_1/flow.safetensors
+9f33654cd4b5156c1a71b9bd01a8bc3cde64d5a19979ce951d818342eeb69c9b  runs/gsf_1ep_linear/export/epoch_1/model.safetensors
+40b4b27ad2a1a5a4fa1fca89b55b411b39a0beaa4c34d4c25335786ab1256ef9  runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/flow.safetensors
+c15930d6cff545a25aad05afe29f81e89b750f1eb4705191473e1d12d8f85e9b  runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/head.safetensors
+c9b9a73b54c60f89efab5ce15c70454367ec1b30cc328bbc25cc5f9563d6fcb5  runs/fmlm2_v3_6ep_bv_lt/export/epoch_1/model.safetensors
+b3fd239bb0f2e19eee270407c0c41fe75012f6e9f9270559d2426e019ce8cec9  runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/flow.safetensors
+b10c48c786c1ff5aa26624953174f48feba3ad2c10b31143691a4e91fa071806  runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/head.safetensors
+fbc731478b05f229b4e00afed51d29e452f45ced4f588f4e8f85413a05eeeded  runs/fmlm2_v3_6ep_bv_emb/export/epoch_1/model.safetensors
+e6d49e24d000733a4b8d9ea5be4a4e660924163512515f45114af9c2551f2946  runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow.safetensors
+39248be7fbbbf3833d300044811fe41539eaa3177644112b719816b519ec70f2  runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/head.safetensors
+d843fb298102997a57d16a9d0924b5dfbb7058f30a5dc8045d39a3a9f8d72847  runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/model.safetensors

init/dflash100k_epoch6/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
+}

init/dflash100k_epoch6/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

init/dflash100k_epoch6/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

init/dflash100k_epoch6/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/dspark_a/export/epoch_1/dspark_config.json

@@ -0,0 +1,11 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "dflash",
+    "path": "init/dflash100k_epoch6",
+    "mode": null
+  },
+  "epoch": 1,
+  "step": 2859,
+  "label": "epoch_1"
+}

runs/dspark_joint_6p/export/epoch_6/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/export/epoch_6/dspark_config.json

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

runs/dspark_joint_6p/export/epoch_6/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.0",
+  "epoch": 6,
+  "step": 17156,
+  "label": "epoch_6",
+  "objective": "dflash",
+  "kd": true,
+  "tv": 0.0,
+  "head_rank": 256
+}

runs/fmlm2_v3/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/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/export/epoch_1/dspark_config.json

@@ -0,0 +1,16 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3/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/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/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/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/export/epoch_6/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/export/epoch_6/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/export/epoch_6/dspark_config.json

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

runs/fmlm2_v3_6ep/export/epoch_6/flow_config.json

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

runs/fmlm2_v3_6ep/export/epoch_6/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/export/epoch_6/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/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/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/export/epoch_1/dspark_config.json

@@ -0,0 +1,18 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_bv/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
+}

runs/fmlm2_v3_6ep_bv/export/epoch_1/flow_config.json

@@ -0,0 +1,12 @@
+{
+  "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
+}

runs/fmlm2_v3_6ep_bv/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/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_emb/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_emb/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_emb/export/epoch_1/dspark_config.json

@@ -0,0 +1,21 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_bv_emb/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": false,
+  "emb_mse": 1.0,
+  "fix_embed_norm": true
+}

runs/fmlm2_v3_6ep_bv_emb/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": false,
+  "emb_mse": 1.0,
+  "fix_embed_norm": true
+}

runs/fmlm2_v3_6ep_bv_emb/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_emb/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/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_sc/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_sc/export/epoch_1/dspark_config.json

@@ -0,0 +1,27 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_bv_lt_sc/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,
+  "self_cond": true,
+  "two_step_frac": 0.5,
+  "two_step_u": [
+    0.2,
+    0.6
+  ]
+}

runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1/flow_config.json

@@ -0,0 +1,21 @@
+{
+  "learn_embed": true,
+  "self_cond": 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,
+  "two_step_frac": 0.5,
+  "two_step_u": [
+    0.2,
+    0.6
+  ]
+}

runs/fmlm2_v3_6ep_bv_lt_sc/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_sc/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.1495, 0.1882, 0.2257, 0.2633, 0.302, 0.34, 0.3789, 0.4187, 0.4597, 0.5018, 0.5463, 0.5937, 0.6456, 0.7056, 0.7812, 0.8759, 1.0], [0.0, 0.0361, 0.0727, 0.1095, 0.1465, 0.1838, 0.221, 0.2587, 0.2964, 0.3341, 0.3719, 0.4109, 0.4514, 0.4923, 0.5354, 0.5813, 0.6319, 0.6911, 0.765, 0.8664, 1.0], [0.0, 0.0357, 0.0716, 0.1078, 0.1441, 0.1807, 0.2172, 0.2541, 0.2911, 0.3287, 0.3663, 0.4048, 0.4446, 0.4848, 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.1794, 0.2154, 0.2518, 0.2886, 0.3261, 0.3632, 0.4011, 0.4404, 0.4804, 0.5227, 0.5678, 0.6176, 0.675, 0.7473, 0.8522, 1.0], [0.0, 0.0354, 0.0709, 0.1064, 0.1421, 0.1787, 0.2151, 0.2513, 0.2878, 0.3249, 0.3619, 0.3996, 0.4384, 0.4778, 0.5196, 0.5638, 0.613, 0.6692, 0.74, 0.846, 1.0], [0.0, 0.0353, 0.0709, 0.1063, 0.142, 0.178, 0.2141, 0.2501, 0.2862, 0.3231, 0.3603, 0.3979, 0.4365, 0.4755, 0.5171, 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.285, 0.322, 0.3592, 0.397, 0.4354, 0.4744, 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.1768, 0.2126, 0.2482, 0.2842, 0.3208, 0.3581, 0.3957, 0.4344, 0.4729, 0.5142, 0.5577, 0.6055, 0.6596, 0.7293, 0.8359, 1.0], [0.0, 0.0353, 0.0705, 0.106, 0.1412, 0.1768, 0.2124, 0.2482, 0.2842, 0.3209, 0.3579, 0.3954, 0.4337, 0.4725, 0.5129, 0.5564, 0.604, 0.6577, 0.7264, 0.8335, 1.0], [0.0, 0.0353, 0.0707, 0.1064, 0.142, 0.1774, 0.213, 0.2491, 0.285, 0.3219, 0.3586, 0.396, 0.4343, 0.4729, 0.5133, 0.5565, 0.6037, 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.3977, 0.4359, 0.4745, 0.5154, 0.559, 0.6066, 0.6602, 0.7286, 0.8344, 1.0], [0.0, 0.0358, 0.0716, 0.1074, 0.1437, 0.18, 0.2164, 0.2529, 0.2897, 0.327, 0.3645, 0.4025, 0.4412, 0.4804, 0.5217, 0.5654, 0.6125, 0.6669, 0.7343, 0.838, 1.0], [0.0, 0.0372, 0.0744, 0.1116, 0.1488, 0.1861, 0.2237, 0.2614, 0.2991, 0.3375, 0.376, 0.4151, 0.4545, 0.4952, 0.5373, 0.5823, 0.6306, 0.687, 0.7579, 0.8576, 1.0]], "loss_table": [[0.56119, 0.51531, 0.47068, 0.55263, 0.50473, 0.53758, 0.55538, 0.51544, 0.54493, 0.46986, 0.48674, 0.50415, 0.54736, 0.51823, 0.50527, 0.49593, 0.46888, 0.50879, 0.51539, 0.52492, 0.58335, 0.49047, 0.48557, 0.48509, 0.45084, 0.48411, 0.5541, 0.49377, 0.49204, 0.4799, 0.49094, 0.46584, 0.49503, 0.46615, 0.56365, 0.51276, 0.49056, 0.49517, 0.55668, 0.49117, 0.52155, 0.52127, 0.4758, 0.54568, 0.47026, 0.4995, 0.50964, 0.52152, 0.4603, 0.50824, 0.51198, 0.50958, 0.48894, 0.5093, 0.49189, 0.45772, 0.50702, 0.50511, 0.46236, 0.48849, 0.49965, 0.49445, 0.5067, 0.53328, 0.50373, 0.50991, 0.46145, 0.52488, 0.47856, 0.548, 0.48685, 0.46464, 0.4842, 0.49155, 0.49151, 0.50586, 0.49445, 0.46619, 0.5133, 0.48231, 0.48286, 0.46749, 0.48616, 0.47015, 0.47883, 0.48777, 0.44945, 0.48648, 0.46142, 0.43925, 0.43364, 0.48663, 0.46961, 0.48539, 0.45462, 0.49058, 0.45492, 0.48712, 0.48816, 0.44201, 0.42974, 0.42977, 0.43304, 0.43996, 0.47087, 0.47044, 0.43195, 0.46672, 0.44405, 0.44013, 0.44522, 0.44462, 0.42168, 0.40946, 0.44051, 0.46796, 0.43323, 0.41879, 0.41994, 0.4235, 0.37343, 0.44319, 0.41676, 0.38602, 0.3976, 0.37023, 0.39292, 0.41876, 0.41587, 0.41189, 0.39201, 0.39052, 0.3472, 0.39928, 0.38014, 0.40125, 0.34738, 0.36489, 0.40939, 0.35774, 0.38595, 0.35198, 0.379, 0.3833, 0.35264, 0.35548, 0.32757, 0.34672, 0.31019, 0.33501, 0.3197, 0.32606, 0.30812, 0.33331, 0.30325, 0.30957, 0.32851, 0.32894, 0.29006, 0.27152, 0.26891, 0.24516, 0.32611, 0.29991, 0.30365, 0.28786, 0.27318, 0.31758, 0.29727, 0.27777, 0.25194, 0.28947, 0.27808, 0.33391, 0.28067, 0.27002, 0.33239, 0.25961, 0.27999, 0.31037, 0.31225, 0.27331, 0.25288, 0.27138, 0.32576, 0.27794, 0.21954, 0.28632, 0.27295, 0.24758, 0.25139, 0.20987, 0.23597, 0.2439, 0.20726, 0.21974, 0.21536, 0.17396, 0.20115, 0.13929], [1.02749, 0.99413, 0.88382, 1.03209, 0.99497, 0.94723, 1.01132, 0.93396, 0.97084, 0.95737, 0.98852, 0.92746, 0.96026, 0.97138, 0.92231, 0.93434, 0.93014, 0.93832, 0.97842, 0.99088, 1.04702, 0.98275, 0.91353, 0.96961, 0.99091, 0.97774, 1.07641, 0.9932, 0.92511, 0.90531, 0.89396, 0.91029, 0.95653, 0.87818, 0.97446, 0.96612, 0.94413, 0.95912, 1.00233, 0.99442, 1.00147, 0.99095, 0.94481, 0.97261, 0.90285, 0.9661, 0.94199, 0.92664, 0.92623, 0.97072, 0.94692, 0.95673, 0.95567, 0.9386, 0.94502, 0.9046, 0.93727, 0.92178, 0.89858, 0.89356, 0.95642, 0.88713, 0.94568, 0.93728, 0.95144, 0.95047, 0.91427, 0.98473, 0.96425, 0.9128, 0.97316, 0.89267, 0.90549, 0.9272, 0.88156, 0.94902, 0.87252, 0.88983, 0.92471, 0.96908, 0.91291, 0.87011, 0.88604, 0.89577, 0.91116, 0.8862, 0.88693, 0.92359, 0.8997, 0.92073, 0.83026, 0.93411, 0.89804, 0.90489, 0.86841, 0.85977, 0.8791, 0.82603, 0.87304, 0.82605, 0.82605, 0.81627, 0.84251, 0.82554, 0.8649, 0.88004, 0.7939, 0.75345, 0.78857, 0.88646, 0.82468, 0.84325, 0.78393, 0.78354, 0.78793, 0.77664, 0.75347, 0.76499, 0.74053, 0.73395, 0.67892, 0.74884, 0.71643, 0.69895, 0.74791, 0.72495, 0.66776, 0.72315, 0.72914, 0.7146, 0.70207, 0.71898, 0.62265, 0.65471, 0.67206, 0.68822, 0.57103, 0.64467, 0.6264, 0.60926, 0.61045, 0.58642, 0.61637, 0.59416, 0.62634, 0.57853, 0.58482, 0.55213, 0.50854, 0.53103, 0.538, 0.52987, 0.53473, 0.4947, 0.48096, 0.49089, 0.48671, 0.46847, 0.4525, 0.47561, 0.42892, 0.4265, 0.46792, 0.45673, 0.50162, 0.48511, 0.43698, 0.43922, 0.42574, 0.43303, 0.41195, 0.45342, 0.38385, 0.44989, 0.42327, 0.4459, 0.44132, 0.36121, 0.38419, 0.39766, 0.47356, 0.38195, 0.40331, 0.40416, 0.40434, 0.39764, 0.36778, 0.37744, 0.38921, 0.33838, 0.36874, 0.29224, 0.34498, 0.3626, 0.32648, 0.27729, 0.31409, 0.24373, 0.25922, 0.17659], [1.39225, 1.38317, 1.259, 1.42016, 1.41192, 1.40226, 1.41422, 1.30517, 1.44866, 1.39469, 1.3499, 1.30747, 1.37125, 1.37296, 1.3169, 1.29564, 1.33354, 1.3285, 1.32276, 1.38465, 1.38373, 1.31224, 1.28707, 1.34263, 1.35271, 1.34375, 1.44008, 1.35379, 1.28976, 1.31489, 1.31449, 1.24361, 1.32765, 1.2633, 1.36325, 1.32049, 1.32013, 1.32323, 1.40516, 1.39584, 1.36534, 1.37658, 1.26454, 1.35305, 1.3233, 1.36718, 1.34341, 1.358, 1.29889, 1.35447, 1.30504, 1.34761, 1.39315, 1.28927, 1.34962, 1.23793, 1.34887, 1.31487, 1.28978, 1.28261, 1.41183, 1.2747, 1.34607, 1.32884, 1.29165, 1.30003, 1.33815, 1.39451, 1.32353, 1.27717, 1.34388, 1.27167, 1.33183, 1.31325, 1.26142, 1.30428, 1.27991, 1.27219, 1.33913, 1.29425, 1.23654, 1.21423, 1.22843, 1.28994, 1.31366, 1.22595, 1.19208, 1.25696, 1.21321, 1.27877, 1.17356, 1.21598, 1.21464, 1.29567, 1.19445, 1.14993, 1.2286, 1.23356, 1.19179, 1.13719, 1.17041, 1.13956, 1.18262, 1.08023, 1.21006, 1.17859, 1.08815, 1.07086, 1.13325, 1.16772, 1.09236, 1.07675, 1.11288, 1.01194, 1.04697, 1.10242, 1.05206, 1.03773, 0.97697, 1.03891, 0.97335, 1.00129, 0.98252, 0.95349, 0.98808, 1.00709, 0.92552, 0.94782, 0.94673, 0.90684, 0.90041, 0.89345, 0.8623, 0.84985, 0.85403, 0.85012, 0.84893, 0.79352, 0.79724, 0.78982, 0.81118, 0.71225, 0.77821, 0.78664, 0.70613, 0.74897, 0.71045, 0.69549, 0.60322, 0.66831, 0.57841, 0.66604, 0.67558, 0.57244, 0.60045, 0.57903, 0.55269, 0.56892, 0.55265, 0.56554, 0.55663, 0.50714, 0.57508, 0.54941, 0.58577, 0.53989, 0.51679, 0.53169, 0.57806, 0.50478, 0.49822, 0.51806, 0.49145, 0.49687, 0.47345, 0.54451, 0.48346, 0.48475, 0.4439, 0.49205, 0.49121, 0.4339, 0.45146, 0.45241, 0.43104, 0.45456, 0.417, 0.47354, 0.43923, 0.39325, 0.38542, 0.38227, 0.42154, 0.3882, 0.35145, 0.32254, 0.32812, 0.26989, 0.30771, 0.18003], [1.73336, 1.77419, 1.60824, 1.79866, 1.76127, 1.64596, 1.69885, 1.67246, 1.74121, 1.72757, 1.65966, 1.65464, 1.71458, 1.67932, 1.61068, 1.68018, 1.65019, 1.6794, 1.64367, 1.68983, 1.78302, 1.67707, 1.60772, 1.68308, 1.69781, 1.6229, 1.78054, 1.67056, 1.65926, 1.66433, 1.66169, 1.68702, 1.6216, 1.62401, 1.71475, 1.66125, 1.6086, 1.65482, 1.71859, 1.63234, 1.71368, 1.6311, 1.63464, 1.66508, 1.63718, 1.72618, 1.626, 1.63613, 1.67082, 1.63058, 1.58278, 1.61928, 1.66687, 1.59678, 1.69049, 1.62969, 1.66296, 1.64842, 1.63355, 1.53913, 1.68523, 1.62311, 1.69355, 1.69182, 1.6038, 1.63667, 1.63326, 1.68042, 1.65832, 1.63011, 1.66281, 1.61329, 1.62474, 1.60296, 1.53943, 1.61974, 1.52583, 1.61659, 1.65606, 1.6645, 1.60731, 1.45777, 1.54233, 1.57651, 1.61199, 1.4974, 1.44537, 1.56346, 1.48815, 1.5203, 1.49543, 1.51448, 1.52068, 1.56343, 1.46172, 1.47023, 1.52557, 1.52323, 1.50945, 1.4577, 1.4619, 1.41793, 1.4171, 1.3938, 1.45581, 1.45945, 1.38454, 1.38295, 1.43199, 1.39199, 1.3527, 1.32446, 1.32184, 1.33474, 1.29355, 1.32083, 1.29875, 1.25605, 1.24232, 1.25121, 1.23369, 1.23377, 1.23937, 1.12635, 1.21699, 1.1762, 1.14775, 1.20808, 1.1321, 1.06619, 1.0086, 1.03166, 0.9907, 0.99977, 1.01315, 1.0202, 0.98302, 0.98164, 0.99903, 0.93179, 0.91494, 0.94235, 0.9193, 0.913, 0.80017, 0.86091, 0.78733, 0.81559, 0.70665, 0.74236, 0.73043, 0.74703, 0.78625, 0.60974, 0.68023, 0.68561, 0.64965, 0.70378, 0.64538, 0.64103, 0.62647, 0.56615, 0.66356, 0.59004, 0.60676, 0.60025, 0.57239, 0.54369, 0.64335, 0.55584, 0.54588, 0.574, 0.49618, 0.62571, 0.56887, 0.54385, 0.51874, 0.52379, 0.49667, 0.51213, 0.5629, 0.49254, 0.50361, 0.54188, 0.51397, 0.51668, 0.48534, 0.5363, 0.50521, 0.42926, 0.45612, 0.41549, 0.46229, 0.44711, 0.37801, 0.38786, 0.31866, 0.27806, 0.32821, 0.20207], [1.97221, 2.07917, 1.90568, 2.02115, 2.07632, 1.92022, 1.96113, 1.98868, 2.00031, 2.0309, 2.03579, 1.9555, 1.97624, 1.9189, 1.89861, 1.99217, 1.9372, 1.96982, 1.95129, 1.98334, 2.00439, 1.91523, 1.91812, 1.96346, 1.92367, 1.94087, 2.0503, 1.98548, 1.96518, 1.95613, 1.91138, 1.91212, 1.94548, 1.93767, 1.97481, 1.94195, 1.97398, 1.92282, 1.95767, 1.90952, 2.06171, 1.92447, 1.89943, 1.93595, 1.98535, 1.97513, 1.93567, 1.87587, 1.9654, 1.90174, 1.8447, 1.89628, 1.97616, 1.93423, 1.91549, 1.87239, 1.9654, 1.92297, 1.82736, 1.83968, 1.9041, 1.84714, 1.94757, 1.90818, 1.91484, 1.91857, 1.85837, 1.96636, 1.93312, 1.91335, 1.88962, 1.84149, 1.87979, 1.83274, 1.79627, 1.88203, 1.85131, 1.84305, 1.86853, 1.88015, 1.80809, 1.70431, 1.84597, 1.84894, 1.87968, 1.76917, 1.71911, 1.81175, 1.69622, 1.76861, 1.80406, 1.76892, 1.75209, 1.82857, 1.73366, 1.76684, 1.72315, 1.71924, 1.73843, 1.62969, 1.60016, 1.66262, 1.73882, 1.61215, 1.68327, 1.72866, 1.63762, 1.61038, 1.62611, 1.64484, 1.60555, 1.54617, 1.46811, 1.48549, 1.48353, 1.48055, 1.44342, 1.41558, 1.45233, 1.38254, 1.38054, 1.42312, 1.38414, 1.35321, 1.36887, 1.32684, 1.26467, 1.33325, 1.25185, 1.25192, 1.19844, 1.22877, 1.10518, 1.16137, 1.13951, 1.15497, 1.06842, 1.04362, 1.14343, 1.06507, 1.06031, 0.99818, 0.98519, 0.97402, 0.88654, 0.93542, 0.89322, 0.93632, 0.7976, 0.80172, 0.77535, 0.86737, 0.83405, 0.76748, 0.7279, 0.76508, 0.67555, 0.71294, 0.7013, 0.72661, 0.6379, 0.62166, 0.6707, 0.64705, 0.68149, 0.64289, 0.63288, 0.59493, 0.61144, 0.60749, 0.60304, 0.62473, 0.56459, 0.68442, 0.61646, 0.6154, 0.59574, 0.5561, 0.50343, 0.54565, 0.63653, 0.52606, 0.50546, 0.5883, 0.50033, 0.5499, 0.48586, 0.57167, 0.52681, 0.49375, 0.50394, 0.4612, 0.45598, 0.48885, 0.41651, 0.42239, 0.3833, 0.34197, 0.31808, 0.19837], [2.29864, 2.31986, 2.19206, 2.26343, 2.26304, 2.13567, 2.206, 2.18939, 2.3218, 2.27737, 2.2693, 2.18573, 2.22123, 2.20388, 2.17335, 2.25911, 2.21474, 2.21648, 2.21724, 2.22264, 2.28633, 2.17622, 2.13382, 2.2178, 2.23148, 2.16321, 2.30956, 2.18309, 2.13981, 2.16945, 2.18453, 2.20589, 2.17771, 2.1729, 2.20881, 2.156, 2.20508, 2.27293, 2.22182, 2.22436, 2.31325, 2.11194, 2.1627, 2.21308, 2.17264, 2.26953, 2.20397, 2.10787, 2.16064, 2.17973, 2.09793, 2.22387, 2.15415, 2.12105, 2.16601, 2.20489, 2.15336, 2.14133, 2.13711, 2.04017, 2.20339, 2.16659, 2.13235, 2.04351, 2.09603, 2.16553, 2.10132, 2.1568, 2.19552, 2.15884, 2.14635, 2.08753, 2.13756, 2.09417, 2.0596, 2.17207, 2.1247, 2.10665, 2.11107, 2.01838, 2.09025, 1.94481, 2.03092, 2.08718, 2.06141, 1.96079, 1.99164, 2.05401, 1.97734, 2.05247, 1.96112, 1.97936, 1.9673, 2.04928, 1.94787, 1.95647, 1.91198, 1.94461, 1.87493, 1.88306, 1.8544, 1.82989, 1.88718, 1.86005, 1.85855, 1.90381, 1.7611, 1.80335, 1.78484, 1.76898, 1.76685, 1.68879, 1.64261, 1.69636, 1.64829, 1.70189, 1.59878, 1.56183, 1.60594, 1.54401, 1.57093, 1.62956, 1.51748, 1.48406, 1.5543, 1.52539, 1.3982, 1.45273, 1.44616, 1.31013, 1.36118, 1.30856, 1.28543, 1.24902, 1.30524, 1.25409, 1.24376, 1.20512, 1.29018, 1.17708, 1.20053, 1.14759, 1.08994, 1.05535, 0.99108, 1.01332, 0.96811, 0.99651, 0.85285, 0.92035, 0.84765, 0.89276, 0.91019, 0.82882, 0.77694, 0.84214, 0.77587, 0.79632, 0.72264, 0.74931, 0.73597, 0.66524, 0.74802, 0.71168, 0.69346, 0.7001, 0.66115, 0.71998, 0.70144, 0.63567, 0.65616, 0.65721, 0.65238, 0.68343, 0.66015, 0.66174, 0.63388, 0.56055, 0.56381, 0.58597, 0.61567, 0.54038, 0.63404, 0.58302, 0.55856, 0.57109, 0.54781, 0.62105, 0.50719, 0.53284, 0.54957, 0.49386, 0.48341, 0.48801, 0.42841, 0.41286, 0.41496, 0.31628, 0.34873, 0.23691], [2.47333, 2.51488, 2.42237, 2.5448, 2.56621, 2.36867, 2.39056, 2.44577, 2.54707, 2.5176, 2.49636, 2.3948, 2.41508, 2.41796, 2.39114, 2.46022, 2.40138, 2.50595, 2.42676, 2.47979, 2.53134, 2.44331, 2.42702, 2.49199, 2.46328, 2.39795, 2.56347, 2.35878, 2.3795, 2.40288, 2.42503, 2.31462, 2.41351, 2.3657, 2.44835, 2.29277, 2.41108, 2.42454, 2.45888, 2.40708, 2.44148, 2.32944, 2.37133, 2.40263, 2.40903, 2.48153, 2.4461, 2.3254, 2.395, 2.42041, 2.37204, 2.393, 2.37046, 2.44926, 2.35795, 2.41504, 2.41119, 2.34475, 2.36877, 2.3048, 2.46128, 2.31866, 2.37045, 2.35793, 2.28937, 2.35244, 2.3507, 2.37463, 2.36227, 2.42287, 2.35051, 2.30427, 2.43467, 2.32974, 2.31664, 2.39052, 2.25348, 2.33939, 2.31824, 2.17441, 2.32559, 2.19568, 2.2527, 2.29048, 2.25024, 2.16483, 2.22973, 2.29935, 2.24021, 2.23555, 2.19937, 2.18485, 2.20256, 2.27071, 2.217, 2.15629, 2.19164, 2.11751, 2.0785, 2.10008, 2.05907, 2.01753, 2.04312, 2.07307, 2.03702, 2.10669, 1.99405, 1.98419, 1.965, 1.94109, 1.95432, 1.8732, 1.89056, 1.8544, 1.78689, 1.84835, 1.77666, 1.72712, 1.74414, 1.75412, 1.74767, 1.73925, 1.71992, 1.68684, 1.70403, 1.62638, 1.60408, 1.58292, 1.55805, 1.48876, 1.48207, 1.47697, 1.39723, 1.38635, 1.43562, 1.37171, 1.37683, 1.28543, 1.34584, 1.32906, 1.29851, 1.21462, 1.1346, 1.15748, 1.069, 1.1539, 1.05774, 1.01477, 0.94245, 1.02192, 0.96232, 0.88324, 0.91676, 0.92103, 0.88779, 0.9118, 0.85673, 0.80127, 0.7523, 0.78998, 0.77927, 0.79747, 0.82085, 0.73699, 0.73449, 0.68428, 0.70477, 0.71046, 0.70575, 0.65844, 0.70333, 0.67813, 0.65796, 0.68015, 0.69938, 0.68837, 0.69379, 0.5923, 0.61131, 0.66318, 0.66077, 0.56744, 0.64353, 0.58394, 0.63154, 0.55642, 0.57834, 0.61178, 0.51573, 0.58112, 0.59925, 0.55079, 0.53575, 0.50488, 0.44627, 0.42749, 0.37085, 0.314, 0.36586, 0.22905], [2.62205, 2.68261, 2.63955, 2.72411, 2.72475, 2.60519, 2.54346, 2.60512, 2.69736, 2.66563, 2.64618, 2.60359, 2.6322, 2.53123, 2.62264, 2.63242, 2.66047, 2.65697, 2.61738, 2.65722, 2.6806, 2.59856, 2.5891, 2.69167, 2.60203, 2.60093, 2.67819, 2.56703, 2.59722, 2.54794, 2.69251, 2.55566, 2.6761, 2.57159, 2.58541, 2.53684, 2.57148, 2.6122, 2.66177, 2.59353, 2.63736, 2.55536, 2.56972, 2.58487, 2.627, 2.66865, 2.56814, 2.57953, 2.60147, 2.56275, 2.62925, 2.59807, 2.66051, 2.58086, 2.57474, 2.55698, 2.53062, 2.56878, 2.5725, 2.49725, 2.60972, 2.45577, 2.52391, 2.49009, 2.53874, 2.44354, 2.51733, 2.57178, 2.56468, 2.47723, 2.52952, 2.51312, 2.59076, 2.51852, 2.45307, 2.50207, 2.49185, 2.46059, 2.47644, 2.36812, 2.48404, 2.41047, 2.43098, 2.51422, 2.42128, 2.32147, 2.39967, 2.46021, 2.3872, 2.41141, 2.46089, 2.38211, 2.37523, 2.38147, 2.33821, 2.31553, 2.26218, 2.26895, 2.26893, 2.27201, 2.1643, 2.18056, 2.25279, 2.24574, 2.22636, 2.20239, 2.12668, 2.06215, 2.10113, 2.10345, 2.03009, 2.0602, 2.05043, 2.00485, 2.01332, 2.01087, 1.94798, 1.87014, 1.95045, 1.87725, 1.86983, 1.85306, 1.78213, 1.79034, 1.78984, 1.73972, 1.73085, 1.6799, 1.64227, 1.64566, 1.53967, 1.56849, 1.5388, 1.44541, 1.51099, 1.49955, 1.47118, 1.4126, 1.44203, 1.29646, 1.35121, 1.31219, 1.25784, 1.18876, 1.12293, 1.1905, 1.11389, 1.1135, 0.98341, 1.06098, 0.97703, 1.01563, 0.94241, 0.97853, 0.91635, 0.95449, 0.92017, 0.88462, 0.83587, 0.8082, 0.83566, 0.77981, 0.7401, 0.7916, 0.79563, 0.75626, 0.6949, 0.75049, 0.79987, 0.66892, 0.68738, 0.72708, 0.68248, 0.74506, 0.70242, 0.66755, 0.6883, 0.65767, 0.64777, 0.71856, 0.65117, 0.62922, 0.66495, 0.60942, 0.62315, 0.60845, 0.58632, 0.62157, 0.61388, 0.63025, 0.61076, 0.51123, 0.50503, 0.49928, 0.45188, 0.41741, 0.38567, 0.38577, 0.37822, 0.23598], [2.82344, 2.9216, 2.82567, 2.9185, 2.8668, 2.82172, 2.80337, 2.84855, 2.93015, 2.86777, 2.83169, 2.81565, 2.84367, 2.80295, 2.83552, 2.89496, 2.88721, 2.86948, 2.78408, 2.89721, 2.87914, 2.8219, 2.74766, 2.88924, 2.84594, 2.79136, 2.93326, 2.81001, 2.83399, 2.8473, 2.89016, 2.8494, 2.78536, 2.84561, 2.84191, 2.76336, 2.84052, 2.75964, 2.88929, 2.82151, 2.87462, 2.76633, 2.83614, 2.80344, 2.83094, 2.81656, 2.75834, 2.79117, 2.82361, 2.75021, 2.8207, 2.83545, 2.80736, 2.74353, 2.77025, 2.77804, 2.70858, 2.83526, 2.73463, 2.70154, 2.82684, 2.63796, 2.77435, 2.61359, 2.63414, 2.65439, 2.71291, 2.78398, 2.77093, 2.70199, 2.69325, 2.67398, 2.69069, 2.71498, 2.55862, 2.72069, 2.66006, 2.70388, 2.63676, 2.67321, 2.60915, 2.55909, 2.63424, 2.76196, 2.63883, 2.51682, 2.63571, 2.63293, 2.50125, 2.59229, 2.64621, 2.61021, 2.58053, 2.5579, 2.54189, 2.54962, 2.47099, 2.46146, 2.38491, 2.47636, 2.34059, 2.38418, 2.42997, 2.37754, 2.4591, 2.3463, 2.34566, 2.22514, 2.22215, 2.21464, 2.19104, 2.1548, 2.1705, 2.19715, 2.20621, 2.11292, 2.14079, 2.04111, 2.07476, 2.09408, 2.0308, 2.04145, 1.98786, 1.90897, 2.00582, 1.79251, 1.80064, 1.81576, 1.75642, 1.81174, 1.71489, 1.6901, 1.60095, 1.54992, 1.66162, 1.54673, 1.61472, 1.55287, 1.55506, 1.4477, 1.41457, 1.45207, 1.29715, 1.3363, 1.27115, 1.24869, 1.18074, 1.20776, 1.14399, 1.12415, 0.99437, 1.11212, 0.98098, 1.06665, 0.97996, 1.0285, 0.95551, 0.91113, 0.88126, 0.8924, 0.82791, 0.82848, 0.76956, 0.79638, 0.81382, 0.81329, 0.72429, 0.78796, 0.75236, 0.75888, 0.7966, 0.77486, 0.73098, 0.70167, 0.67339, 0.69446, 0.78801, 0.75964, 0.67553, 0.74833, 0.75079, 0.67211, 0.68497, 0.63682, 0.67238, 0.66608, 0.65935, 0.65694, 0.63482, 0.62239, 0.55843, 0.5684, 0.47602, 0.53637, 0.46053, 0.48523, 0.4288, 0.37018, 0.39583, 0.26101], [2.9611, 3.08287, 2.98293, 3.10909, 3.07733, 2.95904, 2.93457, 3.06878, 3.08055, 2.99591, 3.03069, 2.95057, 3.10112, 2.99451, 2.93675, 3.02385, 2.98889, 3.03487, 2.96921, 3.07753, 2.98444, 2.96058, 2.9358, 3.09158, 3.05055, 2.8736, 3.05993, 3.03848, 3.07255, 3.0436, 3.06457, 2.97577, 2.9882, 3.02282, 2.97645, 2.94179, 2.99389, 2.95263, 2.958, 3.015, 3.04334, 2.94571, 2.92269, 3.00417, 3.00636, 3.00554, 3.02836, 2.94825, 3.03079, 2.94749, 3.03124, 2.99483, 2.98966, 2.98124, 2.85542, 2.98605, 2.9061, 2.94924, 2.84575, 2.92661, 3.0527, 2.88668, 2.83454, 2.88341, 2.84101, 2.81644, 2.86664, 2.88382, 2.9101, 2.86339, 2.8704, 2.81753, 2.90042, 2.83281, 2.76073, 2.84973, 2.81464, 2.83387, 2.84556, 2.72554, 2.73619, 2.68028, 2.80301, 2.84466, 2.78999, 2.71287, 2.7426, 2.73018, 2.70683, 2.79533, 2.78942, 2.76309, 2.7558, 2.84492, 2.73003, 2.6073, 2.64493, 2.57321, 2.62395, 2.59943, 2.5211, 2.53286, 2.49494, 2.51124, 2.51437, 2.48926, 2.53404, 2.44253, 2.44828, 2.31764, 2.3772, 2.31368, 2.36681, 2.3013, 2.22028, 2.21924, 2.2187, 2.19089, 2.21203, 2.16653, 2.1384, 2.20992, 2.04681, 2.02272, 2.11313, 1.9776, 1.9147, 1.99929, 1.84474, 1.95383, 1.76363, 1.84214, 1.67808, 1.63827, 1.72072, 1.68475, 1.629, 1.60254, 1.54192, 1.48769, 1.57464, 1.46134, 1.37563, 1.31073, 1.36368, 1.28508, 1.2197, 1.29985, 1.20088, 1.15225, 1.0602, 1.12173, 1.07387, 1.11061, 1.04738, 0.99588, 1.02365, 0.97259, 0.85406, 0.90654, 0.92332, 0.89467, 0.8705, 0.87664, 0.79093, 0.81433, 0.7722, 0.8891, 0.84641, 0.78665, 0.79596, 0.79116, 0.78928, 0.6904, 0.7991, 0.76042, 0.74613, 0.74402, 0.70618, 0.77327, 0.75131, 0.64108, 0.71002, 0.65924, 0.6029, 0.70776, 0.61978, 0.72274, 0.60844, 0.57189, 0.63841, 0.56002, 0.58267, 0.53578, 0.49985, 0.4707, 0.42026, 0.38197, 0.43416, 0.23554], [3.19343, 3.20239, 3.13184, 3.29134, 3.21132, 3.11995, 3.09529, 3.24823, 3.29723, 3.14682, 3.14646, 3.1329, 3.24856, 3.10378, 3.07151, 3.12282, 3.19309, 3.19693, 3.10269, 3.24792, 3.15043, 3.20357, 3.08494, 3.26689, 3.22521, 3.09461, 3.25712, 3.19606, 3.11945, 3.15419, 3.18382, 3.16207, 3.13179, 3.19103, 3.09808, 3.11588, 3.18865, 3.13781, 3.13282, 3.11824, 3.23687, 3.08696, 3.10788, 3.11991, 3.13781, 3.2001, 3.13992, 3.10072, 3.18447, 3.08145, 3.16941, 3.11074, 3.06105, 3.13778, 3.0721, 3.16506, 3.06357, 3.10333, 3.00613, 3.07913, 3.17032, 2.96229, 3.02222, 3.02243, 3.06246, 2.97466, 3.0716, 3.04566, 3.01666, 3.04437, 3.06654, 2.94824, 3.14086, 3.03724, 2.91487, 2.99628, 2.9893, 2.90494, 2.99482, 2.92591, 2.93631, 2.90161, 2.95293, 2.98694, 2.95252, 2.85441, 2.85244, 2.90698, 2.80754, 2.89038, 3.00048, 2.86277, 2.91839, 2.92407, 2.80664, 2.80812, 2.81496, 2.85871, 2.81034, 2.85791, 2.67021, 2.61999, 2.66776, 2.59758, 2.72355, 2.6511, 2.69812, 2.55582, 2.51354, 2.41579, 2.46481, 2.4325, 2.48665, 2.43015, 2.40232, 2.3875, 2.30892, 2.29689, 2.24484, 2.33208, 2.24647, 2.27019, 2.20752, 2.15437, 2.27889, 1.99457, 2.00791, 2.07646, 2.01973, 2.02311, 1.93199, 1.86452, 1.8723, 1.7433, 1.77017, 1.72637, 1.6836, 1.68443, 1.61228, 1.60001, 1.65086, 1.54809, 1.49522, 1.39344, 1.39878, 1.38507, 1.27633, 1.29452, 1.23953, 1.28451, 1.12874, 1.09957, 1.08936, 1.12668, 1.13036, 1.05055, 1.00614, 1.01822, 0.93145, 0.91467, 0.91562, 0.92828, 0.87074, 0.85061, 0.85112, 0.86131, 0.81349, 0.83301, 0.8567, 0.78408, 0.80496, 0.78961, 0.76554, 0.74296, 0.77431, 0.86486, 0.75224, 0.7512, 0.68724, 0.75759, 0.83855, 0.7594, 0.75463, 0.75155, 0.68391, 0.75308, 0.66528, 0.66632, 0.64424, 0.57768, 0.61612, 0.60781, 0.64547, 0.59501, 0.51737, 0.48495, 0.41408, 0.36772, 0.44028, 0.23706], [3.37848, 3.42879, 3.33091, 3.39188, 3.38804, 3.26827, 3.23759, 3.35025, 3.40336, 3.29478, 3.30456, 3.3201, 3.36179, 3.25783, 3.28224, 3.27278, 3.29207, 3.35897, 3.26968, 3.31527, 3.36831, 3.3121, 3.22301, 3.38262, 3.32892, 3.24783, 3.37825, 3.30449, 3.27904, 3.38997, 3.35976, 3.33226, 3.36358, 3.31112, 3.32567, 3.27556, 3.34927, 3.29176, 3.33904, 3.28748, 3.35276, 3.2818, 3.22126, 3.21895, 3.29947, 3.38266, 3.21947, 3.227, 3.30285, 3.27223, 3.3048, 3.27837, 3.35545, 3.22739, 3.21509, 3.38732, 3.20734, 3.21172, 3.19823, 3.15898, 3.32153, 3.14329, 3.22848, 3.14964, 3.22607, 3.1951, 3.26725, 3.23237, 3.212, 3.24434, 3.19042, 3.12762, 3.25183, 3.14324, 3.10626, 3.15565, 3.08017, 3.19381, 3.12706, 3.11159, 3.15583, 3.00394, 3.08655, 3.06999, 3.18238, 3.00297, 3.05686, 3.14318, 2.98167, 3.15618, 3.10211, 2.99226, 3.04744, 3.07893, 2.89549, 2.97301, 3.00353, 2.96183, 2.9383, 2.99557, 2.81427, 2.7625, 2.87004, 2.7608, 2.91354, 2.79844, 2.78535, 2.73082, 2.6789, 2.5577, 2.63928, 2.60217, 2.62457, 2.50563, 2.55959, 2.55472, 2.45179, 2.40545, 2.46134, 2.42751, 2.35681, 2.35609, 2.2164, 2.30853, 2.39848, 2.10081, 2.11653, 2.13155, 2.08801, 2.05349, 2.04765, 1.99422, 1.9496, 1.7998, 1.93071, 1.87551, 1.85249, 1.79733, 1.77826, 1.70695, 1.64547, 1.56944, 1.54078, 1.50708, 1.4977, 1.45235, 1.32053, 1.42295, 1.24867, 1.28797, 1.19094, 1.2521, 1.19669, 1.16552, 1.22602, 1.09105, 1.03965, 1.09217, 0.96607, 0.96412, 0.97236, 0.91114, 0.95512, 0.89627, 0.89711, 0.92361, 0.81243, 0.91208, 0.90011, 0.81552, 0.7966, 0.84239, 0.80989, 0.77283, 0.90465, 0.83769, 0.81245, 0.79307, 0.72788, 0.76584, 0.89947, 0.78561, 0.75995, 0.77691, 0.77273, 0.68918, 0.67826, 0.74644, 0.68779, 0.64606, 0.65096, 0.60094, 0.60415, 0.58055, 0.56501, 0.51147, 0.47754, 0.40428, 0.41933, 0.27882], [3.45355, 3.54564, 3.46045, 3.47482, 3.54128, 3.40295, 3.38395, 3.51731, 3.54198, 3.4677, 3.45926, 3.47986, 3.50114, 3.41294, 3.3949, 3.42764, 3.43492, 3.57866, 3.52511, 3.48281, 3.56999, 3.43716, 3.40737, 3.53623, 3.40866, 3.39351, 3.54632, 3.42026, 3.42956, 3.47911, 3.52776, 3.4627, 3.51609, 3.4934, 3.46886, 3.45405, 3.53167, 3.43643, 3.4697, 3.42576, 3.44893, 3.3991, 3.42485, 3.38318, 3.45848, 3.49479, 3.4675, 3.36715, 3.48009, 3.42149, 3.42647, 3.368, 3.48576, 3.41926, 3.31024, 3.46671, 3.39544, 3.4812, 3.34525, 3.32926, 3.43706, 3.29181, 3.38639, 3.37053, 3.31654, 3.35588, 3.31142, 3.33817, 3.3834, 3.27398, 3.28695, 3.24862, 3.36089, 3.30773, 3.28546, 3.33375, 3.37633, 3.32734, 3.28535, 3.27155, 3.31794, 3.18002, 3.29065, 3.29064, 3.24901, 3.19316, 3.19182, 3.25863, 3.25058, 3.26374, 3.19666, 3.13149, 3.1887, 3.21755, 3.09441, 3.04741, 3.07941, 3.10911, 3.07514, 3.09431, 2.96334, 2.89985, 2.99169, 2.87842, 2.94058, 2.93809, 2.92805, 2.75098, 2.86046, 2.75927, 2.71974, 2.72135, 2.79333, 2.69422, 2.65643, 2.61313, 2.62757, 2.55643, 2.53148, 2.46925, 2.47008, 2.51467, 2.39483, 2.39482, 2.48174, 2.25869, 2.32899, 2.38975, 2.17881, 2.18753, 2.19015, 2.18823, 2.00955, 1.95328, 2.00004, 2.00302, 1.93387, 1.87135, 1.87497, 1.79297, 1.77509, 1.67069, 1.63145, 1.6394, 1.56791, 1.54688, 1.53246, 1.50176, 1.42783, 1.33179, 1.28377, 1.27669, 1.2848, 1.24649, 1.27292, 1.18801, 1.10184, 1.11624, 1.06025, 1.04594, 1.02034, 0.99592, 1.02115, 0.91945, 0.92029, 0.94527, 0.93618, 0.94401, 0.8688, 0.92614, 0.86992, 0.8675, 0.85046, 0.896, 0.96437, 0.83422, 0.86613, 0.88575, 0.85782, 0.8376, 0.89819, 0.81351, 0.78291, 0.80305, 0.79095, 0.77775, 0.6833, 0.7569, 0.7446, 0.64494, 0.70293, 0.62498, 0.66985, 0.63109, 0.57789, 0.54638, 0.51493, 0.45151, 0.43952, 0.31857], [3.61767, 3.66768, 3.57593, 3.69979, 3.66217, 3.63199, 3.61033, 3.59745, 3.71156, 3.60465, 3.66011, 3.60192, 3.66626, 3.57384, 3.59725, 3.58856, 3.68127, 3.70509, 3.6534, 3.61837, 3.60968, 3.64579, 3.5151, 3.69107, 3.50227, 3.54887, 3.6568, 3.55652, 3.56703, 3.55841, 3.64193, 3.53288, 3.61845, 3.6559, 3.6339, 3.5125, 3.60767, 3.55384, 3.58367, 3.56711, 3.55331, 3.54517, 3.56861, 3.5635, 3.56746, 3.60169, 3.6058, 3.50728, 3.63786, 3.51844, 3.53564, 3.46318, 3.56557, 3.5769, 3.48261, 3.62617, 3.50404, 3.58098, 3.50374, 3.54791, 3.52408, 3.44594, 3.44886, 3.47724, 3.47529, 3.46958, 3.46426, 3.44886, 3.41027, 3.48644, 3.41059, 3.46414, 3.60672, 3.36225, 3.37614, 3.45554, 3.4215, 3.40851, 3.39349, 3.51987, 3.37851, 3.30514, 3.42004, 3.42084, 3.40329, 3.32041, 3.35407, 3.36702, 3.36984, 3.3939, 3.37188, 3.27173, 3.33237, 3.36051, 3.24738, 3.19448, 3.25411, 3.1887, 3.2024, 3.25358, 3.17007, 3.09129, 3.09577, 2.99909, 3.09046, 3.12736, 3.02587, 3.02227, 2.97045, 2.86802, 2.88487, 2.8753, 2.96061, 2.8623, 2.8252, 2.85061, 2.74172, 2.83872, 2.75648, 2.72822, 2.61215, 2.78517, 2.52749, 2.52331, 2.64812, 2.47071, 2.42974, 2.44227, 2.32075, 2.3522, 2.29646, 2.26832, 2.19471, 2.17318, 2.13813, 2.21669, 2.17031, 2.10956, 2.05189, 1.91065, 1.94063, 1.85436, 1.82479, 1.71666, 1.69216, 1.69863, 1.5901, 1.59375, 1.57871, 1.49146, 1.36535, 1.36509, 1.3421, 1.359, 1.36772, 1.30941, 1.21289, 1.21979, 1.14362, 1.16513, 1.13338, 1.17627, 1.09769, 1.03718, 1.06797, 1.07579, 1.0388, 1.09203, 1.06692, 1.03043, 0.97296, 0.97274, 0.9153, 0.95496, 0.9102, 0.90138, 0.94482, 0.94602, 0.84488, 0.9589, 1.02387, 0.86499, 0.86798, 0.84083, 0.80761, 0.91846, 0.72983, 0.85296, 0.75472, 0.68084, 0.75782, 0.68198, 0.69679, 0.68707, 0.68281, 0.54033, 0.5708, 0.46796, 0.45181, 0.30403], [3.71864, 3.75291, 3.74025, 3.77315, 3.81466, 3.78138, 3.66739, 3.72655, 3.74299, 3.74819, 3.73442, 3.70516, 3.83666, 3.72228, 3.72477, 3.6922, 3.7497, 3.87364, 3.70365, 3.71922, 3.7552, 3.78525, 3.64908, 3.79071, 3.81507, 3.67623, 3.81076, 3.73953, 3.75207, 3.69181, 3.80447, 3.734, 3.77079, 3.73707, 3.80888, 3.6117, 3.72263, 3.62904, 3.72405, 3.83132, 3.7294, 3.64268, 3.7213, 3.65721, 3.73948, 3.75624, 3.69537, 3.68082, 3.69817, 3.6603, 3.66401, 3.66773, 3.82251, 3.70773, 3.65904, 3.73125, 3.69362, 3.72113, 3.64926, 3.65679, 3.655, 3.54995, 3.68218, 3.59637, 3.65499, 3.55792, 3.69086, 3.66125, 3.60391, 3.68028, 3.6113, 3.59844, 3.6975, 3.58898, 3.5349, 3.57887, 3.52024, 3.60692, 3.5482, 3.59944, 3.57185, 3.52229, 3.60637, 3.57342, 3.50053, 3.54751, 3.56167, 3.59107, 3.45306, 3.57478, 3.48849, 3.5337, 3.47604, 3.52519, 3.42277, 3.31377, 3.3611, 3.33201, 3.45154, 3.4749, 3.33156, 3.35206, 3.28419, 3.2494, 3.31909, 3.25372, 3.23122, 3.21778, 3.1627, 3.16857, 3.02649, 3.04922, 3.12445, 3.1432, 3.06561, 3.05601, 2.93882, 2.95704, 2.94729, 2.89347, 2.88088, 3.00531, 2.863, 2.82636, 2.85597, 2.74513, 2.57129, 2.72223, 2.58495, 2.50907, 2.47334, 2.46889, 2.43868, 2.38063, 2.33232, 2.42052, 2.36982, 2.34768, 2.21901, 2.12374, 2.21332, 2.1125, 1.96428, 1.97443, 2.01231, 1.95897, 1.84386, 1.88863, 1.77166, 1.81798, 1.62885, 1.63529, 1.63138, 1.61451, 1.61945, 1.56175, 1.55023, 1.44895, 1.41075, 1.41981, 1.46513, 1.45225, 1.27924, 1.23816, 1.29849, 1.28279, 1.34228, 1.2222, 1.25968, 1.20982, 1.24478, 1.22223, 1.18534, 1.20022, 1.15334, 1.21234, 1.14324, 1.17848, 1.10185, 1.08935, 1.18534, 1.07214, 1.08967, 1.13217, 1.00414, 1.05055, 0.9873, 1.08748, 0.93732, 0.98312, 1.05548, 0.93158, 0.95163, 0.83166, 0.82053, 0.68444, 0.69067, 0.56812, 0.5515, 0.42088]]}

runs/fmlm2_v3_6ep_bv_lt_sc/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_ce/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_ce/export/epoch_1/dspark_config.json

@@ -0,0 +1,16 @@
+{
+  "rank": 256,
+  "backbone": {
+    "kind": "flow",
+    "path": "runs/fmlm2_v3_6ep_ce/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_ce/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/full_scratch/export/epoch_6/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/full_scratch/export/epoch_6/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/full_scratch/export/epoch_6/flow_config.json

@@ -0,0 +1,6 @@
+{
+  "learn_embed": true,
+  "epoch": 6,
+  "step": 17156,
+  "label": "epoch_6"
+}

runs/full_scratch/export/epoch_6/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/full_scratch/export/epoch_6/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