DFlash2 alone on Metal starts again without <mtp>/draft_vocab.bin: the packaged default subset (data/draft_vocab.bin) is used, as before the merge; CUDA, HIP and SYCL keep the error; a file named with --mtp-draft-vocab that is missing is still an error (draft_vocab_test)

CMakeLists.txt

@@ -1902,6 +1902,9 @@ if(STRATA_BUILD_TESTS)
   add_executable(window_rule_test src/spec/window_rule_test.cpp)   # which rule picks the window / whose q judges it, by the active draft source
   target_include_directories(window_rule_test PRIVATE include)
   add_test(NAME window_rule_test COMMAND window_rule_test)
+  add_executable(draft_vocab_test src/spec/draft_vocab_test.cpp)   # the DFlash2-alone start's draft-vocabulary choice (a build's packaged default subset; the error elsewhere)
+  target_include_directories(draft_vocab_test PRIVATE include)
+  add_test(NAME draft_vocab_test COMMAND draft_vocab_test)
   add_executable(spec_start_test src/spec/spec_start_test.cpp)
   target_include_directories(spec_start_test PRIVATE include)
   add_test(NAME spec_start_test COMMAND spec_start_test)

include/strata/spec/draft_vocab.hpp

@@ -0,0 +1,35 @@
+// include/strata/spec/draft_vocab.hpp - which draft-vocabulary file a DFlash2-alone start reads its draft head from.
+//
+// The subset is `--mtp-draft-vocab FILE`, else `<--mtp DIR>/draft_vocab.bin`.  A start without it would run the drafter's head over the whole vocabulary without a word
+// (+0.7 ms a round measured on Metal), so a missing file is an error at start.  One exception: a build that ships its own default subset and can name it (the Metal
+// engine: data/draft_vocab.bin beside the build, the file Window::df2_bind has always fallen back to; the Mac's MTP folder is empty by design) passes that path
+// as `packaged_default`; it is used only when the person did not name a file, and only when it exists.  CUDA, HIP and SYCL pass "": a missing file stays an error.
+// `whole_head_ok` is the person's own STRATA_METAL_DRAFT_VOCAB=full (the whole head, asked for on purpose).  Header only, no device types.
+#pragma once
+
+#include <functional>
+#include <string>
+
+namespace strata::spec {
+
+struct Df2Vocab {
+    bool ok = false;
+    std::string path;        ///< the file the drafter binds (the wanted path, even when it is missing and the whole head was asked for)
+    bool packaged = false;   ///< `path` is the build's packaged default, not <mtp>/draft_vocab.bin
+    bool whole_head = false; ///< no file, on purpose: the drafter runs its head over the whole vocabulary
+    std::string why;         ///< ok == false: the line for stderr
+};
+
+inline Df2Vocab df2_vocab_resolve(const std::string& mtp_dir, const std::string& explicit_file, const std::string& packaged_default, bool whole_head_ok,
+                                  const std::function<bool(const std::string&)>& exists) {
+    Df2Vocab r;
+    const std::string wanted = explicit_file.empty() ? mtp_dir + "/draft_vocab.bin" : explicit_file;
+    if (exists(wanted)) { r.ok = true; r.path = wanted; return r; }
+    if (explicit_file.empty() && !packaged_default.empty() && exists(packaged_default)) { r.ok = true; r.path = packaged_default; r.packaged = true; return r; }
+    if (whole_head_ok) { r.ok = true; r.path = wanted; r.whole_head = true; return r; }
+    r.path = wanted;
+    r.why = "--draft dflash2: the draft vocabulary " + wanted + " does not exist (--mtp DIR needs DIR/draft_vocab.bin, or --mtp-draft-vocab FILE); the drafter would run its head over the whole vocabulary";
+    return r;
+}
+
+}  // namespace strata::spec

metal/src/engine/window.cpp

