From d0bec6c2b73ab609130c28836e25de848432e141 Mon Sep 17 00:00:00 2001
From: Drew <262217085+Amidwestnoob@users.noreply.github.com>
Date: Thu, 30 Jul 2026 11:46:44 -0500
Subject: [PATCH] fix(draft): embed Qwen3.6 SWA metadata in Q8 converter

---
 server/README.md                       |   2 +-
 server/scripts/quantize_draft_q8.py    |  34 +++++-
 server/tests/test_quantize_draft_q8.py | 139 +++++++++++++++++++++++++
 3 files changed, 172 insertions(+), 3 deletions(-)
 create mode 100644 server/tests/test_quantize_draft_q8.py

diff --git a/server/README.md b/server/README.md
index 39ff3c3b2..1e96b697d 100644
--- a/server/README.md
+++ b/server/README.md
@@ -537,7 +537,7 @@ cmake --build build --target test_dflash -j
 
 > **Multi-GPU / distro note.** On a host with more than one AMD GPU, pin the bench to the target with `HIP_VISIBLE_DEVICES`. On distros that link PIE executables by default (e.g. Fedora's system ROCm under `/usr`), add `-DCMAKE_EXE_LINKER_FLAGS=-no-pie` to the `cmake` configure line, and point at the toolchain with `-DCMAKE_HIP_COMPILER_ROCM_ROOT=/usr -DROCM_PATH=/usr` if ROCm lives under `/usr` rather than `/opt/rocm`.
 
-**Drafter recipe for max decode:** target = Qwen3.5-27B Q4_K_M, drafter = same gen quantized to Q8_0 via `server/scripts/quantize_draft_q8.py`. Matching Q8_0 GGUF on the unsloth Qwen3.6 target needs `DFLASH27B_DRAFT_SWA=2048` for sliding-window correctness.
+**Drafter recipe for max decode:** target = Qwen3.5-27B Q4_K_M, drafter = same gen quantized to Q8_0 via `server/scripts/quantize_draft_q8.py`. For the unsloth Qwen3.6 target, pass `--qwen36-swa` when creating the Q8_0 draft so its GGUF embeds the required 2048-token, 4-of-5-layer sliding-window configuration. Older Qwen3.6 drafts without that metadata still need `DFLASH27B_DRAFT_SWA=2048` at runtime.
 
 See also: [`docs/HIP_PERF_PLAN.md`](docs/HIP_PERF_PLAN.md) (perf sweeps), [`docs/MIXED_BACKEND.md`](docs/MIXED_BACKEND.md) (mixed CUDA+HIP runs).
 
diff --git a/server/scripts/quantize_draft_q8.py b/server/scripts/quantize_draft_q8.py
index 7f8e3648c..b8f72fdbc 100644
--- a/server/scripts/quantize_draft_q8.py
+++ b/server/scripts/quantize_draft_q8.py
@@ -13,6 +13,9 @@
     python3 scripts/quantize_draft_q8.py \
         models/draft/model.safetensors \
         models/draft/draft-q8_0.gguf
+
+For a Qwen3.6 target, add ``--qwen36-swa`` so the output carries the
+draft sliding-window architecture metadata required by the server.
 """
 
 import argparse
@@ -48,6 +51,20 @@
 
 Q8_0_BLOCK_SIZE     = 32   # elements per Q8_0 block
 
+QWEN36_SWA_WINDOW   = 2048
+QWEN36_SWA_PATTERN  = (True, True, True, True, False)
+
+
+def add_qwen36_swa_metadata(writer, enabled: bool) -> None:
+    """Embed Qwen3.6 draft SWA metadata without changing Qwen3.5 defaults."""
+    if not enabled:
+        return
+    writer.add_uint32(f"{ARCH}.attention.sliding_window", QWEN36_SWA_WINDOW)
+    writer.add_array(
+        f"{ARCH}.attention.sliding_window_pattern",
+        list(QWEN36_SWA_PATTERN),
+    )
+
 
 # ──────────────────────────────────────────────────────────────────────
 # Tensor name mapping  —  DFlash safetensors -> llama.cpp GGUF
@@ -116,14 +133,24 @@ def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray:
 # Main
 # ──────────────────────────────────────────────────────────────────────
 
-def main():
+def build_arg_parser() -> argparse.ArgumentParser:
     ap = argparse.ArgumentParser(
         description="Quantize DFlash draft BF16 safetensors to Q8_0 GGUF")
     ap.add_argument("safetensors", type=Path,
                     help="Input BF16 safetensors (e.g. models/draft/model.safetensors)")
     ap.add_argument("out_gguf", type=Path,
                     help="Output Q8_0 GGUF (e.g. models/draft/draft-q8_0.gguf)")
-    args = ap.parse_args()
+    ap.add_argument(
+        "--qwen36-swa",
+        action="store_true",
+        help=("embed Qwen3.6 draft SWA metadata (window 2048; layers "
+              "0-3 sliding, layer 4 full attention)"),
+    )
+    return ap
+
+
+def main():
+    args = build_arg_parser().parse_args()
 
     if not args.safetensors.exists():
         print(f"[error] safetensors not found: {args.safetensors}", file=sys.stderr)
@@ -150,6 +177,9 @@ def main():
     writer.add_uint32(f"{ARCH}.vocab_size",              VOCAB)
     writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", RMS_EPS)
     writer.add_float32(f"{ARCH}.rope.freq_base",         ROPE_THETA)
+    add_qwen36_swa_metadata(writer, args.qwen36_swa)
+    if args.qwen36_swa:
+        print("[info] Qwen3.6 draft SWA: layers 0-3 window=2048; layer 4 full attention")
 
     # DFlash-specific hyperparameters
     writer.add_uint32(f"{ARCH}.dflash.n_target_layers", N_TARGET_LAYERS)
diff --git a/server/tests/test_quantize_draft_q8.py b/server/tests/test_quantize_draft_q8.py
new file mode 100644
index 000000000..4fd201eda
--- /dev/null
+++ b/server/tests/test_quantize_draft_q8.py
@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""Regression tests for Q8 draft conversion metadata profiles."""
+
+import importlib.util
+import json
+import struct
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+
+SERVER_DIR = Path(__file__).resolve().parents[1]
+SCRIPT = SERVER_DIR / "scripts" / "quantize_draft_q8.py"
+SPEC = importlib.util.spec_from_file_location("quantize_draft_q8", SCRIPT)
+MODULE = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(MODULE)
+
+
+class RecordingWriter:
+    def __init__(self):
+        self.calls = []
+
+    def add_uint32(self, key, value):
+        self.calls.append(("uint32", key, value))
+
+    def add_array(self, key, value):
+        self.calls.append(("array", key, value))
+
+
+class Qwen36SwaMetadataTest(unittest.TestCase):
+    @staticmethod
+    def _write_minimal_safetensors(path):
+        header = json.dumps(
+            {
+                "hidden_norm.weight": {
+                    "dtype": "F32",
+                    "shape": [1],
+                    "data_offsets": [0, 4],
+                }
+            },
+            separators=(",", ":"),
+        ).encode()
+        payload = struct.pack("<Q", len(header)) + header + struct.pack("<f", 1.0)
+        path.write_bytes(payload)
+
+    def test_profile_is_opt_in(self):
+        writer = RecordingWriter()
+        MODULE.add_qwen36_swa_metadata(writer, False)
+        self.assertEqual(writer.calls, [])
+
+    def test_profile_writes_exact_loader_keys(self):
+        writer = RecordingWriter()
+        MODULE.add_qwen36_swa_metadata(writer, True)
+        self.assertEqual(
+            writer.calls,
+            [
+                (
+                    "uint32",
+                    "qwen35-dflash-draft.attention.sliding_window",
+                    2048,
+                ),
+                (
+                    "array",
+                    "qwen35-dflash-draft.attention.sliding_window_pattern",
+                    [True, True, True, True, False],
+                ),
+            ],
+        )
+
+    def test_cli_default_preserves_qwen35_behavior(self):
+        args = MODULE.build_arg_parser().parse_args(["in.safetensors", "out.gguf"])
+        self.assertFalse(args.qwen36_swa)
+
+    def test_cli_enables_qwen36_profile(self):
+        args = MODULE.build_arg_parser().parse_args(
+            ["in.safetensors", "out.gguf", "--qwen36-swa"]
+        )
+        self.assertTrue(args.qwen36_swa)
+
+    def test_converter_cli_writes_profile_only_when_requested(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp = Path(tmp)
+            source = tmp / "minimal.safetensors"
+            self._write_minimal_safetensors(source)
+            for enabled in (False, True):
+                with self.subTest(enabled=enabled):
+                    output = tmp / f"draft-{enabled}.gguf"
+                    command = [sys.executable, str(SCRIPT), str(source), str(output)]
+                    if enabled:
+                        command.append("--qwen36-swa")
+                    subprocess.run(command, check=True, capture_output=True, text=True)
+                    reader = MODULE.gguf.GGUFReader(output)
+                    window = reader.get_field(
+                        "qwen35-dflash-draft.attention.sliding_window"
+                    )
+                    pattern = reader.get_field(
+                        "qwen35-dflash-draft.attention.sliding_window_pattern"
+                    )
+                    if enabled:
+                        self.assertEqual(window.contents(), 2048)
+                        self.assertEqual(
+                            pattern.contents(), [True, True, True, True, False]
+                        )
+                    else:
+                        self.assertIsNone(window)
+                        self.assertIsNone(pattern)
+
+    def test_gguf_round_trip_preserves_types_and_values(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "metadata.gguf"
+            writer = MODULE.gguf.GGUFWriter(path, MODULE.ARCH)
+            MODULE.add_qwen36_swa_metadata(writer, True)
+            writer.write_header_to_file()
+            writer.write_kv_data_to_file()
+            writer.write_tensors_to_file()
+            writer.close()
+
+            reader = MODULE.gguf.GGUFReader(path)
+            window = reader.get_field(
+                "qwen35-dflash-draft.attention.sliding_window"
+            )
+            pattern = reader.get_field(
+                "qwen35-dflash-draft.attention.sliding_window_pattern"
+            )
+            self.assertIsNotNone(window)
+            self.assertIsNotNone(pattern)
+            self.assertEqual(window.contents(), 2048)
+            self.assertEqual(pattern.contents(), [True, True, True, True, False])
+            self.assertEqual([kind.name for kind in window.types], ["UINT32"])
+            self.assertEqual(
+                [kind.name for kind in pattern.types], ["ARRAY", "BOOL"]
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file