diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index e9e843113..2a20efd46 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,11 +1853,20 @@ A draft's `config.json` states its layers, intermediate size, rotary base and the target layers whose outputs it reads (`dflash_config.target_layer_ids`, ascending, at most eight); its hidden size, vocabulary and mask token are its target's, and its attention (32 query and 8 KV heads of 128), selector (rank -256, 16 candidates), block of 8 rows and 2048-token window are those the kernels -are built for (`readDraftLayout` in `ModelDescriptor.mm`); the kernels take any -hidden width of whole 16-channel convolution groups. Every family loads and runs -its draft alike: a decode cycle drafts the whole block, the anchor and its seven -proposals, and the target verifies the block's first rows, as many as the cycle +256, 16 candidates) and 2048-token window are those the kernels are built for, +and its block at least their 8 rows, of which a cycle drafts the first (the +Qwen3.5-122B-A10B draft was trained to fill 16; `readDraftLayout` in +`ModelDescriptor.mm`); the kernels take any hidden width of whole 16-channel +convolution groups. A draft's rows after its anchor read the target's token +table row of the mask token, unless its release ships that token's BF16 +embedding beside its weights: `mask_embedding.safetensors`, one tensor +`mask_embedding` of the hidden size, which installation fetches with the weights +(`DFlashDraft::addMaskRows`). A quantized table can lose the row: +Qwen3.5-122B-A10B's is of order 1e-6, below what the fp16 block scales of Q8_0 +and the K-quants hold, so its GGUFs read zeros there, and its draft accepted 2.3 +tokens a cycle on code without the embedding, 4.7 with it. Every family loads +and runs its draft alike: a decode cycle drafts eight rows, the anchor and its +seven proposals, and the target verifies their first rows, as many as the cycle chose for the lane (its verify rows, `ops::DraftSelector::addVerifyRows`), never more than the lane may still retain: the tokens its output budget leaves, at most all eight. A dense target verifies every row up to that limit: its 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/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/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/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/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..7b1658bc2 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" ) @@ -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( @@ -1052,7 +1074,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 +1222,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/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/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/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/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/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/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(); 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/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 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) {} };