@@ -3237,6 +3237,18 @@ bool mtp_load_pending(Window* w, std::string& err);
 void set_model_shard(const char* shard1) { g_shard1 = shard1 ? shard1 : ""; }
 void set_batch_slots(int n) { g_slot_cap = n > 0 ? std::min(n, 8) : 0; }
 
+// The draft vocabulary the build ships (data/draft_vocab.bin beside the build, else the compiled-in data folder): what a DFlash2-alone start without <mtp>/draft_vocab.bin
+// uses (generate.cpp), and what Window::df2_bind falls back to.  STRATA_METAL_DRAFT_VOCAB names another file, or "full" for the whole head.
+std::string draft_vocab_default() {
+    const char* e = std::getenv("STRATA_METAL_DRAFT_VOCAB");
+    if (e != nullptr && e[0] != 0 && std::string(e) != "full") return e;
+    return strata_metal::resource(nullptr, "data", STRATA_METAL_DATA_DEFAULT) + "/draft_vocab.bin";
+}
+bool draft_vocab_whole_head() {
+    const char* e = std::getenv("STRATA_METAL_DRAFT_VOCAB");
+    return e != nullptr && std::string(e) == "full";
+}
+
 // One decode window per process: the prompt path (Prefill::init) may bring it up before the Verifier does.
 namespace {
 Window* g_the_window = nullptr;

src/program/generate.cpp

@@ -104,6 +104,7 @@
 #include "strata/program/draft_kv_plan.hpp"
 #include "strata/program/message_boundary.hpp"
 #include "strata/spec/draft_policy.hpp"
+#include "strata/spec/draft_vocab.hpp"
 #include "strata/spec/window_rule.hpp"
 #include "strata/spec/draft_source.hpp"
 #include "strata/spec/suffix_drafter.hpp"
@@ -151,6 +152,8 @@ void set_model_shard(const char* shard1);   // metal/src/engine/window.cpp
 void set_batch_slots(int n);   // metal/src/engine/window.cpp: how many --batch slots the window keeps state for
 // batch slot b <-> the main sequence: the running state (unless kv_only) and the first `upto` cells of K/V
 bool window_slot_copy_main(bool into_slot, int slot, int64_t upto, bool kv_only, std::string& err);
+std::string draft_vocab_default();   // metal/src/engine/window.cpp: the packaged default draft subset (data/draft_vocab.bin)
+bool draft_vocab_whole_head();       // STRATA_METAL_DRAFT_VOCAB=full
 }
 #endif
 
@@ -5300,13 +5303,24 @@ int main(int argc, char** argv) {
             std::fprintf(stderr, "strata generate: --draft dflash2 drafts without the MTP layer (--draft-fallback mtp loads it too) and needs --draft-model PACK, --spec 2 or more and --mtp-draft-vocab FILE (or --mtp DIR, for DIR/draft_vocab.bin)\n");
             return 2;
         }
