Qwen3.5 MoE support

dflash/CMakeLists.txt

@@ -120,6 +120,8 @@ add_library(dflash27b STATIC
     src/kv_cache.cpp
     src/kv_quant.cpp
     src/delta_net_chunked.cpp
+    src/moe_experts.cpp
+    src/qwen35moe_target_graph.cpp
     # Laguna-XS.2 (Poolside) target arch
     src/laguna_target_loader.cpp
     src/laguna_target_graph.cpp
@@ -323,8 +325,8 @@ if(DFLASH27B_TESTS)
         target_link_libraries(test_generate PRIVATE dflash27b ggml ggml-cuda)
     endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash.cpp")
-        add_executable(test_dflash test/test_dflash.cpp)
-        target_include_directories(test_dflash PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
+        add_executable(test_dflash test/test_dflash.cpp test/test_dflash_moe.cpp)
+        target_include_directories(test_dflash PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${CMAKE_CURRENT_SOURCE_DIR}/test)
         target_link_libraries(test_dflash PRIVATE dflash27b ggml ggml-cuda)
         # test_dflash uses cudaMemcpyAsync / cudaMemcpy2DAsync directly for the
         # --fast-rollback path (per-step SSM intermediate state commit). Needs
@@ -337,4 +339,5 @@ if(DFLASH27B_TESTS)
             target_link_libraries(test_dflash PRIVATE OpenMP::OpenMP_CXX)
         endif()
     endif()
+
 endif()

dflash/scripts/convert_dflash_to_gguf.py

@@ -39,16 +39,13 @@
 import gguf
 
 # ──────────────────────────────────────────────────────────────────────
-# DFlash 27B draft architecture constants
+# DFlash draft architecture constants (defaults for 27B; overridden by
+# auto-detection from tensor shapes when possible)
 # ──────────────────────────────────────────────────────────────────────
 
 ARCH                = "qwen35-dflash-draft"
-HIDDEN              = 5120
 N_LAYER             = 5
-N_HEAD              = 32          # query heads
-N_HEAD_KV           = 8
 HEAD_DIM            = 128
-INTERMEDIATE        = 17408
 VOCAB               = 248320
 N_TARGET_LAYERS     = 5            # fc projects 5*hidden -> hidden
 ROPE_THETA          = 1_000_000.0
@@ -155,27 +152,75 @@ def main():
     n_entries = sum(1 for k in header if k != "__metadata__")
     print(f"[info]   {n_entries} tensor entries")
 
+    # Auto-detect dimensions from tensor shapes
+    fc_shape = header.get("fc.weight", {}).get("shape")
+    q_shape = header.get("layers.0.self_attn.q_proj.weight", {}).get("shape")
+    k_shape = header.get("layers.0.self_attn.k_proj.weight", {}).get("shape")
+    gate_shape = header.get("layers.0.mlp.gate_proj.weight", {}).get("shape")
+
+    if fc_shape:
+        # fc.weight: [hidden, n_target_layers * hidden] (PyTorch: [out, in])
+        HIDDEN = fc_shape[0]
+        detected_n_target = fc_shape[1] // HIDDEN
+        print(f"[info] auto-detected hidden={HIDDEN} n_target_layers={detected_n_target}")
+    else:
+        HIDDEN = 5120  # fallback for 27B
+        detected_n_target = N_TARGET_LAYERS
+        print("[warn] fc.weight not found, using default HIDDEN=5120")
+
+    if q_shape:
+        q_dim = q_shape[0]
+        N_HEAD = q_dim // HEAD_DIM
+    else:
+        N_HEAD = 32
+
+    if k_shape:
+        kv_dim = k_shape[0]
+        N_HEAD_KV = kv_dim // HEAD_DIM
+    else:
+        N_HEAD_KV = 8
+
+    if gate_shape:
+        INTERMEDIATE = gate_shape[0]
+    else:
+        INTERMEDIATE = 17408
+
+    # Count actual layers from tensor names
+    layer_indices = set()
+    for name in header:
+        if name.startswith("layers."):
+            parts = name.split(".")
+            if len(parts) >= 2:
+                try:
+                    layer_indices.add(int(parts[1]))
+                except ValueError:
+                    pass
+    if layer_indices:
+        actual_n_layer = max(layer_indices) + 1
+    else:
+        actual_n_layer = N_LAYER
+
+    print(f"[info] dimensions: hidden={HIDDEN} n_head={N_HEAD} n_head_kv={N_HEAD_KV} "
+          f"ffn={INTERMEDIATE} n_layer={actual_n_layer}")
+
     writer = gguf.GGUFWriter(args.out_gguf, ARCH)
 
     # Architecture metadata
-    writer.add_string("general.name", "Qwen3.5-27B-DFlash-Draft")
+    writer.add_string("general.name", "DFlash-Draft")
     writer.add_uint32(f"{ARCH}.context_length",          CTX_LEN)
     writer.add_uint32(f"{ARCH}.embedding_length",        HIDDEN)
-    writer.add_uint32(f"{ARCH}.block_count",             N_LAYER)
+    writer.add_uint32(f"{ARCH}.block_count",             actual_n_layer)
     writer.add_uint32(f"{ARCH}.feed_forward_length",     INTERMEDIATE)
     writer.add_uint32(f"{ARCH}.attention.head_count",    N_HEAD)
     writer.add_uint32(f"{ARCH}.attention.head_count_kv", N_HEAD_KV)
-    # llama.cpp uses key_length / value_length to override the default
-    # n_embd_head = n_embd / n_head heuristic (DFlash has n_embd=5120
-    # but head_dim=128 so n_head*head_dim=4096 != n_embd).
     writer.add_uint32(f"{ARCH}.attention.key_length",    HEAD_DIM)
     writer.add_uint32(f"{ARCH}.attention.value_length",  HEAD_DIM)
     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)
 
     # DFlash-specific hyperparameters
-    writer.add_uint32(f"{ARCH}.dflash.n_target_layers", N_TARGET_LAYERS)
+    writer.add_uint32(f"{ARCH}.dflash.n_target_layers", detected_n_target)
     writer.add_uint32(f"{ARCH}.dflash.block_size",      BLOCK_SIZE)
     writer.add_uint32(f"{ARCH}.dflash.mask_token_id",   MASK_TOKEN_ID)
 

dflash/src/dflash_graph.h

@@ -18,6 +18,9 @@ struct DraftGraphInputs {
     // hidden states. Used for DFlash integration where the draft shares the
     // target's lm_head.
     ggml_tensor * lm_head;
+    // Optional: YaRN RoPE frequency factors [head_dim/2] f32.
+    // If non-null, passed to ggml_rope_ext to correct for YaRN scaling.
+    ggml_tensor * rope_freq_factors = nullptr;
 };
 
 struct DraftGraphOutputs {