grammar: grammar-constrained decoding on the Zig side over xgrammar 0.2.8's C++ core (DeepSeek-V4.1 stack 5/14)

NOTICE

@@ -6,3 +6,6 @@ TensorFold is licensed under the Apache License, Version 2.0 (LICENSE) from 0.6.
 Releases up to 0.5.0 were MIT-licensed, and code written before 0.6.0 keeps its
 MIT notice in LICENSES/MIT.txt. Third-party code and its licenses are listed in
 THIRD_PARTY_NOTICES.md.
+
+This product includes XGrammar (zig/vendor/xgrammar), Copyright (c) 2024 by XGrammar Contributors,
+licensed under the Apache License, Version 2.0.

THIRD_PARTY_NOTICES.md

@@ -76,6 +76,12 @@ checked bit for bit against ExLlamaV3's dequantization.
 
 Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.py`) follow the cache quantization scheme of [ExLlamaV3](https://github.com/turboderp-org/exllamav3) `-cq 8` and `-cq 4` (MIT License, Copyright (c) 2025 Turboderp, text below): groups of 32, one fp16 absmax scale per group, the group rotated by a 32-point Hadamard, midpoint-grid codes, `compand_a == 0`. 8-bit stores each code as a signed int8 (`q - 128`). 4-bit stores two unsigned codes per byte, low nibble first (the same bits as ExLlamaV3's little-endian packing, a uint8 tensor rather than their uint32 words). Their dequantizer folds another `1/sqrt(32)` into the scale and applies the unnormalized butterfly on the way out; this cache applies the normalized H32 to the query and to the merged output instead, and leaves the stored codes rotated. Scales match their quantizer bit for bit. Reconstructed values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit for bit. The quantizer and the attention dequant are written for TensorFold and checked against an independent reference of that arithmetic.
 
+The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts
+the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante
+(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by
+sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors,
+Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file.
+
 ## Vendored code and weights
 
 `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified `dflash/model_mlx.py` from
@@ -87,6 +93,15 @@ Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.
 The MTP layer TensorFold drafts with comes from that checkpoint's last shard (MIT), converted by
 `families/deepseek_v4/convert.py`.
 
+`zig/vendor/xgrammar/` is the C++ core of [XGrammar](https://github.com/mlc-ai/xgrammar) 0.2.8 (tag `v0.2.8`, commit
+`97787376faee5ed8466cfad57c99855e4ce2f6aa`): the sdist's `cpp/`, `include/` and `3rdparty/` sources, with no local
+patches (build files, tests and bindings left out; see its `README.md`). The Zig grammar module
+(`zig/src/core/grammar/`) builds it with Zig's C++ toolchain and calls it through its own C ABI (`xgr_c.cc`).
+XGrammar is Apache License 2.0, Copyright (c) 2024 by XGrammar Contributors (its `LICENSE` and `NOTICE` sit beside
+it). It bundles [DLPack](https://github.com/dmlc/dlpack)'s header (Apache License 2.0, `3rdparty/dlpack/LICENSE`) and
+[picojson](https://github.com/kazuho/picojson) (BSD-2-Clause, Copyright 2009-2010 Cybozu Labs, Inc. and 2011-2014
+Kazuho Oku; the text is in `3rdparty/picojson/picojson.h`).
+
 TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card states Apache-2.0.
 The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use
 without derivatives. Each checkpoint keeps its own license.

build.zig

@@ -4,6 +4,9 @@ const cuda_build = @import("zig/build/cuda.zig");
 const dist_build = @import("zig/build/dist.zig");
 const qwen27_build = @import("zig/build/qwen27.zig");
 const hip_build = @import("zig/build/hip.zig");
+const tp_build = @import("zig/build/tp.zig");
+const sessions_build = @import("zig/build/sessions.zig");
+const grammar_build = @import("zig/build/grammar.zig");
 
 comptime {
     const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n");
@@ -172,6 +175,9 @@ pub fn build(b: *std.Build) void {
     const warning_run = b.addRunArtifact(warning);
     test_step.dependOn(&warning_run.step);
     b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step);
+    tp_build.add(b, target, optimize, test_step);
+    sessions_build.add(b, test_step);
+    grammar_build.add(b, test_step);
 }
 
 /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4).

tools/zig/lean_check.py

@@ -19,6 +19,7 @@
     "zig/src/families/flashnext/roles_gen.zig",
     "zig/kernels/metal/flashnext/sources_gen.zig",
 )
