From 0341410151be123db9cbee0bb7140a084dd1b945 Mon Sep 17 00:00:00 2001 From: Zhijian Liu <5782437+zhijian-liu@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:27:58 -0700 Subject: [PATCH 1/3] Read a split GGUF whose first part holds the metadata alone llama.cpp's gguf-split writes the first part of a split GGUF with the metadata alone, ending at its header: its data section is empty and the padding before it is not written. The parser required that padding and refused the part as truncated. A part of no tensors that ends at its header now has an empty data section. unsloth's Ornith-1.5-122B GGUFs are split this way. Co-Authored-By: Claude Opus 5.5 --- dev/tests/engine/gguf_file_test.cpp | 28 ++++++++++++++++++++++++++++ runtime/model/GgufFile.cpp | 7 ++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/dev/tests/engine/gguf_file_test.cpp b/dev/tests/engine/gguf_file_test.cpp index 6c8c4c469..1d3603cf6 100644 --- a/dev/tests/engine/gguf_file_test.cpp +++ b/dev/tests/engine/gguf_file_test.cpp @@ -97,6 +97,33 @@ void checkSplit(const std::filesystem::path &directory) { "the GGUF's split files hold another number of tensors than split.tensors.count", "accepted a split.tensors.count the parts do not hold"); } +// A GGUF split as llama.cpp's gguf-split writes it: the first part holds the +// metadata alone and ends at its header, unpadded (its data section is +// empty), and the tensors are in the parts after it. +void checkMetadataPart(const std::filesystem::path &directory) { + const std::vector tensors{filled("a", 8, 0xa1), filled("b", 16, 0xb2)}; + const std::vector paths{directory / "metadata-00001-of-00002.gguf", + directory / "metadata-00002-of-00002.gguf"}; + // An architecture name of the length that leaves the header unaligned. + std::string architecture = "fixture"; + gguf::Bytes first = part({gguf::stringKey("general.architecture", architecture)}, 0, 2, 2, {}); + if (first.size() % gguf::kAlignment == 0) + first = part({gguf::stringKey("general.architecture", architecture += "s")}, 0, 2, 2, {}); + splash::test::writeFile(paths[0], first); + splash::test::writeFile(paths[1], part({}, 1, 2, 2, tensors)); + const splash::model::OpenGguf split(paths); + const splash::model::GgufFile &file = split.file(); + require(file.files() == 2 && file.tensors().size() == 2 && file.metadata().architecture() == architecture && + file.source(0).dataOffset() == first.size(), + "a split GGUF's metadata-only first part was not read as an empty data section"); + for (const gguf::Tensor &expected : tensors) { + const splash::model::GgufTensor &tensor = file.require(expected.name); + gguf::Bytes read(tensor.bytes); + file.source(tensor.file).readData(tensor.offset, read); + require(tensor.file == 1 && read == expected.data, + "split GGUF tensor " + expected.name + " is not read from the part after the metadata"); + } +} } // namespace int main() { @@ -159,6 +186,7 @@ int main() { } rejects("excessive nesting", array(nested), "nesting is too deep"); checkSplit(directory.path()); + checkMetadataPart(directory.path()); std::cout << "GGUF parser bounds tests passed\n"; return 0; } catch (const std::exception &error) { diff --git a/runtime/model/GgufFile.cpp b/runtime/model/GgufFile.cpp index eab91c47a..ae1944f90 100644 --- a/runtime/model/GgufFile.cpp +++ b/runtime/model/GgufFile.cpp @@ -279,7 +279,12 @@ void GgufFile::parse(WeightSource &source, uint32_t file) { tensors_.push_back(std::move(tensor)); } const uint64_t headerEnd = reader.position(); - const uint64_t padding = (alignment - headerEnd % alignment) % alignment; + // A part of no tensors may end at its header, as llama.cpp's gguf-split + // writes the first part, of the metadata alone: its data section is empty, + // and the padding before it is not written. + const uint64_t padding = tensors_.size() == firstTensor && headerEnd == source.bytes() + ? 0 + : (alignment - headerEnd % alignment) % alignment; if (padding > source.bytes() - headerEnd) throw GgufError("GGUF data section is truncated"); source.setDataOffset(headerEnd + padding); const uint64_t dataBytes = source.bytes() - source.dataOffset(); From 2faa0124d9d6d887a00210743af5dd2b3256ab00 Mon Sep 17 00:00:00 2001 From: Zhijian Liu <5782437+zhijian-liu@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:27:58 -0700 Subject: [PATCH 2/3] Serve Qwen3.5-122B-A10B and its fine-tune Ornith-1.5-122B A family of the Qwen3.6-35B-A3B's architecture and layers at a larger size (hidden 3072, 48 layers, 32 query heads over 2 KV heads, 64 GDN value heads, experts of 1024). Both name the qwen3_5_moe_text model type and the qwen35moe GGUF architecture, so a config's sizes pick its family. Its draft is Ornith-1.5-122B's DFlash2, trained to fill a block of 16: a draft's block may now be the kernels' 8 rows or more, of which a cycle drafts the first. Families and kernel shapes are data, so another size of an architecture is one value, and a shape no family had its kernels and one table row: - A layout type is an architecture, whose code its families share. A family is a value of it, named on the shared base layout, and every family is one entry of model::kFamilies, which refusals, family lookup and the tests walk. The 122B is kQwen3_5_122BLayout, a Qwen3_6MoeLayout with no code of its own. - A config's family is the one of its model type whose seven sizes it holds. A config of none is refused as it differs from the family it matches furthest. - GDN and attention kernels are chosen from one table of compiled shapes each (kCompiledGdn, kCompiledAttention), their names built from a base and the shape's suffix. New rows: GDN for 64 value heads (_vh64) and attention for 16 query heads per KV head (_kv2_g16). - The kernel tests take every family's GDN and attention shapes from kFamilies (dev/tests/engine/TestFamilies.hpp), so a new family's kernels are tested with no edit to them. A dense paged-attention tile runs in the fewest parts that divide its eight rows and leave each at most 64 fused rows, the 256-thread tile's (SPLASH_ATTENTION_TILE_PARTS, splash_attention_tile_part): one part up to eight query heads per KV head, two up to sixteen, four up to thirty-two. The kernels, host plans and tests read the one rule; today's tiles are one part and unchanged. Measured on an M5 Max (40 cores), interleaved rounds of the attention sweep, which now takes the 122b shape: the two 64-row parts of a group of sixteen ran its 32K and 128K prefill and its 128K verify 20-40% faster than one 128-row tile of 512 threads (32K prefill 67 against 95 ms, 128K verify 1.21 against 1.70 ms), both parts reading the same pages, the second mostly from cache. Splitting further was slower at long histories: sixteen in four parts (128K verify, four lanes, 5.18 against 4.54 ms) and eight in two (128K prefill 172 against 134 ms). Outputs are bit-identical across part counts. Checked on an M5 Max (128 GB) with unsloth/Ornith-1.5-122B-GGUF:UD-Q4_K_XL and its experts streamed: the runtime oracle (fp32 and bf16 state), the kernel tests under shader validation, the native CPU and Python checks, and text, tool calls and vision over HTTP. Co-Authored-By: Claude Opus 5.5 --- DEVELOPMENT.md | 36 +- README.md | 2 + dev/benchmarks/attention_sweep.mm | 21 +- dev/tests/attention_sweep_cli.sh | 2 +- dev/tests/engine/TestDrafts.hpp | 8 +- dev/tests/engine/TestFamilies.hpp | 55 ++ dev/tests/engine/gdn_decode_metal_test.mm | 12 +- dev/tests/engine/gdn_metal_test.mm | 5 +- dev/tests/engine/model_configuration_test.cpp | 43 +- dev/tests/engine/paged_attention_plan_test.mm | 45 +- .../engine/q8_flash_attention_metal_test.mm | 60 +- .../qwen3.5-122b-a10b/config.json | 930 ++++++++++++++++++ .../qwen3.5-122b-a10b/draft/config.json | 61 ++ .../qwen3.5-122b-a10b/gguf-metadata.json | 30 + dev/tests/install/test_upstream.py | 8 +- install/completions/suggested-models.txt | 1 + install/families.py | 3 + runtime/metal/abi/PagedAttention.h | 14 + .../kernels/common/paged_attention_tile.h | 27 +- .../metal/kernels/decode/attention_qkv.metal | 4 + runtime/metal/kernels/decode/gdn.metal | 5 + .../kernels/decode/paged_attention.metal | 67 +- .../metal/kernels/prefill/attention_qkv.metal | 3 + runtime/metal/kernels/prefill/gdn.metal | 5 + .../kernels/prefill/paged_attention.metal | 59 +- runtime/model/ModelDescriptor.hpp | 12 + runtime/model/ModelDescriptor.mm | 94 +- runtime/model/Qwen3_6Moe.hpp | 38 +- runtime/model/Qwen3_8.hpp | 8 +- runtime/model/Qwen4Exp.hpp | 8 +- runtime/model/QwenHybridLayout.hpp | 12 +- runtime/model/QwenTargetLoader.hpp | 2 +- runtime/ops/GDN.cpp | 52 +- runtime/ops/PagedAttention.cpp | 163 ++- runtime/ops/PagedAttention.hpp | 26 +- 35 files changed, 1617 insertions(+), 304 deletions(-) create mode 100644 dev/tests/engine/TestFamilies.hpp create mode 100644 dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/config.json create mode 100644 dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/draft/config.json create mode 100644 dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/gguf-metadata.json diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index e9e843113..f9b19ea05 100644 --- a/DEVELOPMENT.md +++ b/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,14 +1853,16 @@ 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 -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 +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. 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 projections read their weights once for all of a cycle's rows, so a row costs it little. Each row a MoE target verifies routes its own experts, which a rejected row reads for nothing (a Qwen3.8-Flash-Next cycle of one lane routes about 43 diff --git a/README.md b/README.md index a45ebaa3c..1de978bae 100644 --- a/README.md +++ b/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. diff --git a/dev/benchmarks/attention_sweep.mm b/dev/benchmarks/attention_sweep.mm index d14df5a01..a801bb9b4 100644 --- a/dev/benchmarks/attention_sweep.mm +++ b/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; diff --git a/dev/tests/attention_sweep_cli.sh b/dev/tests/attention_sweep_cli.sh index 88a6b7a06..09a55ed6a 100644 --- a/dev/tests/attention_sweep_cli.sh +++ b/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" diff --git a/dev/tests/engine/TestDrafts.hpp b/dev/tests/engine/TestDrafts.hpp index 22aa628d1..67d9f7108 100644 --- a/dev/tests/engine/TestDrafts.hpp +++ b/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 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 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 targetLayers{3, 17, 31, 47}; diff --git a/dev/tests/engine/TestFamilies.hpp b/dev/tests/engine/TestFamilies.hpp new file mode 100644 index 000000000..f9f069974 --- /dev/null +++ b/dev/tests/engine/TestFamilies.hpp @@ -0,0 +1,55 @@ +#pragma once + +#include "model/ModelDescriptor.hpp" +#include "ops/GDN.hpp" +#include "ops/PagedKv.hpp" + +#include +#include +#include + +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 familyGdnShapes() { + std::vector 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 familyAttentionGeometries() { + std::vector 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 diff --git a/dev/tests/engine/gdn_decode_metal_test.mm b/dev/tests/engine/gdn_decode_metal_test.mm index c494b3ce3..6a2ce5771 100644 --- a/dev/tests/engine/gdn_decode_metal_test.mm +++ b/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 kShapes = splash::test::familyGdnShapes(); double roundBfloat(double value) { return bf16ToFloat(floatToBf16(static_cast(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}) diff --git a/dev/tests/engine/gdn_metal_test.mm b/dev/tests/engine/gdn_metal_test.mm index 890288cd7..e22f70ab8 100644 --- a/dev/tests/engine/gdn_metal_test.mm +++ b/dev/tests/engine/gdn_metal_test.mm @@ -13,6 +13,7 @@ #include "LinearNumerics.hpp" #include "NormReference.hpp" +#include "TestFamilies.hpp" #import @@ -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); diff --git a/dev/tests/engine/model_configuration_test.cpp b/dev/tests/engine/model_configuration_test.cpp index cf22f3be9..d018c6c89 100644 --- a/dev/tests/engine/model_configuration_test.cpp +++ b/dev/tests/engine/model_configuration_test.cpp @@ -71,6 +71,12 @@ constexpr std::string_view kMoeGgufConfig = R"({"model_type":"qwen3_5_moe","text_config":{"model_type":"qwen3_5_moe_text","hidden_size":2048,)" R"("max_position_embeddings":262144,"num_attention_heads":16,"num_key_value_heads":2,)" R"("head_dim":256,"num_hidden_layers":40,"vocab_size":248320}})"; +// Qwen3.5-122B-A10B's, of the same model type, from unsloth's +// Ornith-1.5-122B GGUF (its MTP layer left out of the layers). +constexpr std::string_view kLargeMoeGgufConfig = + R"({"model_type":"qwen3_5_moe","text_config":{"model_type":"qwen3_5_moe_text","hidden_size":3072,)" + R"("max_position_embeddings":262144,"num_attention_heads":32,"num_key_value_heads":2,)" + R"("head_dim":256,"num_hidden_layers":48,"vocab_size":248320}})"; // The descriptor inspection makes of model's assembly. model::ModelDescriptor inspect(const SourceModel &model) { @@ -111,6 +117,15 @@ void testFamilies(const std::filesystem::path &fixtures) { moeDescriptor.visionSource == model::VisionSource::Safetensors && moeDescriptor.draft == test::qwen3_6MoeDraft(), "an MLX Qwen3.6-35B-A3B made another descriptor"); + // Qwen3.5-122B-A10B's model type is the Qwen3.6-35B-A3B's: its hidden size + // names its family. + const model::ModelDescriptor largeMoeDescriptor = + inspect(mlxModel(fixtures, "qwen3.5-122b-a10b").with(&SourceModel::record, R"("none")", R"("safetensors")")); + require(largeMoeDescriptor.target == model::TargetLayout{model::kQwen3_5_122BLayout} && + largeMoeDescriptor.visionSource == model::VisionSource::Safetensors && + largeMoeDescriptor.vision.outputHiddenSize == 3072 && + largeMoeDescriptor.draft == test::qwen3_5_122BDraft(), + "an MLX Qwen3.5-122B-A10B made another descriptor"); // Qwen3.8-Flash-Next's GGUF config, with its draft's. const model::ModelDescriptor flashDescriptor = inspect(mlxModel(fixtures, "qwen3.8-flash-next").with(&SourceModel::record, R"("mlx-affine")", R"("gguf")")); @@ -161,7 +176,7 @@ void testFamilies(const std::filesystem::path &fixtures) { // tell its target from other models of the architecture, which Splash does // not serve, such as Qwen3.5-4B. constexpr std::string_view supported = - "; supported: Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.8-Flash-Next"; + "; supported: Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next"; struct Refused final { std::string_view from, to, difference; }; @@ -176,6 +191,14 @@ void testFamilies(const std::filesystem::path &fixtures) { "no supported model has this architecture (" + std::string(refused.difference) + ")" + std::string(supported), "a config of another model was accepted with " + std::string(refused.to)); + // A config of a model type of several families and the sizes of none is + // refused as it differs from the one it matches furthest. + refuses(mlxModel(fixtures, "qwen3.5-122b-a10b") + .with(&SourceModel::config, R"("num_hidden_layers": 48)", R"("num_hidden_layers": 40)"), + "no supported model has this architecture (text config num_hidden_layers: safetensors 40, " + "Qwen3.5-122B-A10B 48)" + + std::string(supported), + "a config of no family's sizes was accepted"); } // A descriptor states where its weights come from: one that leaves either @@ -580,6 +603,24 @@ void testConfigurationCheck(const std::filesystem::path &fixtures) { require(check("gguf", "none", kDenseGgufConfig, denseMetadata, dense.draft) == "Qwen3.8-27B" && check("gguf", "none", kMoeGgufConfig, moeMetadata, std::nullopt) == "Qwen3.6-35B-A3B", "a GGUF's derived config and metadata were not found to be its family"); + // Qwen3.5-122B-A10B, of the same model type and GGUF architecture as the + // Qwen3.6-35B-A3B, as its sizes tell; its draft, trained to fill a block + // of 16, drafts the kernels' eight rows, and a block of fewer is refused. + // Each family refuses the other's draft. + const SourceModel largeMoe = mlxModel(fixtures, "qwen3.5-122b-a10b"); + const std::string largeMoeMetadata = text(fixtures / "qwen3.5-122b-a10b" / "gguf-metadata.json"); + require(check("mlx-affine", "safetensors", largeMoe.config, std::nullopt, largeMoe.draft) == "Qwen3.5-122B-A10B" && + check("gguf", "none", kLargeMoeGgufConfig, largeMoeMetadata, largeMoe.draft) == "Qwen3.5-122B-A10B", + "a Qwen3.5-122B-A10B's configuration was not found to be its family"); + rejects([&] { static_cast(check("mlx-affine", "none", moe.config, std::nullopt, largeMoe.draft)); }, + "draft config hidden_size mismatch: checkpoint 3072, runtime 2048", "a larger family's draft was accepted"); + rejects([&] { static_cast(check("mlx-affine", "none", largeMoe.config, std::nullopt, moe.draft)); }, + "draft config hidden_size mismatch: checkpoint 2048, runtime 3072", "a smaller family's draft was accepted"); + rejects([&] { + static_cast(check("mlx-affine", "none", largeMoe.config, std::nullopt, + largeMoe.with(&SourceModel::draft, R"("block_size": 16)", R"("block_size": 4)").draft)); + }, "draft config dflash_config block_size is less than the 8 rows a cycle drafts", + "a draft of a block shorter than the kernels' rows was accepted"); rejects([&] { static_cast(check("gguf", "none", kDenseGgufConfig, std::nullopt, std::nullopt)); }, "a GGUF target is checked with its metadata", "a GGUF was checked without its metadata"); const auto replaced = [](std::string metadata, std::string_view from, std::string_view to) { diff --git a/dev/tests/engine/paged_attention_plan_test.mm b/dev/tests/engine/paged_attention_plan_test.mm index 9a64e7d0b..3d340d18c 100644 --- a/dev/tests/engine/paged_attention_plan_test.mm +++ b/dev/tests/engine/paged_attention_plan_test.mm @@ -8,6 +8,7 @@ #include "HostKvExtents.hpp" #include "LinearNumerics.hpp" #include "NormReference.hpp" +#include "TestFamilies.hpp" #include #include @@ -45,25 +46,24 @@ using splash::test::rejects; using splash::test::require; -// The production attention geometries: query heads over the KV heads of a -// layout, whose values the group of twelve keeps token-major for its sparse -// kernels. -struct Geometry final { - uint32_t queryHeads; - uint32_t kvHeads; - bool valuesTokenMajor; +// The production attention geometries, every family's: query heads over the +// KV heads of a layout, whose values a QSA family keeps token-major for its +// sparse kernels. +using Geometry = splash::test::AttentionGeometry; +const std::vector kGeometries = splash::test::familyAttentionGeometries(); - [[nodiscard]] kv::Layout layout(kv::Format format) const { - return {.attentionLayers = 1, .kvHeads = kvHeads, .headDimension = 256, .format = format, - .valuesTokenMajor = valuesTokenMajor}; - } -}; -constexpr std::array kGeometries{{{24, 4, false}, {16, 2, false}, {24, 2, true}}}; +// The parts of a layout's tiles, each a threadgroup of the split grid's x +// beside its KV head. +uint32_t tileParts(uint32_t queryHeads, kv::Layout layout) { + return layout.valuesTokenMajor ? 1 : SPLASH_ATTENTION_TILE_PARTS(queryHeads / layout.kvHeads); +} // A layout's split or reduce pipeline of a phase ("prefill" or "verify"). std::string pipelineName(std::string_view phase, bool split, uint32_t queryHeads, kv::Layout layout) { const bool sparse = layout.valuesTokenMajor; - const std::string suffix = layout.kvHeads == 4 ? "" : queryHeads == 24 ? "_kv2_g12" : "_kv2_g8"; + // The kernels of four KV heads, the first compiled, have no suffix. + const std::string suffix = layout.kvHeads == 4 ? "" : "_kv" + std::to_string(layout.kvHeads) + "_g" + + std::to_string(queryHeads / layout.kvHeads); const std::string format = layout.format == kv::Format::Int8 ? "q8_" : "bf16_"; return std::string(phase) + "_attention_" + (split ? format + (sparse ? "sparse" : "split") : std::string(sparse ? "sparse_reduce" : "reduce")) + @@ -79,31 +79,33 @@ void checkPrefillSlotOrientation(uint32_t queryHeads, kv::Layout layout) { bool unequalAxes = false, partialTile = false, multipleSplits = false; + // The split grid's x: each KV head's tile parts. + const uint32_t heads = layout.kvHeads * tileParts(queryHeads, layout); // Enumerate nonsquare grids and partial final tiles. Each logical partial // belongs to exactly one query tile, KV head and balanced history split. for (const uint32_t rows : {17U, 257U, 2048U}) { const auto plan = ops::PagedAttention::prefillPlan(rows, queryHeads, layout); const uint32_t tiles = (rows + 7) / 8; - require(plan.splitGroups.x == layout.kvHeads && + require(plan.splitGroups.x == heads && plan.splitGroups.y == tiles && plan.splitGroups.z == plan.splits, "prefill split axes must be KV head, query tile, split"); unequalAxes |= tiles != plan.splits; partialTile |= rows % 8 != 0; multipleSplits |= plan.splits > 1; - const uint64_t slots = uint64_t{tiles} * layout.kvHeads * plan.splits; + const uint64_t slots = uint64_t{tiles} * heads * plan.splits; std::vector visits(slots, 0); for (uint32_t z = 0; z < plan.splitGroups.z; ++z) for (uint32_t y = 0; y < plan.splitGroups.y; ++y) for (uint32_t x = 0; x < plan.splitGroups.x; ++x) { // group.y selects the query tile and group.z its balanced split; // scratch is [tile][KV head][split]. - const uint64_t slot = (uint64_t{y} * layout.kvHeads + x) * plan.splits + z; + const uint64_t slot = (uint64_t{y} * heads + x) * plan.splits + z; require(slot < visits.size() && visits[slot] == 0, "prefill grid aliases or exceeds a logical partial slot"); ++visits[slot]; require(slot % plan.splits == z && - (slot / plan.splits) % layout.kvHeads == x && - slot / (uint64_t{plan.splits} * layout.kvHeads) == y, + (slot / plan.splits) % heads == x && + slot / (uint64_t{plan.splits} * heads) == y, "prefill dispatch orientation changed logical scratch ownership"); } require(std::all_of(visits.begin(), visits.end(), @@ -126,7 +128,7 @@ void checkPlans(uint32_t queryHeads, kv::Layout layout) { require(plan.rows == rows && plan.splits == splits, "prefill plan lost actual rows"); require(plan.splitPipeline == prefillSplit && plan.reducePipeline == prefillReduce, "prefill plan runs the wrong pipelines"); - require(plan.splitGroups.x == layout.kvHeads && + require(plan.splitGroups.x == layout.kvHeads * tileParts(queryHeads, layout) && plan.splitGroups.y == tiles && plan.splitGroups.z == splits && plan.reduceGroups.x == layout.kvHeads && plan.reduceGroups.y == 8 * queryHeads / layout.kvHeads && @@ -165,7 +167,8 @@ void checkPlans(uint32_t queryHeads, kv::Layout layout) { require(plan.splits == maximum && plan.workspace.partialsBytes == fused * 256 * 4 && plan.workspace.statisticsBytes == fused * 2 * 4 && - plan.splitGroups.x == (sparse ? layout.kvHeads * plan.splits : layout.kvHeads) && + plan.splitGroups.x == + (sparse ? layout.kvHeads * plan.splits : layout.kvHeads * tileParts(queryHeads, layout)) && plan.splitGroups.y == (sparse ? 1 : plan.splits) && plan.splitGroups.z == lanes && plan.reduceGroups.y == 8 * queryHeads / layout.kvHeads && diff --git a/dev/tests/engine/q8_flash_attention_metal_test.mm b/dev/tests/engine/q8_flash_attention_metal_test.mm index bd81bea12..6e29b2abe 100644 --- a/dev/tests/engine/q8_flash_attention_metal_test.mm +++ b/dev/tests/engine/q8_flash_attention_metal_test.mm @@ -3,6 +3,7 @@ #include "HostKvExtents.hpp" #include "TestChecks.hpp" +#include "TestFamilies.hpp" #include "ops/PagedAttention.hpp" #include "Q8PageFormatReference.hpp" @@ -29,30 +30,55 @@ constexpr uint32_t kRows = kVerifyRows; static_assert(kStride == SPLASH_VERIFY_CHUNK_STRIDE); -// The production GQA geometries. The group size selects the kernel -// specialization; the suffix names its pipelines. The group of twelve runs -// the sparse kernels, which attend through block lists, bound before their -// parameters, and read token-major values. +// The production GQA geometries, every family's (TestFamilies.hpp). The +// group size selects the kernel specialization, which the plans name +// (ops::PagedAttention). A QSA family's group runs the sparse kernels, which +// attend through block lists, bound before their parameters, and read +// token-major values; a group of more than eight runs each tile in parts, +// each a threadgroup of the split with partial slots of its own +// (splash_attention_tile_part). struct Shape { uint32_t kvHeads; uint32_t queryHeadsPerKvHead; - const char *suffix; + bool tokenMajor; uint32_t queryHeads() const { return kvHeads * queryHeadsPerKvHead; } uint32_t fusedRows() const { return kRows * queryHeadsPerKvHead; } - bool sparse() const { return queryHeadsPerKvHead == 12; } + bool sparse() const { return tokenMajor; } + uint32_t parts() const { return sparse() ? 1 : SPLASH_ATTENTION_TILE_PARTS(queryHeadsPerKvHead); } + // The partial of a fused row of the tile of KV head `head` from split + // `split` of `slotSplits`: [KV head][part][split][part's fused row]. + uint64_t partialRow(uint32_t head, uint32_t slotSplits, uint32_t split, uint32_t row) const { + const uint32_t partRows = fusedRows() / parts(); + return ((uint64_t{head} * parts() + row / partRows) * slotSplits + split) * partRows + row % partRows; + } Layout layout() const { return {.attentionLayers = 1, .kvHeads = kvHeads, .headDimension = kHeadDimension, .valuesTokenMajor = sparse()}; } - // A phase's ("verify" or "prefill") split and reduce entries. - std::string split(const char *phase) const { - return std::string(phase) + (sparse() ? "_attention_q8_sparse" : "_attention_q8_split") + suffix; + // A phase's ("verify" or "prefill") split and reduce entries, as its plans + // name them. + std::string split(std::string_view phase) const { + return phase == "verify" ? verifyPlan().splitPipeline : prefillPlan().splitPipeline; + } + std::string reduce(std::string_view phase) const { + return phase == "verify" ? verifyPlan().reducePipeline : prefillPlan().reducePipeline; + } + +private: + splash::ops::PrefillAttentionPlan prefillPlan() const { + return splash::ops::PagedAttention::prefillPlan(kRows, queryHeads(), layout()); } - std::string reduce(const char *phase) const { - return std::string(phase) + (sparse() ? "_attention_sparse_reduce" : "_attention_reduce") + suffix; + splash::ops::VerifyAttentionPlan verifyPlan() const { + const std::array history{0}; + return splash::ops::PagedAttention::verifyPlan(1, queryHeads(), layout(), history); } }; -constexpr std::array kShapes{{{4, 6, ""}, {2, 8, "_kv2_g8"}, {2, 12, "_kv2_g12"}}}; +const std::vector kShapes = [] { + std::vector shapes; + for (const splash::test::AttentionGeometry &geometry : splash::test::familyAttentionGeometries()) + shapes.push_back({geometry.kvHeads, geometry.queryHeads / geometry.kvHeads, geometry.valuesTokenMajor}); + return shapes; +}(); using splash::test::require; @@ -536,7 +562,7 @@ Dispatch dispatch(id device, id queue, [encoder useResource:extent usage:MTLResourceUsageRead]; // The sparse split's grid is one-dimensional (ops::PagedAttention::verifyPlan). [encoder dispatchThreadgroups:shape.sparse() ? MTLSizeMake(shape.kvHeads * splits, 1, width) - : MTLSizeMake(shape.kvHeads, splits, width) + : MTLSizeMake(shape.kvHeads * shape.parts(), splits, width) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; [encoder setComputePipelineState:reduce]; [encoder setBuffer:partials offset:0 atIndex:0]; @@ -589,7 +615,7 @@ Dispatch dispatchPrefill(id device, id queue, [encoder setBytes:¶ms length:sizeof(params) atIndex:paramsIndex]; for (id extent : data.extents) [encoder useResource:extent usage:MTLResourceUsageRead]; - [encoder dispatchThreadgroups:MTLSizeMake(shape.kvHeads, 1, splits) + [encoder dispatchThreadgroups:MTLSizeMake(shape.kvHeads * shape.parts(), 1, splits) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; [encoder setComputePipelineState:reduce]; [encoder setBuffer:partials offset:0 atIndex:0]; @@ -805,7 +831,7 @@ void checkReduce(id device, id queue, for (uint32_t head = 0; head < shape.kvHeads; ++head) for (uint32_t split = 0; split < splits; ++split) for (uint32_t row = 0; row < m; ++row) { - const uint64_t index = (uint64_t{head} * splits + split) * m + row; + const uint64_t index = shape.partialRow(head, splits, split, row); const bool active = row / shape.queryHeadsPerKvHead < activeRows; // Alternate underflow-scale differences and fractional exp weights. const float spacing = row % 2 ? 0.25f : 500.0f; @@ -852,11 +878,11 @@ void checkReduce(id device, id queue, } double maximum = -INFINITY; for (uint32_t split = 0; split < splits; ++split) - maximum = std::max(maximum, double(t[((uint64_t{head} * splits + split) * m + row) * 2])); + maximum = std::max(maximum, double(t[shape.partialRow(head, splits, split, row) * 2])); for (uint32_t dim = 0; dim < d; ++dim) { double numerator = 0, denominator = 0; for (uint32_t split = 0; split < splits; ++split) { - const uint64_t index = (uint64_t{head} * splits + split) * m + row; + const uint64_t index = shape.partialRow(head, splits, split, row); const double weight = std::exp(double(t[index * 2]) - maximum); numerator += weight * p[index * d + dim]; denominator += weight * t[index * 2 + 1]; diff --git a/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/config.json b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/config.json new file mode 100644 index 000000000..5d25dc06b --- /dev/null +++ b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/config.json @@ -0,0 +1,930 @@ +{ + "architectures": [ + "Qwen3_5MoeForConditionalGeneration" + ], + "do_sample": true, + "dtype": "bfloat16", + "eos_token_id": [ + 248046, + 248044 + ], + "generation_config": { + "bos_token_id": 248044, + "do_sample": true, + "eos_token_id": [ + 248046, + 248044 + ], + "pad_token_id": 248044, + "temperature": 0.6, + "top_k": 20, + "top_p": 0.95, + "transformers_version": "4.57.0.dev0" + }, + "hidden_size": 3072, + "image_token_id": 248056, + "model_type": "qwen3_5_moe", + "quantization": { + "group_size": 64, + "bits": 4, + "mode": "affine", + "language_model.model.layers.0.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.0.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.1.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.1.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.2.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.2.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.3.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.3.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.4.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.4.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.5.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.5.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.6.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.6.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.7.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.7.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.8.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.8.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.9.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.9.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.10.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.10.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.11.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.11.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.12.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.12.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.13.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.13.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.14.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.14.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.15.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.15.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.16.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.16.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.17.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.17.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.18.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.18.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.19.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.19.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.20.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.20.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.21.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.21.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.22.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.22.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.23.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.23.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.24.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.24.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.25.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.25.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.26.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.26.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.27.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.27.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.28.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.28.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.29.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.29.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.30.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.30.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.31.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.31.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.32.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.32.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.33.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.33.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.34.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.34.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.35.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.35.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.36.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.36.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.37.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.37.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.38.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.38.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.39.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.39.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.40.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.40.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.41.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.41.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.42.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.42.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.43.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.43.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.44.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.44.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.45.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.45.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.46.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.46.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.47.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.47.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + } + }, + "quantization_config": { + "group_size": 64, + "bits": 4, + "mode": "affine", + "language_model.model.layers.0.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.0.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.1.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.1.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.2.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.2.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.3.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.3.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.4.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.4.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.5.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.5.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.6.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.6.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.7.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.7.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.8.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.8.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.9.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.9.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.10.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.10.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.11.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.11.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.12.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.12.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.13.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.13.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.14.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.14.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.15.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.15.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.16.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.16.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.17.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.17.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.18.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.18.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.19.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.19.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.20.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.20.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.21.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.21.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.22.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.22.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.23.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.23.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.24.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.24.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.25.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.25.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.26.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.26.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.27.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.27.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.28.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.28.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.29.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.29.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.30.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.30.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.31.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.31.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.32.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.32.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.33.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.33.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.34.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.34.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.35.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.35.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.36.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.36.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.37.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.37.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.38.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.38.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.39.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.39.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.40.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.40.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.41.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.41.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.42.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.42.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.43.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.43.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.44.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.44.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.45.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.45.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.46.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.46.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.47.mlp.gate": { + "group_size": 64, + "bits": 8 + }, + "language_model.model.layers.47.mlp.shared_expert_gate": { + "group_size": 64, + "bits": 8 + } + }, + "temperature": 0.6, + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "attn_output_gate": true, + "bos_token_id": null, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 64, + "linear_value_head_dim": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_moe_text", + "moe_intermediate_size": 1024, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 32, + "num_experts": 256, + "num_experts_per_tok": 8, + "num_hidden_layers": 48, + "num_key_value_heads": 2, + "output_router_logits": false, + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "router_aux_loss_coef": 0.001, + "shared_expert_intermediate_size": 1024, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 248320 + }, + "tie_word_embeddings": false, + "top_k": 20, + "top_p": 0.95, + "transformers_version": "5.8.1", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "dtype": "bfloat16", + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen3_5_moe_vision", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 3072, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053 +} \ No newline at end of file diff --git a/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/draft/config.json b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/draft/config.json new file mode 100644 index 000000000..78b353536 --- /dev/null +++ b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/draft/config.json @@ -0,0 +1,61 @@ +{ + "architectures": [ + "DFlash2DraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": null, + "dflash_config": { + "block_size": 16, + "conv_group_size": 16, + "conv_kernel_size": 2, + "mask_token_id": 248077, + "selector_rank": 256, + "selector_top_k": 16, + "target_layer_ids": [ + 1, + 7, + 14, + 20, + 26, + 32, + 39, + 45 + ] + }, + "dtype": "bfloat16", + "eos_token_id": 248044, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 9216, + "is_causal": false, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention" + ], + "max_position_embeddings": 262144, + "max_window_layers": 6, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 6, + "num_key_value_heads": 8, + "num_target_layers": 48, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 10000000, + "rope_type": "default" + }, + "sliding_window": 2048, + "tie_word_embeddings": false, + "transformers_version": "5.12.1", + "use_cache": false, + "use_sliding_window": true, + "vocab_size": 248320 +} diff --git a/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/gguf-metadata.json b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/gguf-metadata.json new file mode 100644 index 000000000..2016a6834 --- /dev/null +++ b/dev/tests/fixtures/model-configs/qwen3.5-122b-a10b/gguf-metadata.json @@ -0,0 +1,30 @@ +{ + "float": { + "qwen35moe.attention.layer_norm_rms_epsilon": 9.999999974752427e-07, + "qwen35moe.rope.freq_base": 10000000.0 + }, + "string": { + "general.architecture": "qwen35moe" + }, + "unsigned": { + "qwen35moe.attention.head_count": 32, + "qwen35moe.attention.head_count_kv": 2, + "qwen35moe.attention.key_length": 256, + "qwen35moe.attention.value_length": 256, + "qwen35moe.block_count": 49, + "qwen35moe.context_length": 262144, + "qwen35moe.embedding_length": 3072, + "qwen35moe.expert_count": 256, + "qwen35moe.expert_feed_forward_length": 1024, + "qwen35moe.expert_shared_feed_forward_length": 1024, + "qwen35moe.expert_used_count": 8, + "qwen35moe.full_attention_interval": 4, + "qwen35moe.nextn_predict_layers": 1, + "qwen35moe.rope.dimension_count": 64, + "qwen35moe.ssm.conv_kernel": 4, + "qwen35moe.ssm.group_count": 16, + "qwen35moe.ssm.inner_size": 8192, + "qwen35moe.ssm.state_size": 128, + "qwen35moe.ssm.time_step_rank": 64 + } +} diff --git a/dev/tests/install/test_upstream.py b/dev/tests/install/test_upstream.py index 95da955e3..9c6977bd1 100644 --- a/dev/tests/install/test_upstream.py +++ b/dev/tests/install/test_upstream.py @@ -38,7 +38,7 @@ # tells it apart. UNSUPPORTED = ( "no supported model has this architecture [(]{}[)]; " - "supported: Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.8-Flash-Next" + "supported: Qwen3.8-27B, Qwen3.6-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next" ) @@ -1052,7 +1052,7 @@ def test_new_commit_rejected_before_download_keeps_the_installed_one(self): f"Warning: keeping the installed {MODEL}@{'a' * 12}; cannot install " f"{MODEL}@{'b' * 40}: no supported model has this architecture (text " "config head_dim: safetensors 128, Qwen3.8-27B 256); supported: Qwen3.8-27B, " - "Qwen3.6-35B-A3B, Qwen3.8-Flash-Next", + "Qwen3.6-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next", warnings, ) self.assertEqual(chosen.link.resolve(), installed) @@ -1200,13 +1200,13 @@ def incomplete(root): def test_a_new_draft_of_another_architecture_keeps_the_installed_draft(self): warnings = self.move_target_and_draft( - lambda p: draft_dir(p, DENSE, **{"dflash_config.block_size": 16}) + lambda p: draft_dir(p, DENSE, **{"dflash_config.block_size": 4}) ) self.assertIn( f"Warning: cannot use the {DENSE.name} draft {DENSE.draft_repo}@" f"{'e' * 12}; keeping the installed one: {DENSE.draft_repo}'s draft " f"configuration is incompatible with {DENSE.name}: draft config " - "dflash_config block_size mismatch: checkpoint 16, runtime 8", + "dflash_config block_size is less than the 8 rows a cycle drafts", warnings, ) diff --git a/install/completions/suggested-models.txt b/install/completions/suggested-models.txt index 276df735d..2fa354244 100644 --- a/install/completions/suggested-models.txt +++ b/install/completions/suggested-models.txt @@ -1,4 +1,5 @@ mlx-community/Qwen3.6-35B-A3B-4bit mlx-community/Qwen3.8-27B-4bit +unsloth/Ornith-1.5-122B-GGUF:UD-Q4_K_XL unsloth/Qwen3.6-35B-A3B-GGUF:UD-Q4_K_M unsloth/Qwen3.8-27B-GGUF:UD-Q4_K_M diff --git a/install/families.py b/install/families.py index 89d4e0c40..28f34d4ac 100644 --- a/install/families.py +++ b/install/families.py @@ -26,6 +26,9 @@ class ModelFamily: FAMILIES = ( ModelFamily("Qwen3.8-27B", "incoai/Qwen3.8-27B-DFlash2"), ModelFamily("Qwen3.6-35B-A3B", "incoai/Qwen3.6-35B-A3B-DFlash2"), + # Its fine-tune Ornith-1.5-122B is the family too, and the draft + # released for it the family's. + ModelFamily("Qwen3.5-122B-A10B", "incoai/Ornith-1.5-122B-DFlash2"), ModelFamily("Qwen3.8-Flash-Next", "incoai/Qwen3.8-Flash-Next-DFlash2"), ) diff --git a/runtime/metal/abi/PagedAttention.h b/runtime/metal/abi/PagedAttention.h index 711986707..37329f6d6 100644 --- a/runtime/metal/abi/PagedAttention.h +++ b/runtime/metal/abi/PagedAttention.h @@ -45,6 +45,20 @@ struct SplashChunkedPrefillParams { static_assert(sizeof(SplashChunkedPrefillParams) == 24, "Chunked store parameters are 24 bytes on both sides"); +// The parts a dense paged-attention tile of a GQA group of `group` query heads +// per KV head runs in (kernels/common/paged_attention_tile.h +// splash_attention_tile_part): the fewest that divide its eight rows and leave +// each at most 64 fused rows, the 256-thread tile's. A group of up to eight +// is one part, sixteen two, thirty-two four. Measured on an M5 Max (40 cores), +// the two 64-row parts of a group of sixteen ran its 32K-token prefill and its +// verify past 128K tokens 20-40% faster than one 128-row tile of 512 threads: +// both parts read the same pages, the second mostly from cache, and the +// smaller threadgroups keep more of the GPU busy. Splitting further, a group +// of eight into two parts or sixteen into four, was slower again at long +// histories. +#define SPLASH_ATTENTION_TILE_PARTS(group) \ + ((group) <= 8u ? 1u : (group) <= 16u ? 2u : (group) <= 32u ? 4u : 8u) + // Sparse attention (kernels/shared/sparse_attention.metal, // ops::AttentionBlocks): a query row attends only to the blocks of // SPLASH_ATTENTION_BLOCK_TOKENS tokens its block mask names, within its causal diff --git a/runtime/metal/kernels/common/paged_attention_tile.h b/runtime/metal/kernels/common/paged_attention_tile.h index 45fb48de6..ae2737d92 100644 --- a/runtime/metal/kernels/common/paged_attention_tile.h +++ b/runtime/metal/kernels/common/paged_attention_tile.h @@ -14,9 +14,10 @@ using namespace mpp::tensor_ops; // store writes current rows to their final slots; committed_tokens controls // visibility, and the next command overwrites rejected verify rows. // -// A threadgroup owns one KV head, query tile and history split. It processes -// the tile's GQA group as M = rows x query heads per KV head. Each split writes -// fp32 partials and statistics for a fixed-order reduce. +// A threadgroup owns one KV head, query tile (or a part of one, +// splash_attention_tile_part) and history split. It processes the GQA group +// as M = rows x query heads per KV head, at most 64. Each split writes fp32 +// partials and statistics for a fixed-order reduce. // // Prefill and verify share this device-operand page loop. constant uint SplashVerifyMaximumSplits = @@ -97,6 +98,26 @@ inline SplashAttentionTile splash_prefill_attention_tile( params.split_count, true}; } +// Part `part` of a tile run as Parts parts of RowsPerTile / Parts rows each +// (SPLASH_ATTENTION_TILE_PARTS: at most 64 fused rows each, the 256-thread page +// loop's), each a threadgroup of the split over the same pages. A part is a +// tile of its own rows: their queries and output, positions and active rows, +// and partial slots of its own, slot_stride (the slot splits of each of the +// tile's slots) apart. A part past the tile's active rows is inactive. One +// part is the tile. +template +inline SplashAttentionTile splash_attention_tile_part(SplashAttentionTile tile, uint part, + uint slot_stride) { + static_assert(Parts > 0 && RowsPerTile % Parts == 0, "parts of whole rows"); + constexpr uint Rows = RowsPerTile / Parts; + const uint first = part * Rows; + if (!tile.active || first >= tile.active_rows) + return {}; + return {tile.offset + ulong(first) * QueryHeadsPerKVHead * SplashKvHeadDimension, + tile.first_slot * Parts + ulong(part) * slot_stride, tile.first_position + first, + min(Rows, tile.active_rows - first), tile.splits, true}; +} + // Split and reduce derive the same balanced partition of each query tile's // visible pages. Causal masking remains per query row inside each split. inline uint splash_attention_pages(uint visible_tokens) { diff --git a/runtime/metal/kernels/decode/attention_qkv.metal b/runtime/metal/kernels/decode/attention_qkv.metal index 5bc75ffb0..04f63f55a 100644 --- a/runtime/metal/kernels/decode/attention_qkv.metal +++ b/runtime/metal/kernels/decode/attention_qkv.metal @@ -54,6 +54,8 @@ VERIFY_ATTENTION_QKV(verify_attention_qkv_f32, 24, 4, float) VERIFY_ATTENTION_QKV(verify_attention_qkv_kv2_g8_f32, 16, 2, float) VERIFY_ATTENTION_QKV(verify_attention_qkv_kv2_g12, 24, 2, bfloat) VERIFY_ATTENTION_QKV(verify_attention_qkv_kv2_g12_f32, 24, 2, float) +VERIFY_ATTENTION_QKV(verify_attention_qkv_kv2_g16, 32, 2, bfloat) +VERIFY_ATTENTION_QKV(verify_attention_qkv_kv2_g16_f32, 32, 2, float) #undef VERIFY_ATTENTION_QKV // Element `element` of every verify lane's rows, lane by lane. @@ -93,6 +95,7 @@ inline void full_attention_gate_decode_phase( VERIFY_ATTENTION_GATE(verify_attention_gate, 24, 4) VERIFY_ATTENTION_GATE(verify_attention_gate_kv2_g8, 16, 2) VERIFY_ATTENTION_GATE(verify_attention_gate_kv2_g12, 24, 2) +VERIFY_ATTENTION_GATE(verify_attention_gate_kv2_g16, 32, 2) #undef VERIFY_ATTENTION_GATE #define ATTENTION_GATE_TABLE(Name, QHeads, KHeads, Layout) \ @@ -115,4 +118,5 @@ VERIFY_ATTENTION_GATE(verify_attention_gate_kv2_g12, 24, 2) ATTENTION_GATE_TABLE(verify_attention_gate_table16, 24, 4, gguf_sg::Table16) ATTENTION_GATE_TABLE(verify_attention_gate_table16_kv2_g8, 16, 2, gguf_sg::Table16) ATTENTION_GATE_TABLE(verify_attention_gate_table16_kv2_g12, 24, 2, gguf_sg::Table16) +ATTENTION_GATE_TABLE(verify_attention_gate_table16_kv2_g16, 32, 2, gguf_sg::Table16) #undef ATTENTION_GATE_TABLE diff --git a/runtime/metal/kernels/decode/gdn.metal b/runtime/metal/kernels/decode/gdn.metal index f327d963b..05c5f3c39 100644 --- a/runtime/metal/kernels/decode/gdn.metal +++ b/runtime/metal/kernels/decode/gdn.metal @@ -454,6 +454,11 @@ GDN_DECODE_ENTRY(verify_gdn_fused_f32, 16, 48, 128, 10240, 16640, float) GDN_DECODE_ENTRY(verify_gdn_fused_vh32_f32, 16, 32, 128, 8192, 12544, float) GDN_DECODE_TABLE_ENTRY(verify_gdn_fused_table16_f32, 16, 48, 128, 10240, 16640, gguf_sg::Table16, float) GDN_DECODE_TABLE_ENTRY(verify_gdn_fused_table16_vh32_f32, 16, 32, 128, 8192, 12544, gguf_sg::Table16, float) +// Qwen3.5-122B-A10B (64 value heads). +GDN_DECODE_ENTRY(verify_gdn_fused_vh64, 16, 64, 128, 12288, 20736, bfloat) +GDN_DECODE_ENTRY(verify_gdn_fused_vh64_f32, 16, 64, 128, 12288, 20736, float) +GDN_DECODE_TABLE_ENTRY(verify_gdn_fused_table16_vh64, 16, 64, 128, 12288, 20736, gguf_sg::Table16, bfloat) +GDN_DECODE_TABLE_ENTRY(verify_gdn_fused_table16_vh64_f32, 16, 64, 128, 12288, 20736, gguf_sg::Table16, float) // Qwen3.8-Flash-Next (16512-column rows). GDN_DECODE_ENTRY(verify_gdn_fused_p16512, 16, 48, 128, 10240, 16512, bfloat) GDN_DECODE_ENTRY(verify_gdn_fused_p16512_f32, 16, 48, 128, 10240, 16512, float) diff --git a/runtime/metal/kernels/decode/paged_attention.metal b/runtime/metal/kernels/decode/paged_attention.metal index 4341ad81e..0382984e2 100644 --- a/runtime/metal/kernels/decode/paged_attention.metal +++ b/runtime/metal/kernels/decode/paged_attention.metal @@ -21,8 +21,8 @@ uint3 group [[threadgroup_position_in_grid]], \ uint thread_index [[thread_index_in_threadgroup]]) -#define PAGED_VERIFY_SCRATCH(Group) \ - constexpr uint M = Group * SPLASH_TARGET_VERIFY_ROWS; \ +#define PAGED_VERIFY_SCRATCH(Group, Parts) \ + constexpr uint M = Group * SPLASH_TARGET_VERIFY_ROWS / Parts; \ constexpr uint N = SplashKvPageTokens; \ alignas(16) threadgroup float scores[M * N]; \ alignas(16) threadgroup bfloat probabilities[M * N]; \ @@ -31,61 +31,77 @@ threadgroup float previous_scale[M]; \ threadgroup atomic_uint rescale; -// group.x is the KV head, group.y the history split and group.z the lane, -// whose parameters select the page table and the query tile. -#define PAGED_VERIFY_TILE_AT(Heads, Group) \ +// group.x is the KV head and the tile's part (splash_attention_tile_part), +// group.y the history split and group.z the lane, whose parameters select the +// page table and the query tile. +#define PAGED_VERIFY_TILE_AT(Heads, Group, Parts) \ + const uint kv_head = group.x / Parts; \ const SplashAttentionTile tile = \ - splash_verify_attention_tile(params[group.z], group.x, \ - group.z); \ + splash_attention_tile_part( \ + splash_verify_attention_tile(params[group.z], kv_head, \ + group.z), \ + group.x % Parts, params[group.z].slot_splits); \ if (!tile.active || group.y >= tile.splits) \ return; -#define PAGED_VERIFY_SPLIT(Name, Heads, Group, CacheElement) \ +#define PAGED_VERIFY_SPLIT(Name, Heads, Group, Parts, CacheElement) \ PAGED_VERIFY_SPLIT_SIGNATURE(Name) { \ - PAGED_VERIFY_SCRATCH(Group) \ - PAGED_VERIFY_TILE_AT(Heads, Group) \ - splash_paged_attention_tile( \ queries + tile.offset, \ SPLASH_LANE_BINDING(group.z, page_table0, page_table1, page_table2, \ page_table3), \ - params[group.z].kv, group.x, tile.first_position, tile.active_rows, \ + params[group.z].kv, kv_head, tile.first_position, tile.active_rows, \ tile.splits, group.y, partials, statistics, tile.first_slot + group.y, \ scores, probabilities, row_max, row_sum, previous_scale, &rescale, \ thread_index); \ } -PAGED_VERIFY_SPLIT(verify_attention_q8_split, 4, 6, int8_t) -PAGED_VERIFY_SPLIT(verify_attention_q8_split_kv2_g8, 2, 8, int8_t) +PAGED_VERIFY_SPLIT(verify_attention_q8_split, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u), int8_t) +PAGED_VERIFY_SPLIT(verify_attention_q8_split_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u), int8_t) +PAGED_VERIFY_SPLIT(verify_attention_q8_split_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u), int8_t) // BF16 shares the page loop and reduction, without quantization scales. -PAGED_VERIFY_SPLIT(verify_attention_bf16_split, 4, 6, bfloat) -PAGED_VERIFY_SPLIT(verify_attention_bf16_split_kv2_g8, 2, 8, bfloat) +PAGED_VERIFY_SPLIT(verify_attention_bf16_split, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u), bfloat) +PAGED_VERIFY_SPLIT(verify_attention_bf16_split_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u), bfloat) +PAGED_VERIFY_SPLIT(verify_attention_bf16_split_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u), bfloat) #undef PAGED_VERIFY_SPLIT #undef PAGED_VERIFY_TILE_AT #undef PAGED_VERIFY_SCRATCH #undef PAGED_VERIFY_SPLIT_SIGNATURE -template +// Each fused row of the whole tile reduces the splits of its part's slots. +template inline void splash_verify_attention_reduce_phase( device const float *partials, device const float *statistics, device bfloat *output, constant SplashVerifyAttentionParams *params, uint3 group, uint thread_index, threadgroup float *weights, threadgroup float *group_values) { constexpr ushort M = SPLASH_TARGET_VERIFY_ROWS * QueryHeadsPerKVHead; + constexpr ushort PartM = M / Parts; constexpr ushort D = SplashKvHeadDimension; const uint fused_row = group.y; - const SplashAttentionTile tile = + const SplashAttentionTile whole = splash_verify_attention_tile(params[group.z], group.x, group.z); - if (!tile.active || fused_row >= M || thread_index >= D) + if (!whole.active || fused_row >= M || thread_index >= D) + return; + const SplashAttentionTile tile = + splash_attention_tile_part( + whole, fused_row / PartM, params[group.z].slot_splits); + if (!tile.active) { + output[whole.offset + ulong(fused_row) * D + thread_index] = bfloat(0.0f); return; - splash_attention_reduce_row( + } + splash_attention_reduce_row( partials, statistics, output + tile.offset, tile.first_position, - tile.active_rows, tile.splits, tile.first_slot, fused_row, thread_index, + tile.active_rows, tile.splits, tile.first_slot, fused_row % PartM, thread_index, weights, group_values); } -#define PAGED_VERIFY_REDUCE(Name, Heads, Group) \ +#define PAGED_VERIFY_REDUCE(Name, Heads, Group, Parts) \ kernel void Name( \ device const float *partials [[buffer(0)]], \ device const float *statistics [[buffer(1)]], \ @@ -95,10 +111,11 @@ inline void splash_verify_attention_reduce_phase( uint thread_index [[thread_index_in_threadgroup]]) { \ threadgroup float weights[SplashVerifyMaximumSplits]; \ threadgroup float group_values[8]; \ - splash_verify_attention_reduce_phase( \ + splash_verify_attention_reduce_phase( \ partials, statistics, output, params, group, thread_index, weights, \ group_values); \ } -PAGED_VERIFY_REDUCE(verify_attention_reduce, 4, 6) -PAGED_VERIFY_REDUCE(verify_attention_reduce_kv2_g8, 2, 8) +PAGED_VERIFY_REDUCE(verify_attention_reduce, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u)) +PAGED_VERIFY_REDUCE(verify_attention_reduce_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u)) +PAGED_VERIFY_REDUCE(verify_attention_reduce_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u)) #undef PAGED_VERIFY_REDUCE diff --git a/runtime/metal/kernels/prefill/attention_qkv.metal b/runtime/metal/kernels/prefill/attention_qkv.metal index faa60e9bc..26487a48f 100644 --- a/runtime/metal/kernels/prefill/attention_qkv.metal +++ b/runtime/metal/kernels/prefill/attention_qkv.metal @@ -31,6 +31,8 @@ PREFILL_ATTENTION_QKV(prefill_attention_qkv_f32, 24, 4, float) PREFILL_ATTENTION_QKV(prefill_attention_qkv_kv2_g8_f32, 16, 2, float) PREFILL_ATTENTION_QKV(prefill_attention_qkv_kv2_g12, 24, 2, bfloat) PREFILL_ATTENTION_QKV(prefill_attention_qkv_kv2_g12_f32, 24, 2, float) +PREFILL_ATTENTION_QKV(prefill_attention_qkv_kv2_g16, 32, 2, bfloat) +PREFILL_ATTENTION_QKV(prefill_attention_qkv_kv2_g16_f32, 32, 2, float) #undef PREFILL_ATTENTION_QKV template @@ -57,4 +59,5 @@ inline void full_attention_gate_prefill_phase( PREFILL_ATTENTION_GATE(prefill_attention_gate, 24, 4) PREFILL_ATTENTION_GATE(prefill_attention_gate_kv2_g8, 16, 2) PREFILL_ATTENTION_GATE(prefill_attention_gate_kv2_g12, 24, 2) +PREFILL_ATTENTION_GATE(prefill_attention_gate_kv2_g16, 32, 2) #undef PREFILL_ATTENTION_GATE diff --git a/runtime/metal/kernels/prefill/gdn.metal b/runtime/metal/kernels/prefill/gdn.metal index 21175e226..82500736a 100644 --- a/runtime/metal/kernels/prefill/gdn.metal +++ b/runtime/metal/kernels/prefill/gdn.metal @@ -192,6 +192,7 @@ inline void gdn_scan_prefill_phase( GDN_SCAN_PREFILL_ENTRY(prefill_gdn_scan, 16, 48, 128) GDN_SCAN_PREFILL_ENTRY(prefill_gdn_scan_vh32, 16, 32, 128) +GDN_SCAN_PREFILL_ENTRY(prefill_gdn_scan_vh64, 16, 64, 128) #undef GDN_SCAN_PREFILL_ENTRY #undef GDN_SCAN_PREFILL_KERNEL @@ -304,6 +305,7 @@ inline void gdn_prepare_prefill_phase( GDN_PREPARE_PREFILL_ENTRY(prefill_gdn_prepare, 16, 48, 128, 10240, 16640) GDN_PREPARE_PREFILL_ENTRY(prefill_gdn_prepare_vh32, 16, 32, 128, 8192, 12544) +GDN_PREPARE_PREFILL_ENTRY(prefill_gdn_prepare_vh64, 16, 64, 128, 12288, 20736) // Qwen3.8-Flash-Next packs the 48-head rows to 16512 columns. GDN_PREPARE_PREFILL_ENTRY(prefill_gdn_prepare_p16512, 16, 48, 128, 10240, 16512) #undef GDN_PREPARE_PREFILL_ENTRY @@ -333,6 +335,9 @@ GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_vh32, 16, 32, 128, 8192, 12544, bfloat) GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_f32, 16, 48, 128, 10240, 16640, float) GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_vh32_f32, 16, 32, 128, 8192, 12544, float) +GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_vh64, 16, 64, 128, 12288, 20736, bfloat) +GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_vh64_f32, 16, 64, 128, 12288, 20736, + float) GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_p16512, 16, 48, 128, 10240, 16512, bfloat) GDN_GATE_PREFILL_ENTRY(prefill_gdn_gate_p16512_f32, 16, 48, 128, 10240, 16512, float) #undef GDN_GATE_PREFILL_ENTRY diff --git a/runtime/metal/kernels/prefill/paged_attention.metal b/runtime/metal/kernels/prefill/paged_attention.metal index 5773d28b5..a02fafdf2 100644 --- a/runtime/metal/kernels/prefill/paged_attention.metal +++ b/runtime/metal/kernels/prefill/paged_attention.metal @@ -3,8 +3,10 @@ // Prefill runs the shared device-operand page loop // (splash_paged_attention_tile) over eight query rows of one KV head. -// Prefill split: KV heads vary first, then query tiles, then history splits. -template +// Prefill split: KV heads (and each one's tile parts, +// splash_attention_tile_part) vary first, then query tiles, then history +// splits. +template inline void splash_prefill_attention_split_phase( device bfloat *queries, device float *partials, device float *statistics, device const SplashKvPage *page_table, @@ -14,19 +16,23 @@ inline void splash_prefill_attention_split_phase( threadgroup float *previous_scale, threadgroup atomic_uint *rescale, uint3 group, uint thread_index) { const uint split = group.z; + const uint kv_head = group.x / Parts; const SplashAttentionTile tile = - splash_prefill_attention_tile(params, group.x, group.y); + splash_attention_tile_part( + splash_prefill_attention_tile(params, kv_head, group.y), + group.x % Parts, params.split_count); if (!tile.active || split >= tile.splits) return; splash_paged_attention_tile( - queries + tile.offset, page_table, params.kv, group.x, tile.first_position, + SPLASH_PREFILL_ATTENTION_TILE_ROWS / Parts, CacheElement>( + queries + tile.offset, page_table, params.kv, kv_head, tile.first_position, tile.active_rows, tile.splits, split, partials, statistics, tile.first_slot + split, scores, probabilities, row_max, row_sum, previous_scale, rescale, thread_index); } -template +// Each fused row of the whole tile reduces the splits of its part's slots. +template inline void splash_prefill_attention_reduce_phase( device const float *partials, device const float *statistics, device bfloat *output, @@ -34,20 +40,28 @@ inline void splash_prefill_attention_reduce_phase( uint thread_index, threadgroup float *weights, threadgroup float *group_values) { constexpr ushort M = SPLASH_PREFILL_ATTENTION_TILE_ROWS * QueryHeadsPerKVHead; + constexpr ushort PartM = M / Parts; constexpr ushort D = SplashKvHeadDimension; const uint fused_row = group.y; - const SplashAttentionTile tile = + const SplashAttentionTile whole = splash_prefill_attention_tile(params, group.x, group.z); - if (!tile.active || fused_row >= M || thread_index >= D) + if (!whole.active || fused_row >= M || thread_index >= D) + return; + const SplashAttentionTile tile = + splash_attention_tile_part( + whole, fused_row / PartM, params.split_count); + if (!tile.active) { + output[whole.offset + ulong(fused_row) * D + thread_index] = bfloat(0.0f); return; + } splash_attention_reduce_row( + SPLASH_PREFILL_ATTENTION_TILE_ROWS / Parts>( partials, statistics, output + tile.offset, tile.first_position, - tile.active_rows, tile.splits, tile.first_slot, fused_row, thread_index, + tile.active_rows, tile.splits, tile.first_slot, fused_row % PartM, thread_index, weights, group_values); } -#define PAGED_PREFILL_SPLIT(Name, Heads, Group, CacheElement) \ +#define PAGED_PREFILL_SPLIT(Name, Heads, Group, Parts, CacheElement) \ kernel void Name( \ device bfloat *queries [[buffer(0)]], \ device float *partials [[buffer(1)]], \ @@ -56,7 +70,7 @@ inline void splash_prefill_attention_reduce_phase( constant SplashPrefillAttentionParams ¶ms [[buffer(4)]], \ uint3 group [[threadgroup_position_in_grid]], \ uint thread_index [[thread_index_in_threadgroup]]) { \ - constexpr uint M = Group * SPLASH_PREFILL_ATTENTION_TILE_ROWS; \ + constexpr uint M = Group * SPLASH_PREFILL_ATTENTION_TILE_ROWS / Parts; \ constexpr uint N = SplashKvPageTokens; \ alignas(16) threadgroup float scores[M * N]; \ alignas(16) threadgroup bfloat probabilities[M * N]; \ @@ -64,20 +78,22 @@ inline void splash_prefill_attention_reduce_phase( threadgroup float row_sum[M]; \ threadgroup float previous_scale[M]; \ threadgroup atomic_uint rescale; \ - splash_prefill_attention_split_phase( \ + splash_prefill_attention_split_phase( \ queries, partials, statistics, page_table, params, scores, \ probabilities, row_max, row_sum, previous_scale, &rescale, group, \ thread_index); \ } -PAGED_PREFILL_SPLIT(prefill_attention_q8_split, 4, 6, int8_t) -PAGED_PREFILL_SPLIT(prefill_attention_q8_split_kv2_g8, 2, 8, int8_t) +PAGED_PREFILL_SPLIT(prefill_attention_q8_split, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u), int8_t) +PAGED_PREFILL_SPLIT(prefill_attention_q8_split_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u), int8_t) +PAGED_PREFILL_SPLIT(prefill_attention_q8_split_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u), int8_t) // BF16 shares the page loop and reduction, without quantization scales. -PAGED_PREFILL_SPLIT(prefill_attention_bf16_split, 4, 6, bfloat) -PAGED_PREFILL_SPLIT(prefill_attention_bf16_split_kv2_g8, 2, 8, bfloat) +PAGED_PREFILL_SPLIT(prefill_attention_bf16_split, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u), bfloat) +PAGED_PREFILL_SPLIT(prefill_attention_bf16_split_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u), bfloat) +PAGED_PREFILL_SPLIT(prefill_attention_bf16_split_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u), bfloat) #undef PAGED_PREFILL_SPLIT -#define PAGED_PREFILL_REDUCE(Name, Heads, Group) \ +#define PAGED_PREFILL_REDUCE(Name, Heads, Group, Parts) \ kernel void Name( \ device const float *partials [[buffer(0)]], \ device const float *statistics [[buffer(1)]], \ @@ -87,10 +103,11 @@ PAGED_PREFILL_SPLIT(prefill_attention_bf16_split_kv2_g8, 2, 8, bfloat) uint thread_index [[thread_index_in_threadgroup]]) { \ threadgroup float weights[SplashPrefillMaximumSplits]; \ threadgroup float group_values[8]; \ - splash_prefill_attention_reduce_phase( \ + splash_prefill_attention_reduce_phase( \ partials, statistics, output, params, group, thread_index, weights, \ group_values); \ } -PAGED_PREFILL_REDUCE(prefill_attention_reduce, 4, 6) -PAGED_PREFILL_REDUCE(prefill_attention_reduce_kv2_g8, 2, 8) +PAGED_PREFILL_REDUCE(prefill_attention_reduce, 4, 6, SPLASH_ATTENTION_TILE_PARTS(6u)) +PAGED_PREFILL_REDUCE(prefill_attention_reduce_kv2_g8, 2, 8, SPLASH_ATTENTION_TILE_PARTS(8u)) +PAGED_PREFILL_REDUCE(prefill_attention_reduce_kv2_g16, 2, 16, SPLASH_ATTENTION_TILE_PARTS(16u)) #undef PAGED_PREFILL_REDUCE diff --git a/runtime/model/ModelDescriptor.hpp b/runtime/model/ModelDescriptor.hpp index e788dad52..1102c01cf 100644 --- a/runtime/model/ModelDescriptor.hpp +++ b/runtime/model/ModelDescriptor.hpp @@ -13,12 +13,24 @@ #include #include #include +#include #include namespace splash::model { using TargetLayout = std::variant; +// The families Splash serves, values of the layout types TargetLayout holds, +// one type for each architecture's code. A family of an architecture is one +// more value here. Refusals name them in this order, and a config's model type +// finds its family among those of its type in it (describeSourceModel). +inline constexpr std::tuple kFamilies{Qwen3_8Layout{}, Qwen3_6MoeLayout{}, kQwen3_5_122BLayout, Qwen4ExpLayout{}}; + +// Calls `use` with each family's layout, in kFamilies' order. +template constexpr void forEachFamily(Use use) { + std::apply([&](const auto &...family) { (use(family), ...); }, kFamilies); +} + // Where a model's weights come from: a safetensors checkpoint (MLX's, Model // Optimizer's or compressed-tensors') or a GGUF, prepared into images when it // loads. The vision tower is None for a model installed with --language-only. diff --git a/runtime/model/ModelDescriptor.mm b/runtime/model/ModelDescriptor.mm index 8f9ee8175..c86e51b2c 100644 --- a/runtime/model/ModelDescriptor.mm +++ b/runtime/model/ModelDescriptor.mm @@ -194,13 +194,6 @@ void requireLayerTypes(NSArray *types, const QwenTargetDimensions &target, } } -// Calls `use` with each family's layout, in TargetLayout's order. -template constexpr void forEachFamily(Use use) { - [&](std::index_sequence) { - (use(std::variant_alternative_t{}), ...); - }(std::make_index_sequence>{}); -} - // Each family's native window, which its config's max_position_embeddings // states, and the window YaRN stretches it to fit the runtime's KV ceiling. template constexpr bool fitsKvCeiling(const Layout &layout) { @@ -230,29 +223,44 @@ void requireLayerTypes(NSArray *types, const QwenTargetDimensions &target, supported); } -// The target's text configuration, of the family its model type names. Every -// one holds the sizes the descriptor shares with it, which tell the family's -// target from other models of its architecture. A safetensors target's -// config.json also holds the rest of what the kernels compute; a GGUF's is -// what the installer derived from the GGUF's metadata, which the GGUF planner -// checks (gguf::requireMetadata). -void validateTextConfig(NSDictionary *text, const QwenTargetDimensions &target, std::string_view family, - TargetSource source) { - const std::string_view name = sourceName(source); - for (const ExpectedNumber &size : std::initializer_list{ - {"hidden_size", target.hiddenSize}, - {"num_hidden_layers", target.layers}, - {"vocab_size", target.vocabularySize}, - {"max_position_embeddings", target.nativeContextTokens}, - {"num_attention_heads", target.attentionQueryHeads}, - {"num_key_value_heads", target.attentionKvHeads}, - {"head_dim", target.attentionHeadDimension}}) { +// How a text configuration differs from a family in the sizes every one +// holds, which tell the family's target from other models of its +// architecture: the first that differs, after `matched` that agree. +struct SizeDifference final { + size_t matched; + std::string difference; +}; +std::optional sizeDifference(NSDictionary *text, const QwenHybridLayout &target, + TargetSource source) { + const std::initializer_list sizes{ + {"hidden_size", target.hiddenSize}, + {"num_hidden_layers", target.layers}, + {"vocab_size", target.vocabularySize}, + {"max_position_embeddings", target.nativeContextTokens}, + {"num_attention_heads", target.attentionQueryHeads}, + {"num_key_value_heads", target.attentionKvHeads}, + {"head_dim", target.attentionHeadDimension}}; + size_t matched = 0; + for (const ExpectedNumber &size : sizes) { const std::string label = std::string("text config ") + size.key; NSNumber *value = requireNumber(text[@(size.key)], label); if (value.doubleValue != size.value) - throw unsupportedModel(label + ": " + std::string(name) + " " + value.description.UTF8String + ", " + - std::string(family) + " " + @(size.value).description.UTF8String); + return SizeDifference{matched, label + ": " + std::string(sourceName(source)) + " " + + value.description.UTF8String + ", " + std::string(target.family) + + " " + @(size.value).description.UTF8String}; + ++matched; } + return std::nullopt; +} + +// The target's text configuration, of its family (describeSourceModel), whose +// sizes it holds (sizeDifference). A safetensors target's config.json also +// holds the rest of what the kernels compute; a GGUF's is what the installer +// derived from the GGUF's metadata, which the GGUF planner checks +// (gguf::requireMetadata). +void validateTextConfig(NSDictionary *text, const QwenHybridLayout &target, TargetSource source) { + const std::string_view name = sourceName(source); + if (const auto difference = sizeDifference(text, target, source)) throw unsupportedModel(difference->difference); if (source != TargetSource::Safetensors) return; requireNumbers(text, name, "text config", {{"linear_num_key_heads", target.gdnKeyHeads}, @@ -535,8 +543,9 @@ void requireQuantization(NSDictionary *config, const QwenTargetDimensions &geome } // A DFlash2 checkpoint's config: the draft's layout. Its attention widths, -// block, window, convolutions and selector must be those the draft kernels -// are built for; its hidden size, vocabulary and mask token the target's; +// window, convolutions and selector must be those the draft kernels are built +// for, and its block at least their rows; its hidden size, vocabulary and +// mask token the target's; // and the target layers it reads (target_layer_ids) some of the target's, // ascending. Its layers, intermediate size and rotary base are its own. DFlashDraftLayout readDraftLayout(NSDictionary *draft, const QwenTargetDimensions &target) { @@ -567,12 +576,18 @@ DFlashDraftLayout readDraftLayout(NSDictionary *draft, const QwenTargetDimension NSDictionary *flash = requireObject(draft, @"dflash_config", "draft config dflash_config"); requireNumbers(flash, "checkpoint", "draft config dflash_config", - {{"block_size", ExecutionLimits::draftQueryRows}, - {"conv_group_size", SPLASH_DRAFT_CONVOLUTION_GROUP}, + {{"conv_group_size", SPLASH_DRAFT_CONVOLUTION_GROUP}, {"conv_kernel_size", SPLASH_DRAFT_CONVOLUTION_TAPS}, {"selector_rank", SPLASH_DRAFT_SELECTOR_RANK}, {"selector_top_k", SPLASH_DRAFT_CANDIDATES}, {"mask_token_id", target.maskToken}}); + // The block it was trained to fill: the kernels' rows or more, of which a + // cycle drafts the first (Ornith-1.5-122B's fills 16). Its weights have no + // part of the block's length. + if (requireWhole(flash[@"block_size"], "draft config dflash_config block_size") < + ExecutionLimits::draftQueryRows) + throw std::invalid_argument("draft config dflash_config block_size is less than the " + + std::to_string(ExecutionLimits::draftQueryRows) + " rows a cycle drafts"); // How a target layer's output row becomes the hidden values the draft reads // of it, as the draft was trained (target_hidden_reduction): a // hyper-connected target's streams by their mean (Qwen4ExpTarget's @@ -632,16 +647,23 @@ ModelDescriptor describeSourceModel(std::string name, std::string_view targetFor NSDictionary *text = config[@"text_config"]; if (![text isKindOfClass:[NSDictionary class]]) throw unsupportedModel("its config has no text_config"); const auto type = requireString(text, @"model_type", "text model type"); - std::optional target; - forEachFamily([&](const auto &layout) { - if (layout.modelType == type) target = layout; - }); - if (!target) throw unsupportedModel("text model type " + type); TargetSource targetSource; // "mlx-affine" names every safetensors target, as installations record it. if (targetFormat == "mlx-affine") targetSource = TargetSource::Safetensors; else if (targetFormat == "gguf") targetSource = TargetSource::Gguf; else throw std::invalid_argument("unsupported target source format: " + std::string(targetFormat)); + // The family of the model type whose sizes the config holds, as + // Qwen3.6-35B-A3B's and Qwen3.5-122B-A10B's tell those two apart; of none, + // the refusal names how it differs from the one it matches furthest. + std::optional target; + std::optional nearest; + forEachFamily([&](const auto &layout) { + if (layout.modelType != type || target) return; + auto difference = sizeDifference(text, layout, targetSource); + if (!difference) target = layout; + else if (!nearest || difference->matched > nearest->matched) nearest = std::move(difference); + }); + if (!target) throw unsupportedModel(nearest ? nearest->difference : "text model type " + type); VisionSource visionSource; if (visionFormat == "none") visionSource = VisionSource::None; else if (visionFormat == "safetensors") visionSource = VisionSource::Safetensors; @@ -654,7 +676,7 @@ ModelDescriptor describeSourceModel(std::string name, std::string_view targetFor std::visit([&](const auto &layout) { if (targetSource == TargetSource::Safetensors && !layout.loadsSafetensors) throw std::invalid_argument(std::string(layout.family) + " loads from a GGUF only"); - validateTextConfig(text, layout, layout.family, targetSource); + validateTextConfig(text, layout, targetSource); if (targetSource == TargetSource::Safetensors) requireQuantization(config, layout); if (draft) dflashLayout = readDraftLayout(draft, layout); }, *target); diff --git a/runtime/model/Qwen3_6Moe.hpp b/runtime/model/Qwen3_6Moe.hpp index 96ce8ce8c..84ca2db46 100644 --- a/runtime/model/Qwen3_6Moe.hpp +++ b/runtime/model/Qwen3_6Moe.hpp @@ -13,15 +13,16 @@ namespace splash::model { +// The Qwen3.5 MoE architecture, whose families are sizes of it: Qwen3.6-35B-A3B +// and Qwen3.5-122B-A10B (kFamilies). struct Qwen3_6MoeLayout final : QwenHybridLayout { - // The family's name, as the installer pairs its draft (install/families.py). - static constexpr std::string_view family = "Qwen3.6-35B-A3B"; // Its config's text model type and its GGUFs' architecture; a safetensors // checkpoint of it loads too. static constexpr std::string_view modelType = "qwen3_5_moe_text"; static constexpr std::string_view ggufArchitecture = "qwen35moe"; static constexpr bool loadsSafetensors = true; + // Qwen3.6-35B-A3B's, its first family (kFamilies). constexpr Qwen3_6MoeLayout() : QwenHybridLayout{{.nativeContextTokens = 262'144, .yarnContextTokens = 1'010'000, @@ -46,10 +47,41 @@ struct Qwen3_6MoeLayout final : QwenHybridLayout { .ffnKind = QwenFfnKind::SparseMoe, .experts = 256, .expertsPerToken = 8, - .expertIntermediateSize = 512}} {} + .expertIntermediateSize = 512}, + "Qwen3.6-35B-A3B"} {} + // A family of the architecture. + constexpr explicit Qwen3_6MoeLayout(const QwenHybridLayout &family) : QwenHybridLayout(family) {} bool operator==(const Qwen3_6MoeLayout &) const = default; }; +// Qwen3.5-122B-A10B and its fine-tunes, such as Ornith-1.5-122B. +inline constexpr Qwen3_6MoeLayout kQwen3_5_122BLayout{ + QwenHybridLayout{{.nativeContextTokens = 262'144, + .yarnContextTokens = 1'010'000, + .layers = 48, + .hiddenSize = 3072, + .vocabularySize = 248320, + .packedGdnWidth = 20736, + .packedFullWidth = 17408, + .convolutionDimension = 12288, + .gdnKeyHeads = 16, + .gdnValueHeads = 64, + .gdnHeadDimension = 128, + .attentionWidth = 8192, + .attentionQueryHeads = 32, + .attentionKvHeads = 2, + .attentionHeadDimension = 256, + .rotaryPairs = SPLASH_TARGET_ROPE_PAIRS, + .rotaryTheta = 10'000'000.0F, + .fullAttentionPeriod = 4, + .maskToken = 248077, + .stopTokens = {248044, 248046}, + .ffnKind = QwenFfnKind::SparseMoe, + .experts = 256, + .expertsPerToken = 8, + .expertIntermediateSize = 1024}, + "Qwen3.5-122B-A10B"}}; + struct Qwen3_6MoeLayerWeights final { ops::NormWeights inputNorm; QwenMixerWeights mixer; diff --git a/runtime/model/Qwen3_8.hpp b/runtime/model/Qwen3_8.hpp index dff25fa80..629fb5383 100644 --- a/runtime/model/Qwen3_8.hpp +++ b/runtime/model/Qwen3_8.hpp @@ -14,14 +14,13 @@ namespace splash::model { struct Qwen3_8Layout final : QwenHybridLayout { - // The family's name, as the installer pairs its draft (install/families.py). - static constexpr std::string_view family = "Qwen3.8-27B"; // Its config's text model type and its GGUFs' architecture; a safetensors // checkpoint of it loads too. static constexpr std::string_view modelType = "qwen3_5_text"; static constexpr std::string_view ggufArchitecture = "qwen35"; static constexpr bool loadsSafetensors = true; + // Qwen3.8-27B's, its first family (kFamilies). constexpr Qwen3_8Layout() : QwenHybridLayout{{.nativeContextTokens = 262'144, .yarnContextTokens = 1'010'000, @@ -44,7 +43,10 @@ struct Qwen3_8Layout final : QwenHybridLayout { .maskToken = 248070, .stopTokens = {248044, 248046}, .ffnKind = QwenFfnKind::Dense, - .intermediateSize = 17408}} {} + .intermediateSize = 17408}, + "Qwen3.8-27B"} {} + // A family of the architecture. + constexpr explicit Qwen3_8Layout(const QwenHybridLayout &family) : QwenHybridLayout(family) {} bool operator==(const Qwen3_8Layout &) const = default; }; diff --git a/runtime/model/Qwen4Exp.hpp b/runtime/model/Qwen4Exp.hpp index 604224958..deb51a2ff 100644 --- a/runtime/model/Qwen4Exp.hpp +++ b/runtime/model/Qwen4Exp.hpp @@ -40,8 +40,6 @@ namespace splash::model { struct SparseRowStoreOptions; struct Qwen4ExpLayout final : QwenHybridLayout { - // The family's name, as the installer pairs its draft (install/families.py). - static constexpr std::string_view family = "Qwen3.8-Flash-Next"; // Its config's text model type and its GGUFs' architecture; no safetensors // checkpoint of it loads. static constexpr std::string_view modelType = "qwen4_exp_text"; @@ -54,6 +52,7 @@ struct Qwen4ExpLayout final : QwenHybridLayout { uint32_t pleEosToken = stopTokens[0]; uint32_t pleImageToken = 248056; + // Qwen3.8-Flash-Next's, its first family (kFamilies). constexpr Qwen4ExpLayout() : QwenHybridLayout{{.nativeContextTokens = 262'144, .yarnContextTokens = 1'010'000, @@ -91,7 +90,10 @@ struct Qwen4ExpLayout final : QwenHybridLayout { .pleNgram = 3, .pleHeadsPerNgram = 8, .pleHeadDimension = 160, - .pleConvolutionTaps = 4}} {} + .pleConvolutionTaps = 4}, + "Qwen3.8-Flash-Next"} {} + // A family of the architecture. + constexpr explicit Qwen4ExpLayout(const QwenHybridLayout &family) : QwenHybridLayout(family) {} bool operator==(const Qwen4ExpLayout &) const = default; }; diff --git a/runtime/model/QwenHybridLayout.hpp b/runtime/model/QwenHybridLayout.hpp index 58e550753..9a0c16c3e 100644 --- a/runtime/model/QwenHybridLayout.hpp +++ b/runtime/model/QwenHybridLayout.hpp @@ -132,12 +132,16 @@ struct QwenTargetDimensions { bool operator==(const QwenTargetDimensions &) const = default; }; -// A family's layout: its dimensions and the cache layouts they make. A -// family sets every value its FFN uses and states what names it among the -// families: its name, its config's model type, the architecture of its GGUFs -// and whether a safetensors checkpoint of it loads (family, modelType, +// A family's layout: its name, its dimensions and the cache layouts they +// make. A layout type is an architecture, whose code its families share; each +// family is a value of it (model::kFamilies) that sets every value its FFN +// uses. The type states its config's model type, the architecture of its GGUFs +// and whether a safetensors checkpoint of it loads (modelType, // ggufArchitecture, loadsSafetensors). struct QwenHybridLayout : QwenTargetDimensions { + // The family's name, as the installer pairs its draft (install/families.py). + std::string_view family; + // The KV pages, which also hold a QSA indexer's pooled keys; a QSA target's // sparse attention gathers blocks of tokens, so its values are token-major. [[nodiscard]] constexpr kv::Layout kvLayout() const noexcept { diff --git a/runtime/model/QwenTargetLoader.hpp b/runtime/model/QwenTargetLoader.hpp index 087528703..5b13a6082 100644 --- a/runtime/model/QwenTargetLoader.hpp +++ b/runtime/model/QwenTargetLoader.hpp @@ -169,7 +169,7 @@ loadQwenTarget(metal::MetalBackend &backend, const Layout &layout, const QwenTar SafetensorsTargetLoader &safetensors = std::get>(files).get(); return read(safetensors, BlockTargetFormat{safetensors.float32Norms(), ops::GdnHeadOrder::Grouped}, experts); } else { - throw WeightStoreError(std::string(Layout::family) + " loads from a GGUF only"); + throw WeightStoreError(std::string(layout.family) + " loads from a GGUF only"); } }(); weights.experts = experts.store; diff --git a/runtime/ops/GDN.cpp b/runtime/ops/GDN.cpp index bf78fe5e8..7d20578c7 100644 --- a/runtime/ops/GDN.cpp +++ b/runtime/ops/GDN.cpp @@ -6,6 +6,7 @@ #include "ops/LaneBindings.hpp" #include +#include #include #include #include @@ -24,35 +25,32 @@ uint32_t headFlags(GdnShape shape, GdnHeadOrder order) noexcept { static_assert(offsetof(GDNDecodeBatchParams, conv_layer_bytes) == 8); -// Value48P16512: Qwen3.8-Flash-Next's 48 value heads, whose packed rows are -// 16512 columns wide. -enum class KernelLayout : uint8_t { Value48, Value32, Value48P16512 }; +// The compiled GDN shapes (kernels/prefill/gdn.metal, kernels/decode/gdn.metal), +// of either gate, and their kernels' name suffix. The scan reads no packed +// rows, so shapes of one head layout share its kernel (scanSuffix). +struct CompiledGdn final { + GdnShape shape; + std::string_view suffix; + std::string_view scanSuffix; +}; +constexpr std::array kCompiledGdn{ + CompiledGdn{{16, 48, 128, 10240, 16640}, "", ""}, + CompiledGdn{{16, 32, 128, 8192, 12544}, "_vh32", "_vh32"}, + // Qwen3.8-Flash-Next's 48 value heads, whose packed rows are 16512 columns wide. + CompiledGdn{{16, 48, 128, 10240, 16512}, "_p16512", ""}, + CompiledGdn{{16, 64, 128, 12288, 20736}, "_vh64", "_vh64"}, +}; -[[nodiscard]] KernelLayout kernelShape(const GdnShape &shape) { +[[nodiscard]] const CompiledGdn &compiledGdn(const GdnShape &shape) { if (!shape.valid()) throw std::invalid_argument("invalid GDN shape"); GdnShape heads = shape; heads.sigmoidGate = false; - if (heads == GdnShape{16, 48, 128, 10240, 16640}) - return KernelLayout::Value48; - if (heads == GdnShape{16, 32, 128, 8192, 12544}) - return KernelLayout::Value32; - if (heads == GdnShape{16, 48, 128, 10240, 16512}) - return KernelLayout::Value48P16512; + for (const CompiledGdn &compiled : kCompiledGdn) + if (compiled.shape == heads) return compiled; throw std::invalid_argument("unsupported compiled GDN shape"); } -// A kernel's variant: the scan reads no packed rows, so all 48-head layouts -// share it (`packed` false). -[[nodiscard]] std::string kernelName(KernelLayout shape, std::string_view value48, - std::string_view value32, bool packed = true) { - if (shape == KernelLayout::Value32) - return std::string(value32); - if (shape == KernelLayout::Value48P16512 && packed) - return std::string(value48) + "_p16512"; - return std::string(value48); -} - uint64_t valueWidth(const GdnShape &shape) { return uint64_t{shape.valueHeads} * shape.headDimension; } // One layer's carried convolution rows and recurrent state, a head dimension // square per value head. @@ -120,8 +118,8 @@ void GDN::addPrefill(metal::CommandGraph &graph, const metal::MetalBackend &back std::span captures) { if (!tokens || !validFormat(state)) throw std::invalid_argument("invalid GDN prefill geometry"); - const KernelLayout kernel = kernelShape(shape); - const std::string gate = normKernel(kernelName(kernel, "prefill_gdn_gate", "prefill_gdn_gate_vh32"), + const CompiledGdn &kernel = compiledGdn(shape); + const std::string gate = normKernel(std::string("prefill_gdn_gate") + std::string(kernel.suffix), buffers.mixerNorm, shape.headDimension); const uint64_t keyWidth = uint64_t{shape.keyHeads} * shape.headDimension; const uint64_t keyRows = tokens * keyWidth * 2; @@ -146,9 +144,9 @@ void GDN::addPrefill(metal::CommandGraph &graph, const metal::MetalBackend &back requireBytes(capture.convolution, carriedBytes(shape), "GDN captured convolution state"); requireBytes(capture.recurrent, recurrentBytes(shape, state), "GDN captured recurrent state"); } - const std::string prepare = kernelName(kernel, "prefill_gdn_prepare", "prefill_gdn_prepare_vh32"); + const std::string prepare = std::string("prefill_gdn_prepare") + std::string(kernel.suffix); const std::string scan = - stateKernel(kernelName(kernel, "prefill_gdn_scan", "prefill_gdn_scan_vh32", false), state); + stateKernel(std::string("prefill_gdn_scan") + std::string(kernel.scanSuffix), state); metal::MetalBuffer convolution = buffers.convolutionIn; metal::MetalBuffer recurrent = buffers.recurrentIn; uint32_t begin = 0; @@ -192,7 +190,7 @@ PreparedInput GDN::addDecode(metal::CommandGraph &graph, GdnDecodeBuffers buffer GdnStateStrides state, GdnHeadOrder order, LinearInput input) { if (!lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH || !state.valid() || buffers.tapeLanes.size() < lanes) throw std::invalid_argument("invalid GDN decode geometry"); - const KernelLayout kernel = kernelShape(shape); + const CompiledGdn &kernel = compiledGdn(shape); // Lane l's rows of the packed and hidden rows are rows [8 l, 8 l + 8). const uint64_t rows = uint64_t{lanes} * SPLASH_TARGET_VERIFY_ROWS; requireBytes(buffers.packed, packedBytes(shape, rows), "GDN packed"); @@ -247,7 +245,7 @@ PreparedInput GDN::addDecode(metal::CommandGraph &graph, GdnDecodeBuffers buffer buffers.hidden}); if (prepare) bindings.insert(bindings.end(), {buffers.linearScratch.input, buffers.linearScratch.sums}); - const std::string name = std::string("verify_gdn_fused") + tableSuffix(input) + kernelName(kernel, "", "_vh32"); + const std::string name = std::string("verify_gdn_fused") + tableSuffix(input) + std::string(kernel.suffix); graph.add(stateKernel(normKernel(name, buffers.mixerNorm, shape.headDimension), state.format), std::move(bindings), params, {shape.valueHeads, lanes, 1}); if (!prepare) return {}; diff --git a/runtime/ops/PagedAttention.cpp b/runtime/ops/PagedAttention.cpp index 9c3b2c9f6..e9e9d954e 100644 --- a/runtime/ops/PagedAttention.cpp +++ b/runtime/ops/PagedAttention.cpp @@ -15,7 +15,11 @@ namespace splash::ops { namespace { -enum class KernelLayout : uint8_t { Kv4Group6, Kv2Group8, Kv2Group12 }; +// The parts a dense tile of a layout's GQA group runs in, each a threadgroup +// of the split (SPLASH_ATTENTION_TILE_PARTS); a sparse tile is one. +uint64_t tileParts(uint32_t queryHeads, kv::Layout layout) { + return layout.valuesTokenMajor ? 1 : SPLASH_ATTENTION_TILE_PARTS(queryHeads / layout.kvHeads); +} // One thread per gate element, across rows x query heads x head dimension. uint64_t gateGroups(uint64_t rows, uint32_t queryHeads, uint32_t headDimension) { @@ -24,57 +28,54 @@ uint64_t gateGroups(uint64_t rows, uint32_t queryHeads, uint32_t headDimension) metal::CommandGraph::kDefaultThreads; } -KernelLayout storageKernelLayout(kv::Layout layout) { - if (!layout.valid() || layout.headDimension != 256) { - throw std::invalid_argument("no KV store kernel for layout"); - } - if (layout.kvHeads == 4) - return KernelLayout::Kv4Group6; - if (layout.kvHeads == 2) - return KernelLayout::Kv2Group8; +// The compiled attention geometries (kernels/{prefill,decode}/attention_qkv.metal, +// paged_attention.metal and paged_attention_store.metal, and +// shared/sparse_attention.metal): query heads over KV heads of 256 dimensions +// in a value order, and their kernels' name suffix, the stores' by KV heads and +// value order. The sparse kernels gather blocks of token-major values +// (Qwen3.8-Flash-Next's); the dense page loop reads dimension-major ones. +struct CompiledAttention final { + uint32_t queryHeads; + uint32_t kvHeads; + bool valuesTokenMajor; + std::string_view suffix; + std::string_view storeSuffix; +}; +constexpr std::array kCompiledAttention{ + CompiledAttention{24, 4, false, "", ""}, + CompiledAttention{16, 2, false, "_kv2_g8", "_kv2_g8"}, + CompiledAttention{32, 2, false, "_kv2_g16", "_kv2_g8"}, + CompiledAttention{24, 2, true, "_kv2_g12", "_kv2_vt"}, +}; + +// The KV store kernels' suffix of a layout's KV heads and value order. +std::string_view storeSuffix(kv::Layout layout) { + if (layout.valid() && layout.headDimension == 256) + for (const CompiledAttention &compiled : kCompiledAttention) + if (compiled.kvHeads == layout.kvHeads && compiled.valuesTokenMajor == layout.valuesTokenMajor) + return compiled.storeSuffix; throw std::invalid_argument("no KV store kernel for layout"); } -KernelLayout attentionKernelLayout(uint32_t queryHeads, kv::Layout layout) { - // Select a compiled GQA variant so kernels need no geometry branches. The - // value order picks the family: the sparse kernels gather blocks of - // token-major values (compiled for 24 query heads over two KV heads, - // Qwen3.8-Flash-Next's), and the dense page loop reads dimension-major ones. - const KernelLayout storage = storageKernelLayout(layout); - if (layout.valuesTokenMajor) { - if (storage != KernelLayout::Kv2Group8 || queryHeads != 24) - throw std::invalid_argument("no sparse attention kernel for layout"); - return KernelLayout::Kv2Group12; - } - if (!((storage == KernelLayout::Kv4Group6 && queryHeads == 24) || - (storage == KernelLayout::Kv2Group8 && queryHeads == 16))) - throw std::invalid_argument("no paged-attention kernel for layout"); - return storage; +// The compiled geometry of `queryHeads` over a layout's KV heads, so the +// kernels need no geometry branches. +const CompiledAttention &compiledAttention(uint32_t queryHeads, kv::Layout layout) { + static_cast(storeSuffix(layout)); + for (const CompiledAttention &compiled : kCompiledAttention) + if (compiled.queryHeads == queryHeads && compiled.kvHeads == layout.kvHeads && + compiled.valuesTokenMajor == layout.valuesTokenMajor) + return compiled; + throw std::invalid_argument(layout.valuesTokenMajor ? "no sparse attention kernel for layout" + : "no paged-attention kernel for layout"); } -// The KV store of a layout's format, KV heads and value order. -std::string_view storePipeline(kv::Layout layout, bool verify) { - const KernelLayout kernel = storageKernelLayout(layout); - const bool bf16 = layout.format == kv::Format::BFloat16; - if (layout.valuesTokenMajor) { - if (kernel != KernelLayout::Kv2Group8) - throw std::invalid_argument("no KV store kernel for the layout's value order"); - return verify ? (bf16 ? "verify_attention_bf16_store_kv2_vt" : "verify_attention_q8_store_kv2_vt") - : (bf16 ? "prefill_attention_bf16_store_kv2_vt" : "prefill_attention_q8_store_kv2_vt"); - } - if (kernel == KernelLayout::Kv4Group6) - return verify ? (bf16 ? "verify_attention_bf16_store" : "verify_attention_q8_store") - : (bf16 ? "prefill_attention_bf16_store" : "prefill_attention_q8_store"); - return verify ? (bf16 ? "verify_attention_bf16_store_kv2_g8" : "verify_attention_q8_store_kv2_g8") - : (bf16 ? "prefill_attention_bf16_store_kv2_g8" : "prefill_attention_q8_store_kv2_g8"); +// A phase's ("prefill" or "verify") attention kernel `name` with a suffix. +std::string kernelName(std::string_view phase, std::string_view name, std::string_view suffix) { + return std::string(phase) + "_attention_" + std::string(name) + std::string(suffix); } -std::string_view pipeline(KernelLayout layout, std::string_view kv4Group6, - std::string_view kv2Group8, std::string_view kv2Group12) noexcept { - return layout == KernelLayout::Kv4Group6 ? kv4Group6 - : layout == KernelLayout::Kv2Group8 ? kv2Group8 - : kv2Group12; -} +// A layout's KV format, as its kernels' names hold it. +std::string kernelFormat(kv::Layout layout) { return layout.format == kv::Format::BFloat16 ? "bf16" : "q8"; } // One dispatch normalizes both the queries and the keys, so their norms must // share a stored type. @@ -168,7 +169,7 @@ uint32_t prefillSplits(uint32_t tiles) { PrefillAttentionPlan PagedAttention::prefillPlan( uint32_t rows, uint32_t queryHeads, kv::Layout layout) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!rows || rows > kv::kChunkedPrefillMaximumRows) throw std::invalid_argument("invalid attention workspace rows"); const uint32_t tiles = kv::prefillAttentionTiles(rows); @@ -178,23 +179,17 @@ PrefillAttentionPlan PagedAttention::prefillPlan( return {rows, splits, attentionWorkspace(uint64_t{tiles} * splits * layout.kvHeads * fusedRows, layout.headDimension), - layout.format == kv::Format::BFloat16 - ? pipeline(kernel, "prefill_attention_bf16_split", - "prefill_attention_bf16_split_kv2_g8", - "prefill_attention_bf16_sparse_kv2_g12") - : pipeline(kernel, "prefill_attention_q8_split", - "prefill_attention_q8_split_kv2_g8", - "prefill_attention_q8_sparse_kv2_g12"), - pipeline(kernel, "prefill_attention_reduce", - "prefill_attention_reduce_kv2_g8", "prefill_attention_sparse_reduce_kv2_g12"), - {layout.kvHeads, tiles, splits}, {layout.kvHeads, fusedRows, tiles}, layout.valuesTokenMajor, + kernelName("prefill", kernelFormat(layout) + (kernel.valuesTokenMajor ? "_sparse" : "_split"), kernel.suffix), + kernelName("prefill", kernel.valuesTokenMajor ? "sparse_reduce" : "reduce", kernel.suffix), + {layout.kvHeads * tileParts(queryHeads, layout), tiles, splits}, + {layout.kvHeads, fusedRows, tiles}, layout.valuesTokenMajor, layout, queryHeads}; } VerifyAttentionPlan PagedAttention::verifyPlan( uint32_t lanes, uint32_t queryHeads, kv::Layout layout, std::span historyTokens) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH) throw std::invalid_argument("invalid attention workspace batch width"); if (historyTokens.size() != lanes) @@ -211,36 +206,27 @@ VerifyAttentionPlan PagedAttention::verifyPlan( } return {lanes, laneSplits, splits, verifyWorkspaceBound(lanes, queryHeads, layout), - layout.format == kv::Format::BFloat16 - ? pipeline(kernel, "verify_attention_bf16_split", - "verify_attention_bf16_split_kv2_g8", - "verify_attention_bf16_sparse_kv2_g12") - : pipeline(kernel, "verify_attention_q8_split", - "verify_attention_q8_split_kv2_g8", - "verify_attention_q8_sparse_kv2_g12"), - pipeline(kernel, "verify_attention_reduce", - "verify_attention_reduce_kv2_g8", "verify_attention_sparse_reduce_kv2_g12"), + kernelName("verify", kernelFormat(layout) + (kernel.valuesTokenMajor ? "_sparse" : "_split"), kernel.suffix), + kernelName("verify", kernel.valuesTokenMajor ? "sparse_reduce" : "reduce", kernel.suffix), // A sparse lane's list fills only the first of the plan's splits. // Its split grid is one-dimensional, x the KV head and then the // split: an M5 Max spread the same order as a two-dimensional grid // (x the KV head, y the split) worse, running it up to 1.7x slower // on most grids of more than 32 splits. layout.valuesTokenMajor ? metal::DispatchSize{layout.kvHeads * splits, 1, lanes} - : metal::DispatchSize{layout.kvHeads, splits, lanes}, + : metal::DispatchSize{layout.kvHeads * tileParts(queryHeads, layout), splits, lanes}, {layout.kvHeads, kv::kVerifyRows * (queryHeads / layout.kvHeads), lanes}, layout.valuesTokenMajor, - pipeline(kernel, {}, {}, - layout.format == kv::Format::BFloat16 ? "verify_attention_bf16_sparse_located_kv2_g12" - : "verify_attention_q8_sparse_located_kv2_g12"), - storePipeline(layout, true), + kernel.valuesTokenMajor ? kernelName("verify", kernelFormat(layout) + "_sparse_located", kernel.suffix) : "", + kernelName("verify", kernelFormat(layout) + "_store", storeSuffix(layout)), {uint64_t{lanes} * 2 * kv::kVerifyRows * layout.kvHeads, 1, 1}, {layout.headDimension, 1, 1}, layout, queryHeads}; } AttentionWorkspace PagedAttention::prefillWorkspace( uint32_t maximumRows, uint32_t queryHeads, kv::Layout layout) { - (void)attentionKernelLayout(queryHeads, layout); + (void)compiledAttention(queryHeads, layout); if (!maximumRows || maximumRows > kv::kChunkedPrefillMaximumRows) throw std::invalid_argument("invalid attention workspace rows"); // Allocation-time bound for every shorter sequence. Encoding resolves its @@ -254,7 +240,7 @@ AttentionWorkspace PagedAttention::prefillWorkspace( AttentionWorkspace PagedAttention::verifyWorkspace( uint32_t lanes, uint32_t queryHeads, kv::Layout layout) { - (void)attentionKernelLayout(queryHeads, layout); + (void)compiledAttention(queryHeads, layout); if (!lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH) throw std::invalid_argument("invalid attention workspace batch width"); return verifyWorkspaceBound(lanes, queryHeads, layout); @@ -267,7 +253,7 @@ void PagedAttention::addPrefillProjection( metal::MetalBuffer queries, metal::MetalBuffer chunkKeys, metal::MetalBuffer chunkValues, uint32_t tokens, uint32_t stride, uint32_t queryHeads, kv::Layout layout) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!tokens || !stride) throw std::invalid_argument("invalid paged prefill projection geometry"); const Staging staging{layout.kvHeads, stride, tokens}; @@ -277,9 +263,7 @@ void PagedAttention::addPrefillProjection( requireBytes(chunkKeys, staging.rowBytes(1, layout.headDimension), "attention key"); requireBytes(chunkValues, staging.valueBytes(layout.headDimension), "attention value"); const FullPrefillParams params{tokens, stride}; - graph.add(qkNormKernel(pipeline(kernel, "prefill_attention_qkv", - "prefill_attention_qkv_kv2_g8", "prefill_attention_qkv_kv2_g12"), - queryNorm, keyNorm, layout.headDimension), + graph.add(qkNormKernel(kernelName("prefill", "qkv", kernel.suffix), queryNorm, keyNorm, layout.headDimension), {std::move(packed), queryNorm.buffer, keyNorm.buffer, std::move(ropeCos), std::move(ropeSin), std::move(queries), std::move(chunkKeys), std::move(chunkValues)}, @@ -291,7 +275,7 @@ void PagedAttention::addPrefillGate( metal::CommandGraph &graph, metal::MetalBuffer packed, metal::MetalBuffer attention, metal::MetalBuffer hidden, uint32_t tokens, uint32_t stride, uint32_t queryHeads, kv::Layout layout) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!tokens || !stride) throw std::invalid_argument("invalid paged prefill gate geometry"); requireBytes(packed, gateBytes(tokens, queryHeads, layout), "attention gate"); @@ -300,8 +284,7 @@ void PagedAttention::addPrefillGate( "attention row"); requireBytes(hidden, uint64_t{tokens} * queryHeads * layout.headDimension * 2, "attention hidden"); const FullPrefillParams params{tokens, stride}; - graph.add(std::string(pipeline(kernel, "prefill_attention_gate", - "prefill_attention_gate_kv2_g8", "prefill_attention_gate_kv2_g12")), + graph.add(kernelName("prefill", "gate", kernel.suffix), {std::move(packed), std::move(attention), std::move(hidden)}, params, {gateGroups(tokens, queryHeads, layout.headDimension), 1, 1}); } @@ -313,7 +296,7 @@ void PagedAttention::addVerifyProjection( metal::MetalBuffer queries, metal::MetalBuffer chunkKeys, metal::MetalBuffer chunkValues, uint32_t queryHeads, kv::Layout layout, uint32_t lanes) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH) throw std::invalid_argument("invalid paged verify projection geometry"); const uint64_t rows = uint64_t{lanes} * kv::kVerifyRows; @@ -323,9 +306,7 @@ void PagedAttention::addVerifyProjection( requireBytes(queries, staging.rowBytes(queryHeads / layout.kvHeads, layout.headDimension), "attention query"); requireBytes(chunkKeys, staging.rowBytes(1, layout.headDimension), "attention key"); requireBytes(chunkValues, staging.valueBytes(layout.headDimension), "attention value"); - graph.add(qkNormKernel(pipeline(kernel, "verify_attention_qkv", - "verify_attention_qkv_kv2_g8", "verify_attention_qkv_kv2_g12"), - queryNorm, keyNorm, layout.headDimension), + graph.add(qkNormKernel(kernelName("verify", "qkv", kernel.suffix), queryNorm, keyNorm, layout.headDimension), {std::move(packed), queryNorm.buffer, keyNorm.buffer, std::move(ropeCos), std::move(ropeSin), std::move(queries), std::move(chunkKeys), std::move(chunkValues)}, @@ -338,7 +319,7 @@ PreparedInput PagedAttention::addVerifyGate( metal::MetalBuffer attention, metal::MetalBuffer hidden, uint32_t queryHeads, kv::Layout layout, uint32_t lanes, LinearScratch scratch, LinearInput input) { - const KernelLayout kernel = attentionKernelLayout(queryHeads, layout); + const CompiledAttention &kernel = compiledAttention(queryHeads, layout); if (!lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH) throw std::invalid_argument("invalid paged verify gate geometry"); const uint32_t rows = lanes * SPLASH_TARGET_VERIFY_ROWS; @@ -349,14 +330,12 @@ PreparedInput PagedAttention::addVerifyGate( if (input != LinearInput::Plain) { const uint32_t width = queryHeads * layout.headDimension; requireTableScratch(scratch, input, width, rows); - std::string name = std::string("verify_attention_gate") + tableSuffix(input); - name += pipeline(kernel, "", "_kv2_g8", "_kv2_g12"); - graph.add(std::move(name), {packed, attention, hidden, scratch.input, scratch.sums}, - {width / 64 * lanes, 1, 1}, {256, 1, 1}); + graph.add(kernelName("verify", std::string("gate") + tableSuffix(input), kernel.suffix), + {packed, attention, hidden, scratch.input, scratch.sums}, {uint64_t{width} / 64 * lanes, 1, 1}, + {256, 1, 1}); return {std::move(hidden), input}; } - graph.add(std::string(pipeline(kernel, "verify_attention_gate", - "verify_attention_gate_kv2_g8", "verify_attention_gate_kv2_g12")), + graph.add(kernelName("verify", "gate", kernel.suffix), {std::move(packed), std::move(attention), hidden}, FullDecodeBatchParams{lanes}, {gateGroups(rows, queryHeads, layout.headDimension), 1, 1}); @@ -388,7 +367,7 @@ void PagedAttention::addPrefillStore( metal::MetalBuffer chunkKeys, metal::MetalBuffer chunkValues, metal::MetalBuffer pageTable, const kv::ChunkedPrefillParams ¶ms, kv::Layout layout) { - const std::string_view store = storePipeline(layout, false); + const std::string store = kernelName("prefill", kernelFormat(layout) + "_store", storeSuffix(layout)); const Staging staging{layout.kvHeads, params.chunk_stride, params.chunk_tokens}; requireBytes(chunkKeys, staging.rowBytes(1, layout.headDimension), "attention key"); requireBytes(chunkValues, staging.valueBytes(layout.headDimension), "attention value"); diff --git a/runtime/ops/PagedAttention.hpp b/runtime/ops/PagedAttention.hpp index 7bbbff77c..ccdc63de6 100644 --- a/runtime/ops/PagedAttention.hpp +++ b/runtime/ops/PagedAttention.hpp @@ -11,8 +11,10 @@ #include #include #include +#include #include #include +#include namespace splash::kv { @@ -113,8 +115,8 @@ struct PrefillAttentionPlan final { const uint32_t tiles; const uint32_t splits; const AttentionWorkspace workspace; - const std::string_view splitPipeline; - const std::string_view reducePipeline; + const std::string splitPipeline; + const std::string reducePipeline; const metal::DispatchSize splitGroups; const metal::DispatchSize reduceGroups; // Whether the kernels attend through block lists (AttentionBlocks). @@ -123,12 +125,12 @@ struct PrefillAttentionPlan final { private: friend class PagedAttention; PrefillAttentionPlan(uint32_t rows, uint32_t splits, AttentionWorkspace workspace, - std::string_view splitPipeline, std::string_view reducePipeline, + std::string splitPipeline, std::string reducePipeline, metal::DispatchSize splitGroups, metal::DispatchSize reduceGroups, bool sparse, kv::Layout layout, uint32_t queryHeads) : layout(layout), queryHeads(queryHeads), rows(rows), tiles(kv::prefillAttentionTiles(rows)), splits(splits), workspace(workspace), - splitPipeline(splitPipeline), reducePipeline(reducePipeline), + splitPipeline(std::move(splitPipeline)), reducePipeline(std::move(reducePipeline)), splitGroups(splitGroups), reduceGroups(reduceGroups), sparse(sparse) {} }; @@ -143,8 +145,8 @@ struct VerifyAttentionPlan final { const std::array laneSplits; const uint32_t splits; const AttentionWorkspace workspace; - const std::string_view splitPipeline; - const std::string_view reducePipeline; + const std::string splitPipeline; + const std::string reducePipeline; const metal::DispatchSize splitGroups; const metal::DispatchSize reduceGroups; // Whether the kernels attend through block lists (AttentionBlocks). @@ -153,23 +155,23 @@ struct VerifyAttentionPlan final { private: friend class PagedAttention; // A sparse plan's split for lists with locations (AttentionBlocks). - const std::string_view locatedSplitPipeline_; - const std::string_view storePipeline_; + const std::string locatedSplitPipeline_; + const std::string storePipeline_; const metal::DispatchSize storeGroups_; const metal::DispatchSize storeThreads_; VerifyAttentionPlan(uint32_t lanes, std::array laneSplits, uint32_t splits, AttentionWorkspace workspace, - std::string_view splitPipeline, std::string_view reducePipeline, + std::string splitPipeline, std::string reducePipeline, metal::DispatchSize splitGroups, metal::DispatchSize reduceGroups, - bool sparse, std::string_view locatedSplitPipeline, std::string_view storePipeline, + bool sparse, std::string locatedSplitPipeline, std::string storePipeline, metal::DispatchSize storeGroups, metal::DispatchSize storeThreads, kv::Layout layout, uint32_t queryHeads) : layout(layout), queryHeads(queryHeads), lanes(lanes), laneSplits(laneSplits), splits(splits), workspace(workspace), - splitPipeline(splitPipeline), reducePipeline(reducePipeline), + splitPipeline(std::move(splitPipeline)), reducePipeline(std::move(reducePipeline)), splitGroups(splitGroups), reduceGroups(reduceGroups), sparse(sparse), - locatedSplitPipeline_(locatedSplitPipeline), storePipeline_(storePipeline), + locatedSplitPipeline_(std::move(locatedSplitPipeline)), storePipeline_(std::move(storePipeline)), storeGroups_(storeGroups), storeThreads_(storeThreads) {} }; From 4dd195e1581d9a5d597e9e0eeff35573c81a8340 Mon Sep 17 00:00:00 2001 From: Zhijian Liu <5782437+zhijian-liu@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:27:58 -0700 Subject: [PATCH 3/3] Read the draft's own mask embedding when its release ships one A draft's rows after its anchor read the target's token table row of the mask token. Qwen3.5-122B-A10B's row, which Ornith-1.5-122B inherits, is of order 1e-6, below what the fp16 block scales of GGUF Q8_0 and the K-quants hold, so a GGUF of the family reads zeros there, and the draft's acceptance on code and reasoning halves (about 2.3 tokens per cycle against 4.9 on the MLX target, which keeps the row). A DFlash2 release may now ship the mask token's BF16 embedding beside its weights, mask_embedding.safetensors holding one tensor mask_embedding of the hidden size. Installation fetches it with the weights; the draft checkpoint keeps it in model.bin after the selector's codebooks; and after each draft input gather one dispatch (draft_mask_rows) sets every lane's mask rows to it. A draft without the file loads and runs as before. With the file beside incoai/Ornith-1.5-122B-DFlash2's weights, unsloth/Ornith-1.5-122B-GGUF:UD-Q4_K_XL drafts 4.7 tokens per cycle on the same coding prompts with every row verified. Resident on an M5 Max, one lane of the backend benchmark decodes 150 tok/s against 71 before (6.40 tokens per cycle against 3.05), and one coding stream over HTTP 105 tok/s against 40. Co-Authored-By: Claude Opus 5.5 --- DEVELOPMENT.md | 19 ++++-- dev/tests/engine/draft_preparation_test.mm | 62 +++++++++++++++++++- dev/tests/install/test_upstream.py | 22 +++++++ install/assembly.py | 4 +- install/upstream.py | 17 ++++-- runtime/metal/abi/Embedding.h | 7 +++ runtime/metal/kernels/shared/embedding.metal | 15 +++++ runtime/model/DFlashDraft.cpp | 9 +++ runtime/model/DFlashDraft.hpp | 8 +++ runtime/model/DraftCheckpoint.cpp | 5 ++ runtime/model/DraftCheckpoint.hpp | 9 +++ runtime/model/Runtime.cpp | 2 + runtime/ops/Embedding.cpp | 11 ++++ runtime/ops/Embedding.hpp | 5 ++ 14 files changed, 182 insertions(+), 13 deletions(-) diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index f9b19ea05..2a20efd46 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -1857,12 +1857,19 @@ target's, and its attention (32 query and 8 KV heads of 128), selector (rank 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. 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 +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 projections read their weights once for all of a cycle's rows, so a row costs it little. Each row a MoE target verifies routes its own experts, which a rejected row reads for nothing (a Qwen3.8-Flash-Next cycle of one lane routes about 43 diff --git a/dev/tests/engine/draft_preparation_test.mm b/dev/tests/engine/draft_preparation_test.mm index 4803a0897..40842ca46 100644 --- a/dev/tests/engine/draft_preparation_test.mm +++ b/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 #include +#include #include using namespace gguf_fixtures; @@ -80,6 +86,59 @@ void checkImage(const std::vector &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 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 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(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; diff --git a/dev/tests/install/test_upstream.py b/dev/tests/install/test_upstream.py index 9c6977bd1..7b1658bc2 100644 --- a/dev/tests/install/test_upstream.py +++ b/dev/tests/install/test_upstream.py @@ -665,6 +665,28 @@ def test_missing_metadata_never_falls_back_to_another_repository(self): self.assertFalse(chosen.models_root.exists()) shutil.rmtree(fake.remote) + def test_a_draft_that_ships_its_mask_embedding_installs_it(self): + # A release may ship its mask token's embedding beside its weights, + # which the draft reads in place of the target's token table row. + fake = FakeHub(self, self.cache) + fake.publish(MODEL, "a" * 40, lambda p: mlx_target(p, DENSE)) + + def release(root): + draft_dir(root, DENSE) + (root / upstream.DRAFT_MASK_EMBEDDING).write_bytes(b"mask") + + fake.publish(DENSE.draft_repo, DRAFT_COMMIT, release) + chosen = selection(self.root) + self.prepare(chosen) + self.assertEqual( + (chosen.link / "draft" / upstream.DRAFT_MASK_EMBEDDING).read_bytes(), + b"mask", + ) + self.assertIn( + "draft/" + upstream.DRAFT_MASK_EMBEDDING, + assembly.verify(chosen.link)["files"], + ) + def test_source_assembly_pairs_the_draft_by_architecture(self): fake = FakeHub(self, self.cache) fake.publish( diff --git a/install/assembly.py b/install/assembly.py index 9a2263c99..0b4afa26b 100644 --- a/install/assembly.py +++ b/install/assembly.py @@ -27,7 +27,9 @@ from the GGUF (server.py main) draft/config.json, draft/.safetensors the DFlash2 checkpoint (ModelDescriptor.mm, - DraftCheckpoint.cpp) + DraftCheckpoint.cpp), with the mask token's + embedding, draft/mask_embedding.safetensors, when + the release ships one vision/ the MLX shards holding vision_tower.* (VisionLoader.cpp, through SafetensorsCheckpoint.mm) vision/mmproj.gguf the GGUF vision projector (VisionLoader.cpp) diff --git a/install/upstream.py b/install/upstream.py index 44267f946..96a2f2a53 100644 --- a/install/upstream.py +++ b/install/upstream.py @@ -45,6 +45,11 @@ # its own, or the image_processor object of processor_config.json, where newer # Transformers releases save it. PROCESSOR_FILES = ("preprocessor_config.json", "processor_config.json") +# A DFlash2 release's BF16 embedding of its mask token, beside its weights, +# which its draft reads in place of the target's token table row +# (DraftCheckpoint.hpp kDraftMaskEmbedding): a quantized table can lose a +# row as small as Qwen3.5-122B-A10B's. +DRAFT_MASK_EMBEDDING = "mask_embedding.safetensors" @dataclass(frozen=True) @@ -771,10 +776,11 @@ def _at(repo): def draft_files(repo, target): """The DFlash2 draft in repo for target's family, downloaded, by assembly path: its config.json and safetensors weights, model.safetensors or the - shards its index names, at the root of the repository or --draft-model - directory, as a DFlash2 release holds them. The weights download only - once the engine accepts its configuration with the target's - (check_model).""" + shards its index names, and the mask token's embedding when the release + ships one (DRAFT_MASK_EMBEDDING), at the root of the repository or + --draft-model directory, as a DFlash2 release holds them. The weights + download only once the engine accepts its configuration with the + target's (check_model).""" family = target.family try: if "config.json" not in repo.files: @@ -799,5 +805,6 @@ def draft_files(repo, target): f"{repo.name}'s draft configuration is incompatible with {family.name}: " f"{error}" ) from error - downloaded = repo.download({"config.json", *weights}) + mask = {DRAFT_MASK_EMBEDDING} & set(repo.files) + downloaded = repo.download({"config.json", *weights, *mask}) return {"draft/" + name: path for name, path in downloaded.items()} diff --git a/runtime/metal/abi/Embedding.h b/runtime/metal/abi/Embedding.h index 91043a6f9..21df9e87d 100644 --- a/runtime/metal/abi/Embedding.h +++ b/runtime/metal/abi/Embedding.h @@ -13,3 +13,10 @@ struct VerifyInputBatchParams { static_assert(sizeof(VerifyInputBatchParams) == 4, "Verify input parameters are 4 bytes on both sides"); + +struct DraftMaskRowsParams { + uint32_t width; +}; + +static_assert(sizeof(DraftMaskRowsParams) == 4, + "Draft mask row parameters are 4 bytes on both sides"); diff --git a/runtime/metal/kernels/shared/embedding.metal b/runtime/metal/kernels/shared/embedding.metal index 5fc0c2e35..9726c8d13 100644 --- a/runtime/metal/kernels/shared/embedding.metal +++ b/runtime/metal/kernels/shared/embedding.metal @@ -78,3 +78,18 @@ kernel void verify_input_tokens( row - 1]; verify_input[index] = min(token, params.vocabulary - 1u); } + +// Each lane's mask rows of its draft input, rows 1 to SPLASH_TARGET_VERIFY_ROWS - 1 (the +// rows after its anchor), set to the mask embedding the draft's checkpoint ships +// (ops/Embedding.cpp addDraftMaskRows): x a span of 256 columns, y the mask row, z the lane. +kernel void draft_mask_rows( + device const bfloat *mask [[buffer(0)]], + device bfloat *rows [[buffer(1)]], + constant DraftMaskRowsParams ¶ms [[buffer(2)]], + uint3 group [[threadgroup_position_in_grid]], + uint thread_index [[thread_index_in_threadgroup]]) { + const uint column = group.x * 256 + thread_index; + if (column >= params.width) + return; + rows[(ulong(group.z) * SPLASH_TARGET_VERIFY_ROWS + 1 + group.y) * params.width + column] = mask[column]; +} diff --git a/runtime/model/DFlashDraft.cpp b/runtime/model/DFlashDraft.cpp index 120ffc8c6..dd5a9571f 100644 --- a/runtime/model/DFlashDraft.cpp +++ b/runtime/model/DFlashDraft.cpp @@ -2,6 +2,7 @@ #include "Checked.hpp" #include "DraftCheckpoint.hpp" #include "metal/abi/DraftAttention.h" +#include "ops/Embedding.hpp" #include #include @@ -343,6 +344,12 @@ void DFlashDraft::addDecode( weights_.selectorProjection, linear.decodePlan(weights_.selectorProjection, lanes)); } +void DFlashDraft::addMaskRows(metal::CommandGraph &graph, metal::MetalBuffer rows, uint32_t lanes) const { + if (!weights_.maskEmbedding) return; + ops::Embedding::addDraftMaskRows(graph, weights_.maskEmbedding, std::move(rows), weights_.layout.hiddenSize, + lanes); +} + void DFlashDraft::addContextCommit( metal::CommandGraph &graph, DFlashContextBuffers buffers, std::span startPositions) const { @@ -447,6 +454,8 @@ DFlashDraftWeights loadDFlashDraftWeights(metal::MetalBackend &backend, file.section(codebookBytes, "predecessor-codebook"); result.successorCodebook = file.section(codebookBytes, "successor-codebook"); + if (files.hasMaskEmbedding()) + result.maskEmbedding = file.section(uint64_t{layout.hiddenSize} * kBFloat16Bytes, "mask-embedding"); file.finish(); result.files.push_back(file.record()); } diff --git a/runtime/model/DFlashDraft.hpp b/runtime/model/DFlashDraft.hpp index b0ab34c73..65797bddd 100644 --- a/runtime/model/DFlashDraft.hpp +++ b/runtime/model/DFlashDraft.hpp @@ -219,6 +219,9 @@ struct DFlashDraftWeights final { ops::Projection selectorProjection; metal::MetalBuffer predecessorCodebook; metal::MetalBuffer successorCodebook; + // The mask token's BF16 embedding the checkpoint ships, or none + // (kDraftMaskEmbedding). + metal::MetalBuffer maskEmbedding; std::vector files; uint64_t actualAllocatedBytes = 0; }; @@ -264,6 +267,11 @@ class DFlashDraft final { const ops::DraftSelectorBuffers &buffers, std::span anchors, std::span policies) const; + // Sets each of `lanes` lanes' mask rows of the draft input `rows` (all but + // its anchor's, row 0 of its draftQueryRows) to the mask embedding the + // checkpoint ships; without one they keep the target table's row. + void addMaskRows(metal::CommandGraph &graph, metal::MetalBuffer rows, uint32_t lanes) const; + void addContextCommit(metal::CommandGraph &graph, DFlashContextBuffers buffers, std::span startPositions) const; diff --git a/runtime/model/DraftCheckpoint.cpp b/runtime/model/DraftCheckpoint.cpp index 9439704a5..1d82c43e7 100644 --- a/runtime/model/DraftCheckpoint.cpp +++ b/runtime/model/DraftCheckpoint.cpp @@ -89,6 +89,7 @@ gguf::Image modelImage(const SafetensorsCheckpoint &checkpoint, const DFlashDraf b.projection({{selector + "hidden_projection", layout.selectorRank}}, hidden); b.stored(selector + "predecessor_codebook", {layout.vocabularySize, layout.selectorRank}); b.stored(selector + "successor_codebook", {layout.vocabularySize, layout.selectorRank}); + if (checkpoint.find(kDraftMaskEmbedding)) b.stored(std::string(kDraftMaskEmbedding), {hidden}); return b.finish(); } @@ -131,4 +132,8 @@ WeightFile DraftCheckpointLoader::layer(uint32_t index) { WeightFile DraftCheckpointLoader::model() { return open(planned_->images.size() - 1); } +bool DraftCheckpointLoader::hasMaskEmbedding() const noexcept { + return planned_->checkpoint.find(kDraftMaskEmbedding) != nullptr; +} + } // namespace splash::model diff --git a/runtime/model/DraftCheckpoint.hpp b/runtime/model/DraftCheckpoint.hpp index 9a429c205..f5ef3abd3 100644 --- a/runtime/model/DraftCheckpoint.hpp +++ b/runtime/model/DraftCheckpoint.hpp @@ -27,6 +27,9 @@ class DraftCheckpointLoader final { [[nodiscard]] WeightFile layer(uint32_t index); [[nodiscard]] WeightFile model(); + // Whether the checkpoint ships its mask token's embedding (kDraftMaskEmbedding), + // which model.bin then holds after the selector's codebooks. + [[nodiscard]] bool hasMaskEmbedding() const noexcept; private: // The checkpoint and its images, which their writers share. @@ -42,6 +45,12 @@ class DraftCheckpointLoader final { std::shared_ptr planned_; }; +// The BF16 [hidden] embedding of the mask token a draft's checkpoint may ship, +// in a file of its own beside its weights (mask_embedding.safetensors): its +// draft rows' input in place of the target's token table row, which a +// quantized table can lose (DFlashDraft::addMaskRows). +inline constexpr std::string_view kDraftMaskEmbedding = "mask_embedding"; + // The images of a draft checkpoint, in the order DFlashDraft.cpp reads them: // the layers, then model.bin. The checkpoint outlives them: their rows read // its tensors. diff --git a/runtime/model/Runtime.cpp b/runtime/model/Runtime.cpp index c84635c2e..3fd45bf8e 100644 --- a/runtime/model/Runtime.cpp +++ b/runtime/model/Runtime.cpp @@ -2990,6 +2990,8 @@ Runtime::decodeAsync(const BatchPlan &plan, impl_->decodeArena->batchSlice(DecodeTensor::DraftRopeSin, width)); impl_->encodeBatchEmbedding(commandGraph, DecodeTensor::DraftInputTokens, DecodeTensor::DraftHidden0, width); + impl_->draftModel.addMaskRows(commandGraph, impl_->decodeArena->batchSlice(DecodeTensor::DraftHidden0, width), + width); impl_->encodeDraftBatchGraph(commandGraph, entries, {logicalPositions.data(), width}, {maximumRetained.data(), width}, constrained); diff --git a/runtime/ops/Embedding.cpp b/runtime/ops/Embedding.cpp index 163adc956..ba0f86d37 100644 --- a/runtime/ops/Embedding.cpp +++ b/runtime/ops/Embedding.cpp @@ -64,4 +64,15 @@ void Embedding::addVerifyInput(metal::CommandGraph &graph, {1, 1, 1}); } +void Embedding::addDraftMaskRows(metal::CommandGraph &graph, metal::MetalBuffer mask, metal::MetalBuffer rows, + uint32_t width, uint32_t lanes) { + if (!width || !lanes || lanes > SPLASH_MAXIMUM_BATCH_WIDTH) + throw std::invalid_argument("invalid draft mask rows"); + requireBytes(mask, uint64_t{width} * sizeof(uint16_t), "draft mask embedding"); + requireBytes(rows, uint64_t{lanes} * SPLASH_TARGET_VERIFY_ROWS * width * sizeof(uint16_t), "draft input row"); + const DraftMaskRowsParams params{width}; + graph.add("draft_mask_rows", {std::move(mask), std::move(rows)}, params, + {(width + 255) / 256, SPLASH_TARGET_VERIFY_ROWS - 1, lanes}, {256, 1, 1}); +} + } // namespace splash::ops diff --git a/runtime/ops/Embedding.hpp b/runtime/ops/Embedding.hpp index c32c710b5..18a5cc9c9 100644 --- a/runtime/ops/Embedding.hpp +++ b/runtime/ops/Embedding.hpp @@ -21,6 +21,11 @@ class Embedding final { metal::MetalBuffer proposedTokens, metal::MetalBuffer verifyInputTokens, uint32_t vocabulary, uint32_t lanes); + // Each of `lanes` lanes' mask rows of its SPLASH_TARGET_VERIFY_ROWS draft + // input rows of `width` values, all but row 0, its anchor's, set to the bf16 + // row `mask`. + static void addDraftMaskRows(metal::CommandGraph &graph, metal::MetalBuffer mask, + metal::MetalBuffer rows, uint32_t width, uint32_t lanes); }; } // namespace splash::ops