-        df2_vocab_file = o.mtp_draft_vocab.empty() ? o.mtp + "/draft_vocab.bin" : o.mtp_draft_vocab;
-        {   // without the vocabulary subset the DFlash2 would run its head over all of the target's rows (+0.7 ms a round measured on Metal) without a word
-            std::error_code vec;
-            if (!std::filesystem::exists(df2_vocab_file, vec)) {
-                std::fprintf(stderr, "strata generate: --draft dflash2: the draft vocabulary %s does not exist (--mtp DIR needs DIR/draft_vocab.bin, or --mtp-draft-vocab FILE); the drafter would run its head over the whole vocabulary\n", df2_vocab_file.c_str());
+        {   // without the vocabulary subset the DFlash2 would run its head over all of the target's rows (+0.7 ms a round measured on Metal) without a word.
+            // Metal alone has a packaged default subset to fall back to (the Mac's MTP folder is empty by design); everywhere else a missing file is the error.
+            std::string packaged;
+            bool whole_head_ok = false;
+#if defined(STRATA_GPU_METAL)
+            packaged = strata::metal_engine::draft_vocab_default();
+            whole_head_ok = strata::metal_engine::draft_vocab_whole_head();
+#endif
+            const strata::spec::Df2Vocab dv = strata::spec::df2_vocab_resolve(o.mtp, o.mtp_draft_vocab, packaged, whole_head_ok, [](const std::string& p) {
+                std::error_code vec;
+                return std::filesystem::exists(p, vec);
+            });
+            if (!dv.ok) {
+                std::fprintf(stderr, "strata generate: %s\n", dv.why.c_str());
                 return 2;
             }
+            if (dv.packaged) std::fprintf(stderr, "strata generate: --draft dflash2: no draft vocabulary in the MTP folder: using the packaged default subset %s\n", dv.path.c_str());
+            df2_vocab_file = dv.path;
         }
         o.mtp.clear();
         df2_solo = true;

src/spec/draft_vocab_test.cpp

@@ -0,0 +1,37 @@
+// draft_vocab_test: the DFlash2-alone start's draft-vocabulary choice (include/strata/spec/draft_vocab.hpp).
+#include "strata/spec/draft_vocab.hpp"
+
+#include <cstdio>
+#include <set>
+#include <string>
+
+using strata::spec::df2_vocab_resolve;
+
+static int g_fail = 0;
+#define CHECK(c) do { if (!(c)) { std::printf("FAIL line %d: %s\n", __LINE__, #c); ++g_fail; } } while (0)
+
+int main() {
+    const std::set<std::string> files = {"/mtp/with/draft_vocab.bin", "/given.bin", "/pkg/data/draft_vocab.bin"};
+    auto exists = [&](const std::string& p) { return files.count(p) > 0; };
+
+    // the MTP folder has the subset: it is used, packaged default or not
+    for (const char* pkg : {"", "/pkg/data/draft_vocab.bin"}) {
+        auto r = df2_vocab_resolve("/mtp/with", "", pkg, false, exists);
+        CHECK(r.ok && r.path == "/mtp/with/draft_vocab.bin" && !r.packaged && !r.whole_head);
+    }
+    // a file the person named wins over the folder's
+    { auto r = df2_vocab_resolve("/mtp/with", "/given.bin", "/pkg/data/draft_vocab.bin", false, exists); CHECK(r.ok && r.path == "/given.bin" && !r.packaged); }
+    // the Mac: the MTP folder has none, the build ships its default: that is used (what Metal did before the merge)
+    { auto r = df2_vocab_resolve("/mtp/empty", "", "/pkg/data/draft_vocab.bin", false, exists); CHECK(r.ok && r.packaged && r.path == "/pkg/data/draft_vocab.bin"); }
+    // CUDA / HIP / SYCL (no packaged default named): a missing file is the error it was
+    { auto r = df2_vocab_resolve("/mtp/empty", "", "", false, exists); CHECK(!r.ok && r.why.find("/mtp/empty/draft_vocab.bin") != std::string::npos); }
+    // no subset can be found at all (the build's default is missing too): the error, on every backend
+    { auto r = df2_vocab_resolve("/mtp/empty", "", "/nowhere/draft_vocab.bin", false, exists); CHECK(!r.ok && !r.packaged); }
+    // a file the person named that does not exist is their mistake: no silent fallback to the default
+    { auto r = df2_vocab_resolve("/mtp/with", "/typo.bin", "/pkg/data/draft_vocab.bin", false, exists); CHECK(!r.ok && r.why.find("/typo.bin") != std::string::npos); }
+    // the whole head asked for on purpose (STRATA_METAL_DRAFT_VOCAB=full) is not an error
+    { auto r = df2_vocab_resolve("/mtp/empty", "", "", true, exists); CHECK(r.ok && r.whole_head); }
+
+    if (g_fail == 0) std::printf("draft_vocab_test: all checks passed\n");
+    return g_fail == 0 ? 0 : 1;
+}