+VENDORED = "zig/vendor/xgrammar/"  # upstream XGrammar 0.2.8, byte-identical (its README); our wrappers stay checked
 
 
 def generator_output(path: Path) -> bool:
@@ -29,9 +30,14 @@ def generator_output(path: Path) -> bool:
     return rel.startswith("zig/kernels/metal/flashnext/") and HEX_METAL.search(path.name) is not None
 
 
+def vendored(path: Path) -> bool:
+    """A file of the vendored XGrammar copy (exactly the VENDORED prefix)."""
+    return path.relative_to(ROOT).as_posix().startswith(VENDORED)
+
+
 def problems(path: Path) -> list[str]:
     """Every rule break in one file, as 'path:line: reason'."""
-    if generator_output(path):
+    if generator_output(path) or vendored(path):
         return []
     lines = path.read_text(errors="replace").splitlines()
     rel, found = path.relative_to(ROOT), []

tools/zig/lean_check_test.py

@@ -0,0 +1,38 @@
+"""The lean checker's exemptions: the vendored prefix is skipped, a file beside it is still checked."""
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import lean_check  # noqa: E402
+
+BAD = "// one\n// two\n"  # two full-line comments in a row: one finding
+
+
+class Vendored(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.root = Path(self.tmp.name)
+        self.saved, lean_check.ROOT = lean_check.ROOT, self.root
+
+    def tearDown(self):
+        lean_check.ROOT = self.saved
+        self.tmp.cleanup()
+
+    def write(self, rel: str) -> Path:
+        path = self.root / rel
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(BAD)
+        return path
+
+    def test_vendored_file_is_skipped(self):
+        self.assertEqual(lean_check.problems(self.write(lean_check.VENDORED + "cpp/fsm.h")), [])
+
+    def test_file_beside_the_vendored_prefix_is_checked(self):
+        for rel in ("zig/vendor/xgrammar_glue.h", "zig/vendor/other/fsm.h", "zig/src/core/grammar/xgr_c.h"):
+            self.assertEqual(len(lean_check.problems(self.write(rel))), 1, rel)
+
+
+if __name__ == "__main__":
+    unittest.main()

zig/build/grammar.zig

@@ -0,0 +1,59 @@
+//! xgrammar 0.2.8's C++ core (zig/vendor/xgrammar, Apache-2.0) with the C ABI of xgr_c.cc, and the `grammar` module.
+
+const std = @import("std");
+
+const xgr_sources = [_][]const u8{
+    "compiled_grammar.cc",        "config.cc",                "earley_parser.cc",            "fsm.cc",
+    "fsm_builder.cc",             "grammar.cc",               "grammar_builder.cc",          "grammar_compiler.cc",
+    "grammar_functor.cc",         "grammar_matcher.cc",       "grammar_parser.cc",           "grammar_printer.cc",
+    "json_schema_converter.cc",   "lark_converter.cc",        "regex_converter.cc",          "structural_tag.cc",
+    "suffix_automata.cc",         "testing.cc",               "tokenizer_info.cc",           "support/logging.cc",
+    "support/recursion_guard.cc", "converter_ext/cohere.cc",  "converter_ext/deepseek.cc",   "converter_ext/glm.cc",
+    "converter_ext/kimi_k3.cc",   "converter_ext/minimax.cc", "converter_ext/minimax_m3.cc", "converter_ext/qwen.cc",
+};
+
+/// xgrammar's static library with the C ABI, always ReleaseFast (a debug Earley parser is too slow for the tests).
+pub fn library(b: *std.Build, target: std.Build.ResolvedTarget) *std.Build.Step.Compile {
+    const root = "zig/vendor/xgrammar/";
+    const m = b.createModule(.{ .target = target, .optimize = .ReleaseFast, .link_libc = true, .link_libcpp = true });
+    m.addIncludePath(b.path(root ++ "include"));
+    m.addIncludePath(b.path(root ++ "cpp"));
+    m.addSystemIncludePath(b.path(root ++ "3rdparty/picojson"));
+    m.addSystemIncludePath(b.path(root ++ "3rdparty/dlpack/include"));
+    m.addIncludePath(b.path("zig/src/core/grammar"));
+    const flags = [_][]const u8{ "-std=c++17", "-DXGRAMMAR_ENABLE_CPPTRACE=0", "-fexceptions", "-Wno-everything" };
+    m.addCSourceFiles(.{ .root = b.path(root ++ "cpp"), .files = &xgr_sources, .flags = &flags });
+    m.addCSourceFile(.{ .file = b.path("zig/src/core/grammar/xgr_c.cc"), .flags = &flags });
+    return b.addLibrary(.{ .name = "tf_xgrammar", .root_module = m, .linkage = .static });
+}
+
+/// The `grammar` module (zig/src/core/grammar/root.zig) linked with xgrammar.
+pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module {
+    const m = b.createModule(.{ .root_source_file = b.path("zig/src/core/grammar/root.zig"), .target = target, .optimize = optimize, .link_libc = true, .link_libcpp = true });
+    m.addIncludePath(b.path("zig/src/core/grammar"));
+    m.linkLibrary(library(b, target));
+    return m;
+}
+
+pub fn add(b: *std.Build, test_step: *std.Build.Step) void {
+    const host = b.graph.host;
+    const t = b.addTest(.{ .root_module = module(b, host, .debug) });
+    const run = b.addRunArtifact(t);
+    run.setCwd(b.path("."));
+    for ([_][]const u8{ "synthetic.golden.json", "synthetic/tokenizer.json" }) |f| run.addFileInput(b.path(b.fmt("zig/tests/grammar/fixtures/{s}", .{f})));
+    b.step("test-grammar", "Grammar masks / cuts / compiles vs xgrammar 0.2.8's Python package on recorded goldens (no GPU)").dependOn(&run.step);
+    test_step.dependOn(&run.step);
+}
+
+/// The mask kernel as sm_80 PTX (JIT-compiled by the driver), for `@embedFile` as the import `grammar_mask_ptx`.
+pub fn maskPtx(b: *std.Build) std.Build.LazyPath {
+    const nvptx = b.resolveTargetQuery(.{ .cpu_arch = .nvptx64, .os_tag = .cuda, .cpu_model = .{ .explicit = &std.Target.nvptx.cpu.sm_80 } });
+    const obj = b.addObject(.{ .name = "tf_grammar_mask", .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/core/grammar/mask_kernel.zig"), .target = nvptx, .optimize = .ReleaseFast }) });
+    return obj.getEmittedAsm();
+}
+
+/// Adds `grammar` (and the mask kernel's PTX as `grammar_mask_ptx`) to a module that serves grammars.
+pub fn addTo(b: *std.Build, m: *std.Build.Module, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) void {
+    m.addImport("grammar", module(b, target, optimize));
+    m.addAnonymousImport("grammar_mask_ptx", .{ .root_source_file = maskPtx(b) });
+}

zig/build/sessions.zig

@@ -0,0 +1,13 @@
+//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS.
+
+const std = @import("std");
+
+pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module {
+    return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true });
+}
+
+pub fn add(b: *std.Build, test_step: *std.Build.Step) void {
+    const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) }));
+    b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step);
+    test_step.dependOn(&run.step);
+}

