Serve Qwen3.5-122B-A10B and its fine-tune Ornith-1.5-122B

DEVELOPMENT.md

@@ -1631,10 +1631,20 @@ weight file is downloaded for either. Tensor shapes, and an MLX vision tower's
 weights, which must not be quantized (`VisionLoader.cpp`), are checked only
 when the model loads.
 `families.FAMILIES` names the draft trained for each family; repository names
-and model-card `base_model` fields play no part. An MLX target must declare
-`quantization` in `config.json` in formats it loads ([MLX
-targets](#mlx-targets)). Every start checks the configuration again. Remote
-Python code is not loaded.
+and model-card `base_model` fields play no part. Families of one model type,
+Qwen3.6-35B-A3B and Qwen3.5-122B-A10B (`qwen3_5_moe_text`, GGUF `qwen35moe`),
+are told apart by their sizes, and a fine-tune of either, such as
+Ornith-1.5-35B-A3B or Ornith-1.5-122B, is its family. A layout type
+(`Qwen3_8.hpp`, `Qwen3_6Moe.hpp`, `Qwen4Exp.hpp` in `runtime/model/`) is an
+architecture, whose code its families share, and each family is a value of it in
+one list, `model::kFamilies` (`ModelDescriptor.hpp`): another size of an
+architecture is one more value there, its draft in `families.FAMILIES`, and, for
+a GDN or attention shape no family had, its kernel instances and a row of the
+compiled shapes in `ops/GDN.cpp` or `ops/PagedAttention.cpp`. The kernel tests
+take every family's shapes from that list (`dev/tests/engine/TestFamilies.hpp`).
+An MLX target must declare `quantization` in `config.json` in formats it loads
+([MLX targets](#mlx-targets)). Every start checks the configuration again.
+Remote Python code is not loaded.
 
 ```bash
 splash serve --model mlx-community/Qwen3.6-35B-A3B-4bit
@@ -1843,11 +1853,20 @@ A draft's `config.json` states its layers, intermediate size, rotary base and
 the target layers whose outputs it reads (`dflash_config.target_layer_ids`,
 ascending, at most eight); its hidden size, vocabulary and mask token are its
 target's, and its attention (32 query and 8 KV heads of 128), selector (rank
-256, 16 candidates), block of 8 rows and 2048-token window are those the kernels
-are built for (`readDraftLayout` in `ModelDescriptor.mm`); the kernels take any
-hidden width of whole 16-channel convolution groups. Every family loads and runs
-its draft alike: a decode cycle drafts the whole block, the anchor and its seven
-proposals, and the target verifies the block's first rows, as many as the cycle
+256, 16 candidates) and 2048-token window are those the kernels are built for,
+and its block at least their 8 rows, of which a cycle drafts the first (the
+Qwen3.5-122B-A10B draft was trained to fill 16; `readDraftLayout` in
+`ModelDescriptor.mm`); the kernels take any hidden width of whole 16-channel
+convolution groups. A draft's rows after its anchor read the target's token
+table row of the mask token, unless its release ships that token's BF16
+embedding beside its weights: `mask_embedding.safetensors`, one tensor
+`mask_embedding` of the hidden size, which installation fetches with the weights
+(`DFlashDraft::addMaskRows`). A quantized table can lose the row:
+Qwen3.5-122B-A10B's is of order 1e-6, below what the fp16 block scales of Q8_0
+and the K-quants hold, so its GGUFs read zeros there, and its draft accepted 2.3
+tokens a cycle on code without the embedding, 4.7 with it. Every family loads
+and runs its draft alike: a decode cycle drafts eight rows, the anchor and its
+seven proposals, and the target verifies their first rows, as many as the cycle
 chose for the lane (its verify rows, `ops::DraftSelector::addVerifyRows`), never
 more than the lane may still retain: the tokens its output budget leaves, at
 most all eight. A dense target verifies every row up to that limit: its

README.md

@@ -68,6 +68,7 @@ automatically:
 | --- | --- | --- |
 | Qwen3.8-27B | `unsloth/Qwen3.8-27B-GGUF:UD-Q4_K_M` | `mlx-community/Qwen3.8-27B-4bit` |
 | Qwen3.6-35B-A3B | `unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_M` | `mlx-community/Qwen3.6-35B-A3B-4bit` |
+| Qwen3.5-122B-A10B, such as Ornith-1.5-122B | `unsloth/Ornith-1.5-122B-GGUF:UD-Q4_K_XL` | — |
 | Qwen3.8-Flash-Next | `ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF:Q2_0` | — |
 
 Unsloth GGUF variants span **1–8 bits**, including mixed-precision UD formats
@@ -82,6 +83,7 @@ checkpoints and NVFP4 GGUFs also load. Smaller variants run on
 [24 GB Macs](docs/performance.md#smaller-ggufs-on-24-gb-macs).
 [27B variants](https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/tree/main) ·
 [35B variants](https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/tree/main) ·
+[Ornith-1.5-122B variants](https://huggingface.co/unsloth/Ornith-1.5-122B-GGUF/tree/main) ·
 [Flash-Next variants](https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF/tree/main)
 
 Vision and the tokenizer come from the target model's source.

dev/benchmarks/attention_sweep.mm

@@ -12,7 +12,7 @@
 // case prints a digest of its output: two builds that fill the same pages must
 // print the same digests, whatever their storage.
 //
-// usage: attention-sweep METALLIB [--histories 0,2048,...] [--shapes 27b,35b,fn]
+// usage: attention-sweep METALLIB [--histories 0,2048,...] [--shapes 27b,35b,122b,fn]
 //                        [--lanes 1,4] [--repeat N] [--phases both|verify|prefill]
 //                        [--compare-metallib PATH] [--kv-format int8|bf16]
 //                        [--extent-pages N] [--qsa shared|independent|overlapping|pages|indexer]
@@ -94,7 +94,7 @@
 constexpr uint32_t kLayer = 1;
 // The memory plan's extents of the full models: 16, 10 and 12 attention
 // layers.
-constexpr uint32_t kModelLayers27b = 16, kModelLayers35b = 10, kModelLayersFn = 12;
+constexpr uint32_t kModelLayers27b = 16, kModelLayers35b = 10, kModelLayers122b = 12, kModelLayersFn = 12;
 // Qwen3.8-Flash-Next's QSA: the indexer's top_k of 2048 tokens in blocks of 4.
 constexpr uint32_t kQsaTopBlocks = 512, kQsaBlockTokens = SPLASH_ATTENTION_BLOCK_TOKENS;
 // --qsa overlapping: the blocks a tile's rows choose theirs from, three times
@@ -104,16 +104,23 @@
 
 AttentionShape shapeOf(const std::string &shape, kv::Format format) {
   if (shape == "fn") return AttentionShape{24, 2, 256, format, true};
+  if (shape == "122b") return AttentionShape{32, 2, 256, format};
   return shape == "27b" ? AttentionShape{24, 4, 256, format}
                         : AttentionShape{16, 2, 256, format};
 }
 
 uint32_t modelLayers(const std::string &shape) {
-  return shape == "27b" ? kModelLayers27b : shape == "35b" ? kModelLayers35b : kModelLayersFn;
+  return shape == "27b"    ? kModelLayers27b
+         : shape == "35b"  ? kModelLayers35b
+         : shape == "122b" ? kModelLayers122b
+                           : kModelLayersFn;
 }
 
 std::string modelName(const std::string &shape) {
-  return shape == "27b" ? "qwen3.8-27b" : shape == "35b" ? "qwen3.6-35b-a3b" : "qwen3.8-flash-next";
+  return shape == "27b"    ? "qwen3.8-27b"
+         : shape == "35b"  ? "qwen3.6-35b-a3b"
+         : shape == "122b" ? "qwen3.5-122b-a10b"
+                           : "qwen3.8-flash-next";
 }
 
 // A 32-bit mix of a block (and a row), uniform enough to scatter selections.
@@ -560,7 +567,7 @@ uint32_t parseCount(std::string_view text, uint32_t minimum, uint32_t maximum,
 int main(int argc, const char *argv[]) {
   try {
     if (argc < 2) {
-      std::cerr << "usage: attention-sweep METALLIB [--histories LIST] [--shapes 27b,35b,fn] "
+      std::cerr << "usage: attention-sweep METALLIB [--histories LIST] [--shapes 27b,35b,122b,fn] "
                    "[--lanes LIST] [--repeat N] [--phases both|verify|prefill] "
                    "[--compare-metallib PATH] [--kv-format int8|bf16] "
                    "[--extent-pages N] [--qsa shared|independent|overlapping|pages|indexer] "
@@ -617,8 +624,8 @@ int main(int argc, const char *argv[]) {
           const size_t comma = value.find(',', start);
           const std::string shape =
               value.substr(start, comma == std::string::npos ? std::string::npos : comma - start);
-          if (shape != "27b" && shape != "35b" && shape != "fn")
-            throw std::invalid_argument("--shapes takes 27b, 35b or fn");
+          if (shape != "27b" && shape != "35b" && shape != "122b" && shape != "fn")
+            throw std::invalid_argument("--shapes takes 27b, 35b, 122b or fn");
           shapes.push_back(shape);
           if (comma == std::string::npos) break;
           start = comma + 1;

dev/tests/attention_sweep_cli.sh

@@ -28,7 +28,7 @@ for value in -1 1048576 4294967296 10x '' '0,'; do
     reject '--histories requires integers' --histories "$value"
 done
 for value in typo '' '27b,' '27b,typo' ',35b'; do
-    reject '--shapes takes 27b, 35b or fn' --shapes "$value"
+    reject '--shapes takes 27b, 35b, 122b or fn' --shapes "$value"
 done
 for value in 0 16384 -1 1x ''; do
     reject '--extent-pages requires integers' --extent-pages "$value"

dev/tests/engine/TestDrafts.hpp

@@ -13,7 +13,8 @@ namespace splash::test {
 
 // The DFlash2 drafts the configs in dev/tests/fixtures/model-configs
 // describe (model/ModelDescriptor.mm readDraftLayout): the released drafts
-// of Qwen3.8-27B and Qwen3.6-35B-A3B, and one of made-up sizes for
+// of Qwen3.8-27B, Qwen3.6-35B-A3B and Qwen3.5-122B-A10B (Ornith-1.5-122B's,
+// trained to fill a block of 16), and one of made-up sizes for
 // Qwen3.8-Flash-Next.
 [[nodiscard]] inline model::DFlashDraftLayout qwen3_8Draft() {
   constexpr model::Qwen3_8Layout target;
@@ -25,6 +26,11 @@ namespace splash::test {
   constexpr std::array<uint32_t, 8> targetLayers{1, 6, 11, 16, 22, 27, 32, 37};
   return model::draftLayout(target.hiddenSize, target.vocabularySize, 6, 6144, 10'000'000.0F, targetLayers);
 }
+[[nodiscard]] inline model::DFlashDraftLayout qwen3_5_122BDraft() {
+  constexpr model::Qwen3_6MoeLayout target = model::kQwen3_5_122BLayout;
+  constexpr std::array<uint32_t, 8> targetLayers{1, 7, 14, 20, 26, 32, 39, 45};
+  return model::draftLayout(target.hiddenSize, target.vocabularySize, 6, 9216, 10'000'000.0F, targetLayers);
+}
 [[nodiscard]] inline model::DFlashDraftLayout flashNextDraft() {
   constexpr model::Qwen4ExpLayout target;
   constexpr std::array<uint32_t, 4> targetLayers{3, 17, 31, 47};

dev/tests/engine/TestFamilies.hpp

@@ -0,0 +1,55 @@
+#pragma once
+
+#include "model/ModelDescriptor.hpp"
+#include "ops/GDN.hpp"
+#include "ops/PagedKv.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace splash::test {
+
+// The shapes the families Splash serves run (model::kFamilies), each once, in
+// the families' order, which the kernel tests take so that every family's
+// kernels are tested, a new family's with no edit here.
+
+// Each family's GDN shape.
+[[nodiscard]] inline std::vector<ops::GdnShape> familyGdnShapes() {
+  std::vector<ops::GdnShape> shapes;
+  model::forEachFamily([&](const auto &family) {
+    const ops::GdnShape shape{family.gdnKeyHeads,          family.gdnValueHeads,  family.gdnHeadDimension,
+                              family.convolutionDimension, family.packedGdnWidth, family.gdnSigmoidGate};
+    if (std::find(shapes.begin(), shapes.end(), shape) == shapes.end()) shapes.push_back(shape);
+  });
+  return shapes;
+}
+
+// Each family's full attention: its query heads over the KV heads, head
+// dimension and value order of its pages.
+struct AttentionGeometry final {
+  uint32_t queryHeads;
+  uint32_t kvHeads;
+  uint32_t headDimension;
+  bool valuesTokenMajor;
+
+  // Its pages of one attention layer in `format`.
+  [[nodiscard]] kv::Layout layout(kv::Format format) const {
+    return {.attentionLayers = 1, .kvHeads = kvHeads, .headDimension = headDimension, .format = format,
+            .valuesTokenMajor = valuesTokenMajor};
+  }
+  bool operator==(const AttentionGeometry &) const = default;
+};
+[[nodiscard]] inline std::vector<AttentionGeometry> familyAttentionGeometries() {
+  std::vector<AttentionGeometry> geometries;
+  model::forEachFamily([&](const auto &family) {
+    const kv::Layout layout = family.kvLayout();
+    const AttentionGeometry geometry{family.attentionQueryHeads, layout.kvHeads, layout.headDimension,
+                                     layout.valuesTokenMajor};
+    if (std::find(geometries.begin(), geometries.end(), geometry) == geometries.end())
+      geometries.push_back(geometry);
+  });
+  return geometries;
+}
+
+} // namespace splash::test

dev/tests/engine/draft_preparation_test.mm

@@ -11,13 +11,19 @@
 // how to update it.
 #include "GgufFixtures.hpp"
 #include "TestAdmission.hpp"
+#include "TestCheckpoint.hpp"
+#include "TestFiles.hpp"
+#include "metal/abi/ExecutionGeometry.h"
 #include "model/DFlashDraft.hpp"
 #include "model/DraftCheckpoint.hpp"
 #include "model/GgufPreparation.hpp"
 #include "model/SafetensorsCheckpoint.hpp"
 #include "model/WeightImages.hpp"
+#include "ops/Embedding.hpp"
 
+#include <algorithm>
 #include <cstring>
+#include <filesystem>
 #include <vector>
 
 using namespace gguf_fixtures;
@@ -80,6 +86,59 @@ void checkImage(const std::vector<uint8_t> &image, const model::gguf::Image &pla
   }
 }
 
+// A checkpoint that ships its mask token's embedding beside its weights
+// (mask_embedding.safetensors): model.bin holds it after the codebooks, the
+// loader reads it as shipped, and each lane's mask rows of the draft input
+// take it, the lane's anchor row as the table gave it. An embedding of
+// another shape is refused.
+void checkMaskEmbedding(MetalBackend &backend, const std::filesystem::path &root,
+                        const model::DFlashDraftLayout &layout) {
+  const splash::test::TemporaryDirectory directory("splash-draft-mask");
+  for (const auto &entry : std::filesystem::directory_iterator(root))
+    if (entry.path().extension() == ".safetensors")
+      std::filesystem::copy_file(entry.path(), directory.path() / entry.path().filename());
+  const uint64_t rowBytes = uint64_t{layout.hiddenSize} * 2;
+  std::vector<uint8_t> mask(rowBytes);
+  for (uint32_t column = 0; column < layout.hiddenSize; ++column) {
+    const uint16_t bits = uint16_t(0x3400 + column);
+    std::memcpy(mask.data() + uint64_t{column} * 2, &bits, 2);
+  }
+  const std::filesystem::path path = directory.path() / "mask_embedding.safetensors";
+  splash::test::writeSyntheticShard(path, {{"mask_embedding", "BF16", {layout.hiddenSize}, mask}});
+  model::WeightImages images(backend, "mask fixture");
+  model::DraftCheckpointLoader loader(backend, images, directory.path(), layout);
+  check(loader.hasMaskEmbedding(), "the loader finds the mask embedding the checkpoint ships");
+  const model::DFlashDraftWeights draft = model::loadDFlashDraftWeights(backend, loader, layout);
+  check(draft.maskEmbedding && bytesOf(model::contentsOf(draft.maskEmbedding)) == mask,
+        "the draft reads the mask embedding as shipped");
+  // Two lanes of poisoned draft input rows.
+  const uint32_t lanes = 2, rows = lanes * SPLASH_TARGET_VERIFY_ROWS;
+  const auto input =
+      backend.allocateBuffer(rows * rowBytes, splash::metal::BufferStorage::Shared, "draft input rows");
+  std::memset(input.contents(), 0xA5, rows * rowBytes);
+  splash::metal::CommandGraph graph;
+  ops::Embedding::addDraftMaskRows(graph, draft.maskEmbedding, input, layout.hiddenSize, lanes);
+  (void)backend.submitCommandAsync(graph.dispatches()).wait();
+  const std::vector<uint8_t> written = bytesOf(model::contentsOf(input));
+  bool set = true;
+  for (uint32_t row = 0; row < rows; ++row) {
+    const auto begin = written.begin() + row * rowBytes;
+    set = set && (row % SPLASH_TARGET_VERIFY_ROWS
+                      ? std::equal(mask.begin(), mask.end(), begin)
+                      : std::all_of(begin, begin + rowBytes, [](uint8_t byte) { return byte == 0xA5; }));
+  }
+  check(set, "each lane's mask rows take the mask embedding, its anchor row as it was");
+  splash::test::writeSyntheticShard(path, {{"mask_embedding", "BF16", {layout.hiddenSize / 2}, {}}});
+  bool refused = false;
+  try {
+    const model::SafetensorsCheckpoint checkpoint(directory.path());
+    static_cast<void>(model::planDraftImages(checkpoint, layout));
+  } catch (const std::exception &error) {
+    refused = std::string_view(error.what()).find("mask_embedding") != std::string_view::npos;
+  }
+  check(refused, "a mask embedding of another shape is refused");
+}
+
 } // namespace
 
 int main(int argc, char **argv) {
@@ -131,8 +190,9 @@ int main(int argc, char **argv) {
                projection(layer.qkvProjection, layout.qkvSize, layout.hiddenSize) &&
                projection(layer.outputProjection, layout.hiddenSize, layout.attentionSize) &&
                projection(layer.downProjection, layout.hiddenSize, layout.intermediateSize);
-      check(read, "the draft loader reads the draft images");
+      check(read && !draft.maskEmbedding, "the draft loader reads the draft images");
     });
+    guarded("draft mask embedding", [&] { checkMaskEmbedding(backend, root, layout); });
     std::printf("%s (%d failures)\n", failures ? "draft preparation tests FAILED" : "draft preparation tests passed",
                 failures);
     return failures ? 1 : 0;

dev/tests/engine/gdn_decode_metal_test.mm

@@ -15,6 +15,7 @@
 // decode over a bf16 state is checked against the fp32 decode.
 #include "TestBuffers.hpp"
 #include "TestChecks.hpp"
+#include "TestFamilies.hpp"
 #include "metal/MetalBackend.hpp"
 #include "metal/abi/ExecutionGeometry.h"
 #include "metal/abi/GDN.h"
@@ -58,12 +59,9 @@
 constexpr uint32_t kLayers = 2;
 // The scales of the normalized q and k rows.
 constexpr double kQueryScale = 0.0078125, kKeyScale = 0.08838834765;
-// Qwen3.6-35B-A3B's, Qwen3.8-27B's and Qwen3.8-Flash-Next's, whose 16512-column
-// packed rows have kernel variants of their own (_p16512) and whose output
-// gate is sigmoid(z), the others' silu(z).
-constexpr std::array kShapes{GdnShape{16, 48, 128, 10240, 16640},
-                             GdnShape{16, 32, 128, 8192, 12544},
-                             GdnShape{16, 48, 128, 10240, 16512, true}};
+// Every family's shape (Qwen3.8-Flash-Next's gate is sigmoid(z), the others'
+// silu(z)).
+const std::vector<GdnShape> kShapes = splash::test::familyGdnShapes();
 
 double roundBfloat(double value) {
   return bf16ToFloat(floatToBf16(static_cast<float>(value)));
@@ -1051,7 +1049,7 @@ void gateMatchesPrefill(MetalBackend &backend, const GdnShape &shape, bool float
 }
 
 void rejectsInvalid(MetalBackend &backend) {
-  const GdnShape &shape = kShapes[1];
+  const GdnShape &shape = kShapes.front();
   Fixture fixture(backend, shape, 1);
   CommandGraph graph;
   for (const uint32_t lanes : {0U, kMaxLanes + 1})

dev/tests/engine/gdn_metal_test.mm

@@ -13,6 +13,7 @@
 
 #include "LinearNumerics.hpp"
 #include "NormReference.hpp"
+#include "TestFamilies.hpp"
 
 #import <Foundation/Foundation.h>
 
@@ -771,9 +772,7 @@ void bufferExtents(MetalBackend &backend, const GdnShape &shape, bool float32,
 
 void run(const std::string &metallib) {
   MetalBackend backend(metallib);
-  for (const GdnShape &shape :
-       {GdnShape{16, 48, 128, 10240, 16640}, GdnShape{16, 32, 128, 8192, 12544},
-        GdnShape{16, 48, 128, 10240, 16512, true}}) {
+  for (const GdnShape &shape : splash::test::familyGdnShapes()) {
     for (bool float32 : {false, true})
       bufferExtents(backend, shape, float32, GdnStateFormat::Float32);
     bufferExtents(backend, shape, false, GdnStateFormat::BFloat16);