zig/build/tp.zig

@@ -0,0 +1,59 @@
+//! Tensor parallelism's build: host tests on any OS, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel.
+
+const std = @import("std");
+
+/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL).
+fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module {
+    const options = b.addOptions();
+    options.addOption(bool, "with_kernels", false);
+    const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    cuda.addOptions("kernel_options", options);
+    cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize }));
+    return cuda;
+}
+
+/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE.
+pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module {
+    const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    m.addImport("cuda", cuda);
+    m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true }));
+    return m;
+}
+
+/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else.
+fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath {
+    const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" });
+    const out = run.addOutputFileArg("tp_mailbox.fatbin");
+    run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu"));
+    return out;
+}
+
+/// `pub const mailbox: []const u8`, the fatbin's bytes or empty.
+fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module {
+    const wf = b.addWriteFiles();
+    if (image) |file| {
+        _ = wf.addCopyFile(file, "tp_mailbox.fatbin");
+        // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing
+        return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") });
+    }
+    return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") });
+}
+
+pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void {
+    const host = b.graph.host;
+    const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) }));
+    b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step);
+    test_step.dependOn(&tests.step);
+    if (target.result.os.tag != .linux) return;
+    const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel");
+    const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed");
+    const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null;
+    const cuda = runtime(b, target, optimize);
+    const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true });
+    runner.addImport("cuda", cuda);
+    runner.addImport("tp", module(b, target, optimize, cuda));
+    runner.addImport("tp_kernel", kernelModule(b, image));
+    const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner });
+    b.installArtifact(exe);
+    b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step);
+}