diff --git a/Cargo.lock b/Cargo.lock index 3dff4a6d..da982c33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1272,6 +1272,7 @@ dependencies = [ "half", "memmap2", "omni-cuda", + "omni-modernbert", "omni-runtime", "safetensors 0.6.2", "serde", @@ -1282,6 +1283,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-modernbert" +version = "0.1.0" +dependencies = [ + "anyhow", + "serde_json", + "sha2 0.10.9", + "tempfile", +] + [[package]] name = "omni-omnijev-native" version = "0.2.0" @@ -1327,6 +1338,8 @@ dependencies = [ "safetensors 0.8.0", "serde", "serde_json", + "sha2 0.10.9", + "tempfile", ] [[package]] @@ -1353,6 +1366,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-vela-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "half", + "memmap2", + "omni-cuda", + "omni-modernbert", + "safetensors 0.8.0", + "serde_json", + "tempfile", +] + [[package]] name = "once_cell" version = "1.21.4" diff --git a/Cargo.toml b/Cargo.toml index 7910b3dd..766a6c5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/omnijev/native", "src/models/laya", "src/backends/cuda", "src/models/decider/native", "src/models/jemm/native", "src/models/startlux/native", "src/models/clef/native"] +members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/omnijev/native", "src/models/laya", "src/models/modernbert", "src/models/vela/native", "src/backends/cuda", "src/models/decider/native", "src/models/jemm/native", "src/models/startlux/native", "src/models/clef/native"] resolver = "3" diff --git a/README.md b/README.md index 1502a8fb..fe69eb3a 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,8 @@ repository root. | [`docs/`](docs/) | Project documentation and architecture assets. | The frontend, native runtime, the native workers, their shared Qwen3.5/3.8 -prefill implementation and the Laya CUDA backend are Cargo workspace members. +prefill implementation, the ModernBERT family structure crate, the Vela native +profile crate and the Laya CUDA backend are Cargo workspace members. Each model and backend documents its implementation and process boundary. ## Getting Started diff --git a/recipe/vela/compare_reference.py b/recipe/vela/compare_reference.py index 3588e06a..9cbe4c26 100644 --- a/recipe/vela/compare_reference.py +++ b/recipe/vela/compare_reference.py @@ -16,11 +16,14 @@ response construction to the same pinned code — so the criterion is reference parity: exact match, including JSON key order, with rejections agreeing on status and body (``{"detail": [{"loc", "msg", -"type"}]}``). ``--export`` additionally freezes the oracle fixture — -request, SystemOne response and the predict()-level answers — for the -future native encoder slices, and writes it only after a passing, -non-rejected comparison. Requires the pinned package, the reference -dependencies, and fastapi+uvicorn+httpx; see recipe/vela/reference.md. +"type"}]}``). The run also compares the worker's verbatim endpoint +bodies (``/health``, ``/ready``, ``/v1/models`` and the unauthorized +401) against the live reference app. ``--export`` additionally freezes +the oracle fixture — request, SystemOne response, the predict()-level +answers and the generating run's provenance — for the future native +encoder slices, and writes it only after a passing, non-rejected +comparison. Requires the pinned package, the reference dependencies, +and fastapi+uvicorn+httpx; see recipe/vela/reference.md. The default fixture set is the recipe's four example requests plus five synthetic rejections: an invalid span question (whole-request @@ -237,17 +240,56 @@ def compare_case(name, worker, reference): return lines, False +def endpoint_parity(client, gated, worker_endpoints, unauthorized_body, probe_body): + """Compare the worker's verbatim endpoint bodies with the live reference app. + + ``worker_endpoints`` maps the reference app's GET routes to the exact + bodies the worker serves; ``gated`` is a reference app built with an + API key, so an unauthorized POST must answer the worker's 401 body. + A drift in either direction fails the run. Returns report lines and + the failure count. + """ + + lines = [] + failures = 0 + for path, expected in worker_endpoints: + response = client.get(path) + observed = response.json() + if response.status_code != 200 or observed != expected: + failures += 1 + lines.append( + f"endpoint {path}: FAIL: HTTP {response.status_code} " + f"{json.dumps(observed, ensure_ascii=False)[:160]} != the worker's " + f"{json.dumps(expected, ensure_ascii=False)[:160]}" + ) + else: + lines.append(f"endpoint {path}: HTTP 200, exact body match: PASS") + response = gated.post("/v1/systemone", json=probe_body) + observed = response.json() + if response.status_code != 401 or observed != unauthorized_body: + failures += 1 + lines.append( + f"unauthorized POST: FAIL: HTTP {response.status_code} " + f"{json.dumps(observed, ensure_ascii=False)[:160]} != the worker's " + f"{json.dumps(unauthorized_body, ensure_ascii=False)[:160]}" + ) + else: + lines.append("unauthorized POST: 401, exact body match: PASS") + return lines, failures + + def export_allowed(ok, worker): """The oracle freezes served 200 responses, never a rejection.""" return ok and worker[0] == 200 -def export_oracle(name, body, response, predict, path): +def export_oracle(name, body, response, predict, path, source): """Freeze one passing case for the future native encoder slices.""" oracle = { "case": name, + "source": source, "request": body, "system_one": response, "predict": predict, @@ -276,8 +318,25 @@ def main(argv=None): finally: sys.dont_write_bytecode = previous_dont_write_bytecode + from frontend.vela import HEALTH_BODY, MODELS_ENTRY, READY_BODY, UNAUTHORIZED_BODY + app = vela2_serve.create_app(executor.engine, model_name=executor.model_name, api_key=None) client = TestClient(app) + gated = TestClient( + vela2_serve.create_app(executor.engine, model_name=executor.model_name, api_key="parity-key") + ) + source = { + "model": executor.model_name, + "device": args.device, + "generator": "recipe/vela/compare_reference.py --export", + } + # A well-formed probe body: any answer other than the worker's 401 + # means the reference's key gate drifted, not the probe itself. + unauthorized_probe = { + "model": executor.model_name, + "state": "short", + "questions": {"q": {"type": "noul", "instructions": "any?"}}, + } served = executor.model_name # The hybrid member's oracle fixtures get their own prefix so both # members' frozen baselines can live in one directory. @@ -286,6 +345,21 @@ def main(argv=None): exported = 0 failed = 0 try: + endpoint_lines, endpoint_failures = endpoint_parity( + client, + gated, + [ + ("/health", HEALTH_BODY), + ("/ready", READY_BODY), + ("/v1/models", {"models": [{**MODELS_ENTRY, "name": executor.model_name}]}), + ], + UNAUTHORIZED_BODY, + unauthorized_probe, + ) + print("endpoint parity:") + print("\n".join(endpoint_lines)) + print() + failed += endpoint_failures for name, body in load_cases(args.fixture, served): raw = json.dumps(body).encode("utf-8") worker = worker_response(raw, executor, decide) @@ -299,7 +373,7 @@ def main(argv=None): target = Path(args.export) if target.is_dir(): target = target / f"{prefix}{name}.json" - export_oracle(name, body, worker[1], executor.engine.predict(req), target) + export_oracle(name, body, worker[1], executor.engine.predict(req), target, source) exported += 1 print() finally: diff --git a/recipe/vela/reference.md b/recipe/vela/reference.md index 559a7afb..03d13280 100644 --- a/recipe/vela/reference.md +++ b/recipe/vela/reference.md @@ -35,8 +35,11 @@ Scope: calibrated readouts; zero generated tokens. Out of scope: the larger hybrid members (4B/9B), the ONNX backend of -the 0.3B, any native Rust/CUDA executor, and `noul_calibration` (off -by default; documented as a reference-server feature only). +the 0.3B, native execution for the reference workers (the ModernBERT +profile's own `cuda`-gated route under +[`native/`](../../src/models/vela/native/README.md) is a separate +slice), and `noul_calibration` (off by default; documented as a +reference-server feature only). ## Architecture and ownership @@ -256,11 +259,15 @@ share one loaded engine: the pinned `frontend.vela.decide` pipeline. The criterion is reference parity: **exact match** including JSON key order for 200s, and agreement on status and body for rejections (needs `pip install -fastapi uvicorn httpx` on top of the requirements file). +fastapi uvicorn httpx` on top of the requirements file). The run also +compares the worker's verbatim endpoint bodies — `/health`, `/ready`, +`/v1/models` and the unauthorized 401 — against the live reference app, +so a drift in either direction fails the run. `--export tests/vela/fixtures/oracle-.json` freezes each passing -200 case — request, SystemOne response and predict()-level answers — as -the frozen comparison baseline, and writes it only after a passing -comparison. +200 case — request, SystemOne response, predict()-level answers and the +generating run's provenance — as +the frozen comparison baseline for the future native encoder +slices, and writes it only after a passing comparison. ~~~sh PYTHONPATH=src .venv/bin/python recipe/vela/compare_reference.py \ diff --git a/recipe/vela/test_compare_reference.py b/recipe/vela/test_compare_reference.py index f76fb69d..bb286616 100644 --- a/recipe/vela/test_compare_reference.py +++ b/recipe/vela/test_compare_reference.py @@ -1,11 +1,12 @@ #!/usr/bin/env python3 """Regression tests for recipe/vela/compare_reference.py. -Checks the pure comparison, case-loading and export helpers without the -pinned package, weights, torch or fastapi: exact-match verdicts -including key order, rejection agreement, the default fixture set's -shape (the invalid-span 422 and over-limit 413 synthetic cases), the -oracle export gate, and the oracle export structure. +Checks the pure comparison, case-loading, endpoint-parity and export +helpers without the pinned package, weights, torch or fastapi: +exact-match verdicts including key order, rejection agreement, the +default fixture set's shape (the invalid-span 422 and over-limit 413 +synthetic cases), the oracle export gate, and the oracle export +structure. python recipe/vela/test_compare_reference.py """ @@ -139,18 +140,73 @@ def check(name: str, condition: bool, detail: str = "") -> None: json.dumps(over)[:120], ) - # 4. The oracle export freezes the full native-comparison surface. + # 4. Endpoint parity fails only on a real body or status drift. + class FakeResponse: + def __init__(self, status_code, body): + self.status_code = status_code + self._body = body + + def json(self): + return self._body + + class FakeClient: + def __init__(self, gets, post=None): + self._gets = gets + self._post = post + + def get(self, path): + return self._gets[path] + + def post(self, path, json=None): + return self._post + + health = {"status": "ok"} + models = {"models": [{"name": tool.MODEL_NAME}]} + endpoints = [("/health", health), ("/v1/models", models)] + unauthorized = {"detail": "Missing or invalid API key"} + probe = {"model": tool.MODEL_NAME, "state": "s", "questions": {}} + + same = FakeClient( + {"/health": FakeResponse(200, dict(health)), "/v1/models": FakeResponse(200, {"models": [{"name": tool.MODEL_NAME}]})}, + post=FakeResponse(401, dict(unauthorized)), + ) + lines, failures = tool.endpoint_parity(same, same, endpoints, unauthorized, probe) + check( + "endpoint parity passes on exact bodies", + failures == 0 and lines[0].endswith("PASS") and lines[1].endswith("PASS") + and lines[2].endswith("PASS"), + "\n".join(lines), + ) + + drifted = FakeClient( + {"/health": FakeResponse(200, {"status": "degraded"}), "/v1/models": FakeResponse(200, {"models": [{"name": "other"}]})}, + post=FakeResponse(401, {"detail": "different"}), + ) + lines, failures = tool.endpoint_parity(drifted, drifted, endpoints, unauthorized, probe) + check("endpoint body drift fails", failures == 3, "\n".join(lines)) + + wrong_status = FakeClient( + {"/health": FakeResponse(503, dict(health)), "/v1/models": FakeResponse(200, {"models": [{"name": tool.MODEL_NAME}]})}, + post=FakeResponse(401, dict(unauthorized)), + ) + lines, failures = tool.endpoint_parity(wrong_status, wrong_status, endpoints, unauthorized, probe) + check("endpoint status drift fails", failures == 1, "\n".join(lines)) + + # 5. The oracle export freezes the full native-comparison surface. body = {"model": tool.MODEL_NAME, "state": "s", "questions": {"q": {"type": "noul", "instructions": "x"}}} + source = {"model": tool.MODEL_NAME, "device": "cpu", + "generator": "recipe/vela/compare_reference.py --export"} with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "oracle.json" oracle = tool.export_oracle( - "case", body, response(), {"answers": [{"id": "q", "type": "choice"}]}, path + "case", body, response(), {"answers": [{"id": "q", "type": "choice"}]}, path, source ) reread = json.loads(path.read_text(encoding="utf-8")) check( "oracle export round-trips", oracle == reread + and oracle["source"] == source and oracle["request"] == body and oracle["system_one"]["answers"]["q"]["noul"] == 0.75 and oracle["predict"]["answers"][0]["id"] == "q", diff --git a/src/backends/cuda/kernels/model_ops.cu b/src/backends/cuda/kernels/model_ops.cu index fad23b1d..e73bf3e4 100644 --- a/src/backends/cuda/kernels/model_ops.cu +++ b/src/backends/cuda/kernels/model_ops.cu @@ -22,11 +22,11 @@ __device__ Stats combine(Stats b, Stats a) { return {0, 0, 0}; } -template +template __device__ Stats stats(Load load, float* buf) { int lane = threadIdx.x, warp = threadIdx.y, t = lane + warp * 32; Stats wd{0, 0, 0}; - for (int i = t; i < 256; i += 128) { + for (int i = t; i < D / 4; i += 128) { #pragma unroll for (int j = 0; j < 4; j++) { float v = load(4 * i + j), delta = v - wd.mean, count = wd.count + 1.f, @@ -56,7 +56,7 @@ __device__ Stats stats(Load load, float* buf) { } if (lane == 0 && warp == 0) { buf[0] = wd.mean; - buf[1] = wd.var / 1024.f; + buf[1] = wd.var / float(D); } __syncthreads(); return {buf[0], buf[1], 0}; @@ -76,20 +76,21 @@ struct FloatLoad { } }; +template __global__ void embed_norm(const int64_t* ids, const half* w, const float* gamma, float* x, BF* y) { int r = blockIdx.x, t = threadIdx.x + threadIdx.y * 32; __shared__ float buf[6]; - HalfLoad load{w + ids[r] * 1024}; - Stats wd = stats(load, buf); + HalfLoad load{w + ids[r] * D}; + Stats wd = stats(load, buf); float inv = rsqrtf(wd.var + 1e-5f); - for (int i = t; i < 256; i += 128) { + for (int i = t; i < D / 4; i += 128) { #pragma unroll for (int k = 0; k < 4; k++) { int j = 4 * i + k; float z = gamma[j] * (inv * (load(j) - wd.mean)); - x[r * 1024 + j] = z; - y[r * 1024 + j] = __float2bfloat16_rn(z); + x[r * D + j] = z; + y[r * D + j] = __float2bfloat16_rn(z); } } } @@ -113,7 +114,7 @@ __global__ void gather_norm(const float* h, const int32_t* indices, int r = blockIdx.x, t = threadIdx.x + threadIdx.y * 32; __shared__ float buf[6]; FloatLoad load{h + indices[r] * 1024}; - Stats wd = stats(load, buf); + Stats wd = stats(load, buf); float inv = rsqrtf(wd.var + 1e-5f); for (int i = t; i < 256; i += 128) { #pragma unroll @@ -178,7 +179,14 @@ __global__ void action_features(const float* h, const BF* logits, extern "C" { int laya_embed(void** p, int B, int L, int M, cudaStream_t s) { - embed_norm<<>>( + embed_norm<1024><<>>( + (int64_t*)p[0], (half*)p[1], (float*)p[2], (float*)p[3], (BF*)p[4]); + return cudaGetLastError(); +} + +// The 768-hidden ModernBERT profile (Vela 2.0 0.3B) on the same glue op. +int laya_embed_h768(void** p, int B, int L, int M, cudaStream_t s) { + embed_norm<768><<>>( (int64_t*)p[0], (half*)p[1], (float*)p[2], (float*)p[3], (BF*)p[4]); return cudaGetLastError(); } diff --git a/src/backends/cuda/laya.backend.json b/src/backends/cuda/laya.backend.json index 3d0452c2..290b83a3 100644 --- a/src/backends/cuda/laya.backend.json +++ b/src/backends/cuda/laya.backend.json @@ -1,6 +1,6 @@ { "name": "laya", - "abi_version": 1, + "abi_version": 2, "status": "experimental", "sources": [ "kernels/runtime.cu", diff --git a/src/backends/cuda/src/lib.rs b/src/backends/cuda/src/lib.rs index a72f008e..828d05de 100644 --- a/src/backends/cuda/src/lib.rs +++ b/src/backends/cuda/src/lib.rs @@ -156,7 +156,7 @@ impl Cuda { /// Buffers must belong to this context and stay alive until synchronization or graph destruction. pub unsafe fn launch(&self, name: &str, args: &[Ptr], b: usize, l: usize) -> Result<()> { ensure!( - b > 0 && b <= 16 && l > 0 && l <= 512 && l.is_multiple_of(16), + b > 0 && b <= 16 && l > 0 && l <= 8192 && l.is_multiple_of(16), "invalid CUDA shape" ); let k = self @@ -354,7 +354,7 @@ impl Kernel { /// or destruction of any graph that captures the launch. pub unsafe fn launch(&self, args: &[Ptr], b: usize, l: usize) -> Result<()> { ensure!( - b > 0 && b <= 16 && l > 0 && l <= 512 && l.is_multiple_of(16), + b > 0 && b <= 16 && l > 0 && l <= 8192 && l.is_multiple_of(16), "invalid CUDA shape" ); self.ctx.check(unsafe { diff --git a/src/backends/cuda/tools/build.py b/src/backends/cuda/tools/build.py index 10424d27..4dbbf23f 100644 --- a/src/backends/cuda/tools/build.py +++ b/src/backends/cuda/tools/build.py @@ -59,7 +59,7 @@ def main(): (args.bundle / "build-command.json").write_text(json.dumps(commands, indent=2)) build_manifest = { - "abi": 1, + "abi": 2, "arch": "sm_90a", "nvcc": subprocess.check_output([nvcc, "--version"], text=True), "sources": sources, diff --git a/src/backends/cuda/tools/export.py b/src/backends/cuda/tools/export.py index b1c0460d..19be95f2 100644 --- a/src/backends/cuda/tools/export.py +++ b/src/backends/cuda/tools/export.py @@ -62,22 +62,27 @@ def integer(value): return int(value) -def patch_geglu_static_registers(body): +def patch_geglu_static_registers(body, pinned=True): """Remove fixed register reservations only from the accepted BN32 lowering.""" - expected = "1ed96990542c07b58dca419dafcc45c6163bb9b6fff04ffca5ecedb2a34e921c" - if hashlib.sha256(body.encode()).hexdigest() != expected: - raise ValueError("long GEGLU BN32 lowering changed") + if pinned: + expected = "1ed96990542c07b58dca419dafcc45c6163bb9b6fff04ffca5ecedb2a34e921c" + if hashlib.sha256(body.encode()).hexdigest() != expected: + raise ValueError("long GEGLU BN32 lowering changed") for call in ( " tl::warpgroup_reg_dealloc<24>();\n", " tl::warpgroup_reg_alloc<240>();\n", ): - if body.count(call) != 1: + count = body.count(call) + if pinned and count != 1: raise ValueError("long GEGLU register calls changed") - body = body.replace(call, "", 1) + if count == 1: + body = body.replace(call, "", 1) + elif count != 0: + raise ValueError("unexpected GEGLU register call count") return body -def export(name, kernel): +def export(name, kernel, l_max=512): source = kernel.get_kernel_source() signature = re.search(r"void main_kernel\((.*?)\);", source, re.S)[1] params = [param.strip() for param in signature.split(",")] @@ -160,8 +165,8 @@ def export(name, kernel): body = source[source.index('extern "C" __global__') :].replace( "main_kernel", symbol ) - if name == "geglu": - body = patch_geglu_static_registers(body) + if name.startswith("geglu"): + body = patch_geglu_static_registers(body, pinned=name == "geglu") # TileLang reuses Q_s for O_s. Its generated wait is inside the key loop, # so zero-key rows can overwrite Q_s while the asynchronous Q load is live. # Wait before entering producer/consumer branches, including zero iterations. @@ -197,7 +202,7 @@ def export(name, kernel): dispatch = " if(B==1 && M<=64)return laya_geglu_short(p,B,L,M,stream);\n" wrapper_name = {"qkv": "qkv_full", "out": "out_full"}.get(name, name) wrapper = f"""extern "C" int laya_{wrapper_name}(void** p,int B,int L,int M,cudaStream_t stream) {{ - if(B<1 || B>16 || L<16 || L>512 || L%16 || M!=B*L)return -1; + if(B<1 || B>16 || L<16 || L>{l_max} || L%16 || M!=B*L)return -1; {dispatch}{chr(10).join(bindings)} {chr(10).join(descriptors)} {symbol}<<>>({','.join(args)}); @@ -243,6 +248,19 @@ def main(): "qkv_short": kernels.gemm_kernel(3072, 1024, bn=64), "attn_full": kernels.attn_kernel(None, None, 16, 64), } + # The second ModernBERT profile (Vela 2.0 0.3B, hidden 768, 12 heads, + # 1152 intermediate, requests packed to 8192 tokens). Its symbols carry + # an _h768 suffix so both profiles coexist in one library. + vela_exports = { + "qkv_h768": kernels.gemm_kernel(2304, 768), + "out_h768": kernels.gemm_kernel(768, 768), + "geglu_h768": kernels.gemm_geglu_kernel(1152, 768, bn=32), + "down_h768": kernels.gemm_kernel(768, 1152, bm=64, bn=64), + "addln_h768": kernels.add_ln_kernel(768), + "attn_full_h768": kernels.attn_kernel(None, None, 12, 64), + "attn_local_h768": kernels.attn_kernel(None, None, 12, 64, window=64), + "rope_h768": rope.build(12, 64, 1, 8), + } if not args.probe_only: exports.update( { @@ -271,8 +289,11 @@ def main(): bodies = [] inits = [] metadata = [] - for name, kernel in exports.items(): - preamble, body, init, details = export(name, kernel) + entries = [(name, kernel, 512) for name, kernel in exports.items()] + if not args.probe_only: + entries += [(name, kernel, 8192) for name, kernel in vela_exports.items()] + for name, kernel, l_max in entries: + preamble, body, init, details = export(name, kernel, l_max) preambles.append(preamble) bodies.append(body) inits.append(init) diff --git a/src/backends/cuda/tools/export_tables.py b/src/backends/cuda/tools/export_tables.py index 95bbb1b9..1fe85ec6 100644 --- a/src/backends/cuda/tools/export_tables.py +++ b/src/backends/cuda/tools/export_tables.py @@ -47,7 +47,7 @@ def main(): files[name] = hashlib.sha256(data).hexdigest() metadata = { - "abi": 1, + "abi": 2, "laya": "0.3.20", "hidden_size": 1024, "head_dim": 64, diff --git a/src/frontend/README.md b/src/frontend/README.md index 4acb3a5e..e52599cb 100644 --- a/src/frontend/README.md +++ b/src/frontend/README.md @@ -38,7 +38,9 @@ The health timeout and response limit must be positive integers. ## HTTP interface - `POST /v1/systemone` forwards the `model`, `state` and `questions` envelope - unchanged. Workers return `choice`, `score` or `noul` decisions; see the + unchanged. Workers return `choice`, `score`, `noul`, `set` and `span` + decisions (the Vela 2.0 worker adds the `set` and `span` kinds with + their Noul views); see the [Jev API reference](https://docs.typesafe.ai/api). - `GET /health` returns the worker's health response, including unhealthy status codes. - Authorization and other end-to-end headers are forwarded. Response status, diff --git a/src/frontend/vela.py b/src/frontend/vela.py index 78f46631..72bae4b5 100644 --- a/src/frontend/vela.py +++ b/src/frontend/vela.py @@ -54,6 +54,10 @@ ), } UNAUTHORIZED_BODY = {"detail": "Missing or invalid API key"} +# Upstream verbatim endpoint bodies; compare_reference.py pins them +# against the live reference app. +HEALTH_BODY = {"status": "ok"} +READY_BODY = {"ready": True} def detect_member(package: str | Path) -> tuple[type, str]: @@ -191,9 +195,9 @@ def do_GET(self): if path == "/health": # Upstream verbatim; the worker binds the port only after # loading and warmup, so answering at all means ready. - self.send_json(200, {"status": "ok"}) + self.send_json(200, HEALTH_BODY) elif path == "/ready": - self.send_json(200, {"ready": True}) + self.send_json(200, READY_BODY) elif path == "/v1/models": if self._unauthorized(): return diff --git a/src/models/laya/Cargo.toml b/src/models/laya/Cargo.toml index 259d3bae..0a449982 100644 --- a/src/models/laya/Cargo.toml +++ b/src/models/laya/Cargo.toml @@ -14,6 +14,7 @@ serde_json = { version = "1", features = ["preserve_order", "arbitrary_precision tokenizers = { version = "0.23.2", default-features = false, features = ["fancy-regex"] } sha2 = "0.10" omni-cuda = { path = "../../backends/cuda", optional = true } +omni-modernbert = { path = "../modernbert" } omni-runtime = { path = "../../runtime", optional = true } tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "signal", "sync"], optional = true } axum = { version = "0.8", optional = true } @@ -48,3 +49,11 @@ path = "../../../tests/laya/decision.rs" [[test]] name = "batch" path = "../../../tests/laya/batch.rs" + +[[test]] +name = "plan" +path = "../../../tests/laya/plan.rs" + +[[test]] +name = "profile" +path = "../../../tests/laya/profile.rs" diff --git a/src/models/laya/src/artifacts.rs b/src/models/laya/src/artifacts.rs index 70185dc6..ba3d101f 100644 --- a/src/models/laya/src/artifacts.rs +++ b/src/models/laya/src/artifacts.rs @@ -32,12 +32,12 @@ pub fn validate_bundle(checkpoint: &Path, bundle: &Path) -> Result<()> { let build: serde_json::Value = serde_json::from_slice(&fs::read(bundle.join("build-manifest.json"))?)?; ensure!( - tables["abi"] == 1 + tables["abi"] == 2 && tables["laya"] == "0.3.20" && tables["hidden_size"] == 1024 && tables["head_dim"] == 64 && tables["max_len"] == 512 - && build["abi"] == 1 + && build["abi"] == 2 && build["arch"] == "sm_90a", "unsupported CUDA bundle" ); diff --git a/src/models/laya/src/config.rs b/src/models/laya/src/config.rs index 8b85c393..1b9d05da 100644 --- a/src/models/laya/src/config.rs +++ b/src/models/laya/src/config.rs @@ -1,7 +1,11 @@ use anyhow::{Context, Result, ensure}; +use omni_modernbert::gate::BackendLimits; +use omni_modernbert::structure::ModernBertStructure; use serde::Deserialize; use std::{collections::HashMap, fs, path::Path}; +use crate::profile::{LayaBundleAdapter, laya_0_3_20}; + #[derive(Debug, Deserialize)] pub struct AgentConfig { pub max_len: usize, @@ -24,59 +28,46 @@ pub struct EncoderConfig { pub rope_parameters: serde_json::Value, } +impl From for EncoderConfig { + fn from(s: ModernBertStructure) -> Self { + Self { + hidden_size: s.hidden_size, + intermediate_size: s.intermediate_size, + num_attention_heads: s.num_attention_heads, + num_hidden_layers: s.num_hidden_layers, + vocab_size: s.vocab_size, + norm_eps: s.norm_eps, + local_attention: s.local_window, + layer_types: s.schedule.iter().map(|k| k.as_str().to_string()).collect(), + rope_parameters: serde_json::json!({ + "full_attention": {"rope_type": "default", "rope_theta": s.rope_full.theta}, + "sliding_attention": {"rope_type": "default", "rope_theta": s.rope_sliding.theta}, + }), + } + } +} + pub struct Config { pub agent: AgentConfig, pub encoder: EncoderConfig, + /// The admitted family structure; geometry and plans derive from it. + pub structure: ModernBertStructure, + /// The registered profile's request envelope, admitted with the + /// structure; the workspace bounds come from here, not inline + /// literals. + pub limits: BackendLimits, } impl Config { pub fn load(dir: &Path) -> Result { let read = |name| fs::read(dir.join(name)).with_context(|| format!("read {name}")); let agent: AgentConfig = serde_json::from_slice(&read("rl_agent_config.json")?)?; - let encoder_bytes = read("encoder/config.json")?; - let raw: serde_json::Value = serde_json::from_slice(&encoder_bytes)?; - ensure!( - raw["model_type"] == "modernbert" - && raw["hidden_activation"] == "gelu" - && raw["attention_bias"] == false - && raw["mlp_bias"] == false - && raw["norm_bias"] == false, - "unsupported encoder activation, bias or model type" - ); - let encoder: EncoderConfig = serde_json::from_slice(&encoder_bytes)?; + let (structure, _meta) = LayaBundleAdapter::read(dir)?; + let profile = laya_0_3_20(); + profile.admit(&structure)?; ensure!( agent.max_len == 512 && agent.head_max_len == 192 && agent.head_layers == 2, "native Laya supports max_len=512, head_max_len=192, head_layers=2" ); - ensure!( - encoder.hidden_size == 1024 - && encoder.intermediate_size == 2624 - && encoder.num_attention_heads == 16 - && encoder.num_hidden_layers == 28 - && encoder.vocab_size == 50368 - && encoder.local_attention == 128 - && encoder.norm_eps == 1e-5, - "unsupported encoder configuration" - ); - let expected: Vec<_> = (0..28) - .map(|i| { - if i % 3 == 0 { - "full_attention" - } else { - "sliding_attention" - } - }) - .collect(); - ensure!( - encoder.layer_types == expected, - "unsupported attention schedule" - ); - for (kind, theta) in [("full_attention", 160000.0), ("sliding_attention", 10000.0)] { - let r = &encoder.rope_parameters[kind]; - ensure!( - r["rope_type"] == "default" && r["rope_theta"].as_f64() == Some(theta), - "unsupported RoPE configuration" - ); - } ensure!( agent.temperature.len() == 3 && agent @@ -86,6 +77,11 @@ impl Config { .all(|t| t.is_finite() && *t > 0.0), "invalid temperatures" ); - Ok(Self { agent, encoder }) + Ok(Self { + encoder: EncoderConfig::from(structure.clone()), + agent, + structure, + limits: profile.limits, + }) } } diff --git a/src/models/laya/src/lib.rs b/src/models/laya/src/lib.rs index 7db9d9ad..b0f42e35 100644 --- a/src/models/laya/src/lib.rs +++ b/src/models/laya/src/lib.rs @@ -2,8 +2,10 @@ pub mod artifacts; pub mod config; pub mod decision; pub mod packing; +pub mod plan; pub mod preprocess; pub mod processing; +pub mod profile; pub mod weights; #[cfg(feature = "serve")] diff --git a/src/models/laya/src/model.rs b/src/models/laya/src/model.rs index 54f6aa67..8212c525 100644 --- a/src/models/laya/src/model.rs +++ b/src/models/laya/src/model.rs @@ -1,8 +1,14 @@ //! Single-GPU Laya executor. All CUDA buffers and graphs remain on the owning worker thread. -use crate::{config::Config, packing::Batch, weights::Weights}; +use crate::config::Config; +use crate::packing::Batch; +use crate::plan::{PlannedArg, PlannedStep, Scratch, plan_encoder}; +use crate::weights::Weights; use anyhow::{Context, Result, ensure}; use half::bf16; use omni_cuda::{Buffer, Cuda, Graph, Kernel, Ptr}; +use omni_modernbert::gate::BackendLimits; +use omni_modernbert::geometry::EncoderScratch; +use omni_modernbert::structure::ModernBertStructure; use std::{ collections::{HashMap, VecDeque}, fs, @@ -134,22 +140,25 @@ struct Workspace { actions: Buffer, } impl Workspace { - fn sizes(b: usize, l: usize) -> Result<[usize; 17]> { - ensure!( - b.is_power_of_two() && b <= 16 && (16..=512).contains(&l) && l.is_multiple_of(16), - "invalid workspace shape" - ); + fn sizes( + b: usize, + l: usize, + structure: &ModernBertStructure, + limits: &BackendLimits, + ) -> Result<[usize; 17]> { + limits.check_request(b, l)?; let m = b * l; + let encoder = EncoderScratch::new(b, l, structure); Ok([ - m * 8, - b * 4, + encoder.ids, + encoder.lens, b * 8, - m * D * 4, - m * D * 2, - m * D * 6, - m * D * 2, - m * 2624 * 2, - m * 4096 * 2, + encoder.x, + encoder.y, + encoder.qkv, + encoder.o, + encoder.gate, + m * 4 * D * 2, 2048 * 4, 17 * 4, 2048 * D * 2, @@ -160,11 +169,22 @@ impl Workspace { 16 * 2 * 2, ]) } - fn required_bytes(b: usize, l: usize) -> Result { - Ok(Self::sizes(b, l)?.iter().sum()) + fn required_bytes( + b: usize, + l: usize, + structure: &ModernBertStructure, + limits: &BackendLimits, + ) -> Result { + Ok(Self::sizes(b, l, structure, limits)?.iter().sum()) } - fn new(c: &Cuda, b: usize, l: usize) -> Result { - let sizes = Self::sizes(b, l)?; + fn new( + c: &Cuda, + b: usize, + l: usize, + structure: &ModernBertStructure, + limits: &BackendLimits, + ) -> Result { + let sizes = Self::sizes(b, l, structure, limits)?; let bytes = sizes.iter().sum(); let mut sizes = sizes.into_iter(); let mut alloc = || c.alloc(sizes.next().expect("workspace buffer layout")); @@ -212,18 +232,6 @@ impl Default for CacheConfig { } } -#[derive(Clone, Copy)] -enum Scratch { - Ids, - Lens, - Types, - X, - Y, - Qkv, - O, - G, - Ff, -} impl Scratch { fn buffer(self, s: &Workspace) -> &Buffer { match self { @@ -300,216 +308,45 @@ enum Step { }, } +fn resolve_step( + cuda: &Cuda, + weights: &HashMap, + planned: &PlannedStep, +) -> Result { + match planned { + PlannedStep::Launch { kernel, args } => { + ensure!(args.len() <= 5, "encoder argument count"); + let mut resolved = [Argument::Weight(std::ptr::null_mut()); 5]; + for (slot, arg) in resolved.iter_mut().zip(args) { + *slot = match arg { + PlannedArg::Weight(name) => Argument::Weight(weights[name].ptr()), + PlannedArg::Scratch(s) => Argument::Scratch(*s), + }; + } + Ok(Step::Launch { + kernel: SelectedKernel::resolve(cuda, kernel)?, + args: resolved, + count: args.len(), + }) + } + PlannedStep::Dump { name, slot, bf } => Ok(Step::Dump { + name: name.clone(), + slot: *slot, + bf: *bf, + }), + } +} + fn prepare_encoder( cuda: &Cuda, weights: &HashMap, + structure: &ModernBertStructure, original_rope: bool, ) -> Result> { - let steps = std::cell::RefCell::new(Vec::new()); - let weight = |name: &str| Argument::Weight(weights[name].ptr()); - let call = |name: &str, args: &[Argument]| -> Result<()> { - ensure!(args.len() <= 5, "encoder argument count"); - let mut resolved = [Argument::Weight(std::ptr::null_mut()); 5]; - resolved[..args.len()].copy_from_slice(args); - steps.borrow_mut().push(Step::Launch { - kernel: SelectedKernel::resolve(cuda, name)?, - args: resolved, - count: args.len(), - }); - Ok(()) - }; - let dump = |name: &str, slot: Scratch, bf: bool| { - steps.borrow_mut().push(Step::Dump { - name: name.to_owned(), - slot, - bf, - }); - Ok::<_, anyhow::Error>(()) - }; - let z = weight("zeros.1024"); - let attention = |label: &str| format!("attn_{label}"); - call( - "embed", - &[ - Argument::Scratch(Scratch::Ids), - weight("encoder.embeddings.tok_embeddings.weight"), - weight("encoder.embeddings.norm.weight"), - Argument::Scratch(Scratch::X), - Argument::Scratch(Scratch::Y), - ], - )?; - dump("embedding", Scratch::X, false)?; - for i in 0..28 { - let p = format!("encoder.layers.{i}"); - let w = |n: &str| weight(&format!("{p}.{n}")); - call( - "qkv", - &[ - Argument::Scratch(Scratch::Y), - w("attn.Wqkv.weight"), - weight("zeros.3072"), - Argument::Scratch(Scratch::Qkv), - ], - )?; - let kind = if i % 3 == 0 { "full" } else { "local" }; - call( - if original_rope { - "rope_original" - } else { - "rope" - }, - &[ - Argument::Scratch(Scratch::Qkv), - weight(&format!("rope_{kind}_cos")), - weight(&format!("rope_{kind}_sin")), - ], - )?; - call( - &attention(if i % 3 == 0 { "full" } else { "local" }), - &[ - Argument::Scratch(Scratch::Qkv), - Argument::Scratch(Scratch::Lens), - Argument::Scratch(Scratch::O), - ], - )?; - call( - "out", - &[ - Argument::Scratch(Scratch::O), - w("attn.Wo.weight"), - z, - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "addln", - &[ - Argument::Scratch(Scratch::X), - Argument::Scratch(Scratch::Y), - w("mlp_norm.weight"), - z, - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "geglu", - &[ - Argument::Scratch(Scratch::Y), - w("mlp.Wi.weight"), - Argument::Scratch(Scratch::G), - ], - )?; - call( - "down", - &[ - Argument::Scratch(Scratch::G), - w("mlp.Wo.weight"), - z, - Argument::Scratch(Scratch::Y), - ], - )?; - let next = if i < 27 { - weight(&format!("encoder.layers.{}.attn_norm.weight", i + 1)) - } else { - weight("encoder.final_norm.weight") - }; - call( - "addln", - &[ - Argument::Scratch(Scratch::X), - Argument::Scratch(Scratch::Y), - next, - z, - Argument::Scratch(Scratch::Y), - ], - )?; - if [0, 1, 2, 27].contains(&i) { - dump(&format!("encoder{i}_residual"), Scratch::X, false)?; - dump(&format!("encoder{i}_normalized"), Scratch::Y, true)?; - } - } - call( - "type", - &[ - Argument::Scratch(Scratch::Y), - weight("type_emb.weight"), - Argument::Scratch(Scratch::Types), - Argument::Scratch(Scratch::X), - ], - )?; - for i in 0..2 { - let p = format!("head.layers.{i}"); - let w = |n: &str| weight(&format!("{p}.{n}")); - call( - "ln_bias", - &[ - Argument::Scratch(Scratch::X), - Argument::Scratch(Scratch::Y), - w("norm1.weight"), - w("norm1.bias"), - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "head_in", - &[ - Argument::Scratch(Scratch::Y), - w("self_attn.in_proj_weight"), - w("self_attn.in_proj_bias"), - Argument::Scratch(Scratch::Qkv), - ], - )?; - call( - &attention("full"), - &[ - Argument::Scratch(Scratch::Qkv), - Argument::Scratch(Scratch::Lens), - Argument::Scratch(Scratch::O), - ], - )?; - call( - "head_out", - &[ - Argument::Scratch(Scratch::O), - w("self_attn.out_proj.weight"), - w("self_attn.out_proj.bias"), - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "addln_bias", - &[ - Argument::Scratch(Scratch::X), - Argument::Scratch(Scratch::Y), - w("norm2.weight"), - w("norm2.bias"), - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "ffn1", - &[ - Argument::Scratch(Scratch::Y), - w("linear1.weight"), - w("linear1.bias"), - Argument::Scratch(Scratch::Ff), - ], - )?; - call( - "ffn2", - &[ - Argument::Scratch(Scratch::Ff), - w("linear2.weight"), - w("linear2.bias"), - Argument::Scratch(Scratch::Y), - ], - )?; - call( - "residual", - &[Argument::Scratch(Scratch::X), Argument::Scratch(Scratch::Y)], - )?; - } - Ok(steps.into_inner()) + plan_encoder(structure, original_rope) + .iter() + .map(|planned| resolve_step(cuda, weights, planned)) + .collect() } pub struct Model { @@ -601,11 +438,12 @@ impl Model { let key = format!("zeros.{n}"); weights.insert(key.clone(), upload(&key, &vec![0; n * 4])?); } + let table_bytes = omni_modernbert::geometry::rope_table_bytes(&config.structure); for kind in ["full", "local"] { for part in ["cos", "sin"] { let key = format!("rope_{kind}_{part}"); let data = fs::read(bundle.join(format!("{key}.f32")))?; - ensure!(data.len() == 512 * 32 * 4, "invalid rotary table size"); + ensure!(data.len() == table_bytes, "invalid rotary table size"); weights.insert(key.clone(), upload(&key, &data)?); } } @@ -615,7 +453,7 @@ impl Model { weights.len() ); } - let plan = prepare_encoder(&cuda, &weights, original_rope)?; + let plan = prepare_encoder(&cuda, &weights, &config.structure, original_rope)?; Ok(Self { config, cuda, @@ -676,7 +514,10 @@ impl Model { ensure!( batch.b <= 16 && batch.l <= 512 - && batch.input_ids.iter().all(|i| *i >= 0 && *i < 50368), + && batch + .input_ids + .iter() + .all(|i| *i >= 0 && (*i as usize) < self.config.structure.vocab_size), "invalid packed input" ); let n = batch.markers.len(); @@ -707,7 +548,12 @@ impl Model { batch.markers.iter().map(Vec::len).sum::() <= 2048, "too many markers" ); - let bytes = Workspace::required_bytes(batch.b, batch.l)?; + let bytes = Workspace::required_bytes( + batch.b, + batch.l, + &self.config.structure, + &self.config.limits, + )?; let cacheable = self.cache_config.max_shapes > 0 && bytes <= self.cache_config.max_bytes; let mut s = if cacheable { self.eager = None; @@ -729,7 +575,13 @@ impl Model { self.cached_bytes -= evicted.bytes; drop(evicted); } - Workspace::new(&self.cuda, batch.b, batch.l)? + Workspace::new( + &self.cuda, + batch.b, + batch.l, + &self.config.structure, + &self.config.limits, + )? } } else { let eager = self @@ -739,7 +591,13 @@ impl Model { // A different fallback shape is dropped before replacement allocation. match eager { Some(workspace) => workspace, - None => Workspace::new(&self.cuda, batch.b, batch.l)?, + None => Workspace::new( + &self.cuda, + batch.b, + batch.l, + &self.config.structure, + &self.config.limits, + )?, } }; let (ids, rest) = s.staging.split_at_mut(batch.input_ids.len() * 8); diff --git a/src/models/laya/src/packing.rs b/src/models/laya/src/packing.rs index 68a0601b..276aa95e 100644 --- a/src/models/laya/src/packing.rs +++ b/src/models/laya/src/packing.rs @@ -14,7 +14,7 @@ pub struct Batch { pub markers: Vec>, } -pub fn pack(prepared: &Prepared) -> Result { +pub fn pack(prepared: &Prepared, vocab: usize) -> Result { let n = prepared.questions.len(); ensure!(n <= 16, "at most 16 questions per CUDA request"); ensure!( @@ -47,7 +47,7 @@ pub fn pack(prepared: &Prepared) -> Result { for (i, q) in prepared.questions.iter().enumerate() { ensure!( !q.ids.is_empty() - && q.ids.iter().all(|id| *id < 50368) + && q.ids.iter().all(|id| (*id as usize) < vocab) && (0..=2).contains(&q.qtype) && !q.markers.is_empty() && q.markers.iter().all(|m| *m < q.ids.len()), diff --git a/src/models/laya/src/plan.rs b/src/models/laya/src/plan.rs new file mode 100644 index 00000000..5263d89f --- /dev/null +++ b/src/models/laya/src/plan.rs @@ -0,0 +1,228 @@ +//! The encoder-plus-head execution plan as pure data: kernel names and +//! buffer roles, before any CUDA symbol or weight pointer is resolved. +//! The encoder half is generated by the family core from the admitted +//! structure; the task-head half stays here. Keeping the plan pure lets +//! CPU fixtures pin the exact step sequence. + +use omni_modernbert::plan::{ + EncoderPlanOptions, PlanAnchor, PlanArg, PlanStep, ScratchSlot, + encoder_plan as family_encoder_plan, +}; +use omni_modernbert::structure::ModernBertStructure; + +/// Scratch buffers the plan addresses; allocation lives in the executor. +#[derive(Clone, Copy)] +pub enum Scratch { + Ids, + Lens, + Types, + X, + Y, + Qkv, + O, + G, + Ff, +} + +impl Scratch { + pub fn name(self) -> &'static str { + match self { + Self::Ids => "ids", + Self::Lens => "lens", + Self::Types => "types", + Self::X => "x", + Self::Y => "y", + Self::Qkv => "qkv", + Self::O => "o", + Self::G => "g", + Self::Ff => "ff", + } + } +} + +#[derive(Clone)] +pub enum PlannedArg { + Weight(String), + Scratch(Scratch), +} + +#[derive(Clone)] +pub enum PlannedStep { + Launch { + kernel: String, + args: Vec, + }, + Dump { + name: String, + slot: Scratch, + bf: bool, + }, +} + +impl PlannedStep { + /// One-line canonical rendering used by the plan fixtures. + pub fn render(&self) -> String { + match self { + Self::Launch { kernel, args } => { + let mut line = format!("launch {kernel}"); + for arg in args { + line.push(' '); + line.push_str(&match arg { + PlannedArg::Weight(name) => name.clone(), + PlannedArg::Scratch(slot) => slot.name().to_string(), + }); + } + line + } + Self::Dump { name, slot, bf } => { + format!("dump {name} {} {}", slot.name(), u8::from(*bf)) + } + } + } +} + +fn map_arg(arg: PlanArg) -> PlannedArg { + match arg { + PlanArg::Weight(name) => PlannedArg::Weight(name), + PlanArg::Scratch(slot) => PlannedArg::Scratch(match slot { + ScratchSlot::Ids => Scratch::Ids, + ScratchSlot::Lens => Scratch::Lens, + ScratchSlot::Types => Scratch::Types, + ScratchSlot::X => Scratch::X, + ScratchSlot::Y => Scratch::Y, + ScratchSlot::Qkv => Scratch::Qkv, + ScratchSlot::O => Scratch::O, + ScratchSlot::Gate => Scratch::G, + }), + } +} + +/// The execution plan for the admitted Laya profile: the family encoder +/// plan with this bundle's debug dumps hooked to the sequence anchors, +/// followed by the task head's two transformer layers. The readout tail +/// stays with eager launches at inference time. +pub fn plan_encoder(structure: &ModernBertStructure, original_rope: bool) -> Vec { + let mut steps = Vec::new(); + for step in family_encoder_plan(structure, &EncoderPlanOptions { original_rope }) { + match step { + PlanStep::Launch { kernel, args } => steps.push(PlannedStep::Launch { + kernel, + args: args.into_iter().map(map_arg).collect(), + }), + PlanStep::Anchor(PlanAnchor::AfterEmbed) => steps.push(PlannedStep::Dump { + name: "embedding".to_string(), + slot: Scratch::X, + bf: false, + }), + PlanStep::Anchor(PlanAnchor::LayerEnd(i)) => { + if [0, 1, 2, 27].contains(&i) { + steps.push(PlannedStep::Dump { + name: format!("encoder{i}_residual"), + slot: Scratch::X, + bf: false, + }); + steps.push(PlannedStep::Dump { + name: format!("encoder{i}_normalized"), + slot: Scratch::Y, + bf: true, + }); + } + } + PlanStep::Anchor(PlanAnchor::AfterEncoder) => {} + } + } + steps.extend(head_plan()); + steps +} + +fn head_plan() -> Vec { + let steps = std::cell::RefCell::new(Vec::new()); + let weight = |name: &str| PlannedArg::Weight(name.to_string()); + let scratch = |slot: Scratch| PlannedArg::Scratch(slot); + let call = |kernel: &str, args: &[PlannedArg]| { + steps.borrow_mut().push(PlannedStep::Launch { + kernel: kernel.to_string(), + args: args.to_vec(), + }); + }; + let attention = |label: &str| format!("attn_{label}"); + for i in 0..2 { + let p = format!("head.layers.{i}"); + let w = |n: &str| weight(&format!("{p}.{n}")); + call( + "ln_bias", + &[ + scratch(Scratch::X), + scratch(Scratch::Y), + w("norm1.weight"), + w("norm1.bias"), + scratch(Scratch::Y), + ], + ); + call( + "head_in", + &[ + scratch(Scratch::Y), + w("self_attn.in_proj_weight"), + w("self_attn.in_proj_bias"), + scratch(Scratch::Qkv), + ], + ); + call( + &attention("full"), + &[ + scratch(Scratch::Qkv), + scratch(Scratch::Lens), + scratch(Scratch::O), + ], + ); + call( + "head_out", + &[ + scratch(Scratch::O), + w("self_attn.out_proj.weight"), + w("self_attn.out_proj.bias"), + scratch(Scratch::Y), + ], + ); + call( + "addln_bias", + &[ + scratch(Scratch::X), + scratch(Scratch::Y), + w("norm2.weight"), + w("norm2.bias"), + scratch(Scratch::Y), + ], + ); + call( + "ffn1", + &[ + scratch(Scratch::Y), + w("linear1.weight"), + w("linear1.bias"), + scratch(Scratch::Ff), + ], + ); + call( + "ffn2", + &[ + scratch(Scratch::Ff), + w("linear2.weight"), + w("linear2.bias"), + scratch(Scratch::Y), + ], + ); + call("residual", &[scratch(Scratch::X), scratch(Scratch::Y)]); + } + steps.into_inner() +} + +/// Renders a full plan, one step per line, for fixture comparison. +pub fn render_plan(plan: &[PlannedStep]) -> String { + plan.iter() + .map(PlannedStep::render) + .collect::>() + .join("\n") + + "\n" +} diff --git a/src/models/laya/src/processing.rs b/src/models/laya/src/processing.rs index a05c7737..21b62b23 100644 --- a/src/models/laya/src/processing.rs +++ b/src/models/laya/src/processing.rs @@ -14,6 +14,7 @@ use crate::{ pub struct Processor { preprocessor: Preprocessor, config: Arc, + vocab: usize, } pub struct PreparedRequest { @@ -28,9 +29,11 @@ pub struct ResponseContext { impl Processor { pub fn load(checkpoint: &Path) -> Result { + let config = Config::load(checkpoint)?; Ok(Self { preprocessor: Preprocessor::load(&checkpoint.join("tokenizer/tokenizer.json"))?, - config: Arc::new(Config::load(checkpoint)?.agent), + vocab: config.structure.vocab_size, + config: Arc::new(config.agent), }) } @@ -38,7 +41,7 @@ impl Processor { let request = Request::from_json(std::str::from_utf8(raw)?)?; let prepared = self.preprocessor.prepare(&request)?; Ok(PreparedRequest { - inputs: packing::pack(&prepared)?, + inputs: packing::pack(&prepared, self.vocab)?, context: ResponseContext { prepared, config: self.config.clone(), diff --git a/src/models/laya/src/profile.rs b/src/models/laya/src/profile.rs new file mode 100644 index 00000000..9a79bd85 --- /dev/null +++ b/src/models/laya/src/profile.rs @@ -0,0 +1,280 @@ +//! The Laya bundle dialect and its one registered family profile. +//! +//! The two-file recipe bundle (`rl_agent_config.json` + +//! `encoder/config.json`) is a Laya checkpoint format, so it is parsed +//! here rather than in the family crate. The dialect reads exactly the +//! keys the native worker has always admitted: dimensions, explicit +//! `layer_types` schedule, `local_attention`, the two-entry +//! `rope_parameters`, `norm_eps`, `vocab_size`. HF-side keys the +//! worker does not consume (`max_position_embeddings`, `head_dim`) are +//! recorded as unrecognized instead of gated — the bundle's rotary +//! tables exercise 512 positions and the head dim is hidden/heads, as +//! before the family split. + +use anyhow::{Context, Result, ensure}; +use omni_modernbert::gate::{BackendLimits, RegisteredProfile}; +use omni_modernbert::structure::{ + Activation, AttentionKind, FieldOrigin, ModernBertStructure, RopeSpec, +}; +use serde_json::Value as Json; +use std::collections::BTreeMap; +use std::path::Path; + +/// Reads the two-file Laya bundle layout: the encoder half of +/// `rl_agent_config.json` + `encoder/config.json`. Agent fields are +/// task bundle and stay in the Laya crate; this adapter reads +/// `encoder/config.json` only. +pub struct LayaBundleAdapter; + +impl LayaBundleAdapter { + pub fn read(dir: &Path) -> Result<(ModernBertStructure, SourceMeta)> { + let path = dir.join("encoder/config.json"); + let bytes = std::fs::read(&path).with_context(|| format!("{}", path.display()))?; + let config: Json = + serde_json::from_slice(&bytes).with_context(|| format!("parse {}", path.display()))?; + parse_laya_bundle(&config, "encoder/config.json") + } +} + +/// Marker and file evidence for one parsed source, mirroring the family +/// adapter's metadata. +#[derive(Debug, Clone, PartialEq)] +pub struct SourceMeta { + pub config_files: Vec, + pub model_type: String, + /// Top-level keys the dialect does not consume. Recorded, never a + /// rejection. + pub unrecognized_keys: Vec, +} + +fn uint(config: &Json, key: &str) -> Result { + config[key] + .as_u64() + .and_then(|v| usize::try_from(v).ok()) + .with_context(|| format!("{key} is missing or not an unsigned integer")) +} + +fn float(config: &Json, key: &str) -> Result { + config[key] + .as_f64() + .with_context(|| format!("{key} is missing or not a number")) +} + +fn require_modernbert(config: &Json) -> Result<()> { + ensure!( + config["model_type"].as_str() == Some("modernbert"), + "model_type {:?} is not modernbert", + config["model_type"] + ); + ensure!( + config["hidden_activation"].as_str() == Some("gelu"), + "hidden_activation {:?} is not gelu", + config["hidden_activation"] + ); + Ok(()) +} + +fn unrecognized(config: &Json, known: &[&str]) -> Vec { + config + .as_object() + .map(|object| { + object + .keys() + .filter(|k| !known.contains(&k.as_str())) + .cloned() + .collect() + }) + .unwrap_or_default() +} + +fn parse_laya_bundle(config: &Json, source: &str) -> Result<(ModernBertStructure, SourceMeta)> { + require_modernbert(config)?; + for key in ["attention_bias", "mlp_bias", "norm_bias"] { + ensure!( + config[key] == Json::Bool(false), + "{key} must be stated and false" + ); + } + let layer_types = config["layer_types"] + .as_array() + .context("layer_types is missing or not an array")?; + let schedule = layer_types + .iter() + .map(|t| { + t.as_str() + .map(AttentionKind::parse) + .unwrap_or_else(|| Err(anyhow::anyhow!("layer_types entry {t:?} is not a string"))) + }) + .collect::>>()?; + let rope = |kind: &str, provenance: &mut BTreeMap| { + let entry = &config["rope_parameters"][kind]; + ensure!( + entry["rope_type"].as_str() == Some("default"), + "rope_parameters.{kind} rope_type {:?} is not default", + entry["rope_type"] + ); + provenance.insert( + format!("rope_{kind}"), + FieldOrigin::Source(format!("{source}#rope_parameters.{kind}")), + ); + Ok(RopeSpec { + theta: entry["rope_theta"].as_f64().with_context(|| { + format!("rope_parameters.{kind}.rope_theta is missing or not a number") + })?, + scaling: None, + }) + }; + let mut provenance = BTreeMap::new(); + let rope_full = rope("full_attention", &mut provenance)?; + let rope_sliding = rope("sliding_attention", &mut provenance)?; + let hidden = uint(config, "hidden_size")?; + let heads = uint(config, "num_attention_heads")?; + provenance.insert( + "max_positions".into(), + FieldOrigin::Default("Laya bundle rotary tables exercise 512 positions"), + ); + provenance.insert( + "head_dim".into(), + FieldOrigin::Derived("hidden_size / num_attention_heads".into()), + ); + for key in [ + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "norm_eps", + "num_attention_heads", + "schedule", + "local_window", + "vocab_size", + ] { + let json_key = match key { + "schedule" => "layer_types", + "local_window" => "local_attention", + _ => key, + }; + provenance.insert( + key.into(), + FieldOrigin::Source(format!("{source}#{json_key}")), + ); + } + provenance.insert( + "type_embedding_rows".into(), + FieldOrigin::Default("the Laya bundle checkpoint carries a [3, hidden] type embedding"), + ); + let structure = ModernBertStructure { + hidden_size: hidden, + intermediate_size: uint(config, "intermediate_size")?, + num_hidden_layers: uint(config, "num_hidden_layers")?, + activation: Activation::Gelu, + norm_eps: float(config, "norm_eps")? as f32, + num_attention_heads: heads, + head_dim: hidden / heads, + schedule, + local_window: uint(config, "local_attention")?, + rope_full, + rope_sliding, + vocab_size: uint(config, "vocab_size")?, + type_embedding_rows: Some(3), + max_positions: 512, + provenance, + }; + structure.validate()?; + let meta = SourceMeta { + config_files: vec![source.into()], + model_type: "modernbert".into(), + unrecognized_keys: unrecognized( + config, + &[ + "model_type", + "hidden_activation", + "attention_bias", + "mlp_bias", + "norm_bias", + "hidden_size", + "intermediate_size", + "num_attention_heads", + "num_hidden_layers", + "vocab_size", + "norm_eps", + "local_attention", + "layer_types", + "rope_parameters", + ], + ), + }; + Ok((structure, meta)) +} + +/// The Laya 0.3.20 encoder as shipped, ABI 1. +pub fn laya_0_3_20() -> RegisteredProfile { + RegisteredProfile { + name: "laya-0.3.20", + structure: laya_structure(), + limits: BackendLimits { + hidden_sizes: vec![1024], + batch_max: 16, + len_min: 16, + len_max: 512, + len_step: 16, + yarn_scaling: false, + }, + backend: "laya_cuda", + abi_version: 1, + } +} + +fn laya_structure() -> ModernBertStructure { + let mut provenance = BTreeMap::new(); + for field in [ + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "norm_eps", + "num_attention_heads", + "head_dim", + "schedule", + "local_window", + "rope_full", + "rope_sliding", + "vocab_size", + "type_embedding_rows", + "max_positions", + ] { + provenance.insert(field.to_string(), origin_of_profile()); + } + ModernBertStructure { + hidden_size: 1024, + intermediate_size: 2624, + num_hidden_layers: 28, + activation: Activation::Gelu, + norm_eps: 1e-5, + num_attention_heads: 16, + head_dim: 64, + schedule: (0..28) + .map(|i| { + if i % 3 == 0 { + AttentionKind::Full + } else { + AttentionKind::Sliding + } + }) + .collect(), + local_window: 128, + rope_full: RopeSpec { + theta: 160000.0, + scaling: None, + }, + rope_sliding: RopeSpec { + theta: 10000.0, + scaling: None, + }, + vocab_size: 50368, + type_embedding_rows: Some(3), + max_positions: 512, + provenance, + } +} + +fn origin_of_profile() -> FieldOrigin { + FieldOrigin::Source("registered profile laya-0.3.20".into()) +} diff --git a/src/models/laya/src/weights.rs b/src/models/laya/src/weights.rs index d3e5a1f8..19858b20 100644 --- a/src/models/laya/src/weights.rs +++ b/src/models/laya/src/weights.rs @@ -1,6 +1,7 @@ use anyhow::{Result, ensure}; use half::{bf16, f16}; use memmap2::Mmap; +use omni_modernbert::geometry; use safetensors::{Dtype, SafeTensors}; use std::{fs::File, path::Path}; @@ -93,37 +94,21 @@ pub struct TensorSpec { pub shape: Vec, } pub fn checkpoint_tensors() -> Vec { - const D: usize = 1024; - let mut tensors = Vec::new(); - let mut add = |name: &str, shape: &[usize]| { - tensors.push(TensorSpec { - name: name.into(), - shape: shape.into(), - }); + // Encoder tensors come from the family structure; the task head keeps + // its own inventory. + let structure = crate::profile::laya_0_3_20().structure; + let mut tensors: Vec<_> = geometry::encoder_tensors(&structure) + .into_iter() + .map(|(name, shape)| TensorSpec { name, shape }) + .collect(); + let mut add = |name: String, shape: Vec| { + tensors.push(TensorSpec { name, shape }); }; - add("encoder.embeddings.tok_embeddings.weight", &[50368, D]); - add("encoder.embeddings.norm.weight", &[D]); - add("encoder.final_norm.weight", &[D]); - for i in 0..28 { - let p = format!("encoder.layers.{i}"); - if i > 0 { - add(&format!("{p}.attn_norm.weight"), &[D]); - } - add(&format!("{p}.mlp_norm.weight"), &[D]); - for (name, shape) in [ - ("attn.Wqkv.weight", vec![3 * D, D]), - ("attn.Wo.weight", vec![D, D]), - ("mlp.Wi.weight", vec![5248, D]), - ("mlp.Wo.weight", vec![D, 2624]), - ] { - add(&format!("{p}.{name}"), &shape); - } - } - add("type_emb.weight", &[3, D]); + const D: usize = 1024; for i in 0..2 { let p = format!("head.layers.{i}"); for n in ["norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias"] { - add(&format!("{p}.{n}"), &[D]); + add(format!("{p}.{n}"), vec![D]); } for (n, rows, cols) in [ ("self_attn.in_proj_weight", 3 * D, D), @@ -131,7 +116,7 @@ pub fn checkpoint_tensors() -> Vec { ("linear1.weight", 4 * D, D), ("linear2.weight", D, 4 * D), ] { - add(&format!("{p}.{n}"), &[rows, cols]); + add(format!("{p}.{n}"), vec![rows, cols]); } for (n, len) in [ ("self_attn.in_proj_bias", 3 * D), @@ -139,11 +124,11 @@ pub fn checkpoint_tensors() -> Vec { ("linear1.bias", 4 * D), ("linear2.bias", D), ] { - add(&format!("{p}.{n}"), &[len]); + add(format!("{p}.{n}"), vec![len]); } } for n in ["scorer.0.weight", "scorer.0.bias"] { - add(n, &[D]); + add(n.into(), vec![D]); } for (p, n, k) in [ ("scorer.1", D, D), @@ -151,13 +136,13 @@ pub fn checkpoint_tensors() -> Vec { ("act_head.0", 256, 1028), ("act_head.2", 2, 256), ] { - add(&format!("{p}.weight"), &[n, k]); - add(&format!("{p}.bias"), &[n]); + add(format!("{p}.weight"), vec![n, k]); + add(format!("{p}.bias"), vec![n]); } // Laya 0.3.20 common.py registers this legacy buffer but forward does not // consume it. Agent decoding uses fitted config temperatures instead. // Keep it in the inventory for complete checkpoint accounting. - add("temperature", &[3]); + add("temperature".into(), vec![3]); tensors } diff --git a/src/models/modernbert/Cargo.toml b/src/models/modernbert/Cargo.toml new file mode 100644 index 00000000..2f08f654 --- /dev/null +++ b/src/models/modernbert/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "omni-modernbert" +version = "0.1.0" +edition = "2024" +publish = false + +[dependencies] +anyhow = "1" +serde_json = "1" +sha2 = "0.10" + +[dev-dependencies] +tempfile = "3" + +[[test]] +name = "structure" +path = "../../../tests/modernbert/structure.rs" + +[[test]] +name = "plan" +path = "../../../tests/modernbert/plan.rs" diff --git a/src/models/modernbert/README.md b/src/models/modernbert/README.md new file mode 100644 index 00000000..a41b1b8d --- /dev/null +++ b/src/models/modernbert/README.md @@ -0,0 +1,23 @@ +# ModernBERT encoder family core + +`omni-modernbert` owns the family data shared by the ModernBERT-based +encoders: structure parameters parsed from checkpoint configs with +per-field provenance ([`structure`](src/structure.rs)), the tensor +geometry derived from a structure ([`geometry`](src/geometry.rs)), +registered capability profiles that decide whether a CUDA bundle may +execute a structure ([`gate`](src/gate.rs)), and the named-kernel +execution plan for the encoder stack ([`plan`](src/plan.rs)). + +The crate is family data only. It never touches CUDA, device memory, or +weight files: loading weights, resolving kernel symbols to device +pointers, and every task head stay with the model crates that consume a +profile. The Laya native worker runs its encoder on this core +([`src/models/laya`](../laya/README.md)), and the Vela 2.0 0.3B native +profile is the second registered family member +([`src/models/vela/native`](../vela/native/README.md)). + +Structure parsing and capability admission are CPU-verifiable without a +checkpoint or device: the tests under +[`tests/modernbert/`](../../../tests/modernbert/) pin adapter +provenance, family invariants, the frozen gate's accept/reject set, and +the serialized execution plans as golden fixtures. diff --git a/src/models/modernbert/src/adapter.rs b/src/models/modernbert/src/adapter.rs new file mode 100644 index 00000000..87e18bc5 --- /dev/null +++ b/src/models/modernbert/src/adapter.rs @@ -0,0 +1,335 @@ +use crate::structure::{ + Activation, AttentionKind, FieldOrigin, ModernBertStructure, RopeSpec, YarnScaling, +}; +use anyhow::{Context, Result, ensure}; +use serde_json::Value as Json; +use std::collections::BTreeMap; + +/// Marker and file evidence for one parsed source. +#[derive(Debug, Clone, PartialEq)] +pub struct SourceMeta { + pub config_files: Vec, + pub model_type: String, + /// Top-level keys the adapters do not consume. Recorded, never a + /// rejection: a checkpoint may carry recipe or training fields beside + /// the structure. + pub unrecognized_keys: Vec, +} + +/// Reads any single-object ModernBERT config: a stock HF +/// `config.json` or a member object such as the Vela bundle's +/// `encoder_config`. Dialect bundle layouts (for example the Laya +/// two-file recipe) are parsed by the model crates that ship them. +pub struct ModernBertConfigAdapter; + +impl ModernBertConfigAdapter { + pub fn from_value( + config: &Json, + source_name: &str, + ) -> Result<(ModernBertStructure, SourceMeta)> { + parse_hf_config(config, source_name) + } +} + +fn uint(config: &Json, key: &str) -> Result { + config[key] + .as_u64() + .and_then(|v| usize::try_from(v).ok()) + .with_context(|| format!("{key} is missing or not an unsigned integer")) +} + +fn float(config: &Json, key: &str) -> Result { + config[key] + .as_f64() + .with_context(|| format!("{key} is missing or not a number")) +} + +fn require_modernbert(config: &Json) -> Result<()> { + ensure!( + config["model_type"].as_str() == Some("modernbert"), + "model_type {:?} is not modernbert", + config["model_type"] + ); + ensure!( + config["hidden_activation"].as_str() == Some("gelu"), + "hidden_activation {:?} is not gelu", + config["hidden_activation"] + ); + Ok(()) +} + +fn unrecognized(config: &Json, known: &[&str]) -> Vec { + config + .as_object() + .map(|object| { + object + .keys() + .filter(|k| !known.contains(&k.as_str())) + .cloned() + .collect() + }) + .unwrap_or_default() +} + +fn derive_head_dim( + hidden: usize, + heads: usize, + config: &Json, + provenance: &mut BTreeMap, +) -> Result { + let head_dim = match config.get("head_dim") { + Some(value) => value + .as_u64() + .and_then(|v| usize::try_from(v).ok()) + .context("head_dim is not an unsigned integer")?, + None => { + provenance.insert( + "head_dim".into(), + FieldOrigin::Derived("hidden_size / num_attention_heads".into()), + ); + hidden / heads + } + }; + Ok(head_dim) +} + +fn parse_hf_config(config: &Json, source: &str) -> Result<(ModernBertStructure, SourceMeta)> { + require_modernbert(config)?; + let mut provenance = BTreeMap::new(); + for key in ["attention_bias", "mlp_bias", "norm_bias"] { + match config.get(key) { + None => { + provenance.insert( + key.into(), + FieldOrigin::Default("HF ModernBERT states no bias by default"), + ); + } + Some(value) => ensure!( + value == &Json::Bool(false), + "{key} must be false, got {value}" + ), + } + } + let schedule = if let Some(layer_types) = config["layer_types"].as_array() { + let parsed = layer_types + .iter() + .map(|t| { + t.as_str().map(AttentionKind::parse).unwrap_or_else(|| { + Err(anyhow::anyhow!("layer_types entry {t:?} is not a string")) + }) + }) + .collect::>>()?; + provenance.insert( + "schedule".into(), + FieldOrigin::Source(format!("{source}#layer_types")), + ); + parsed + } else { + let every = uint(config, "global_attn_every_n_layers") + .context("neither layer_types nor global_attn_every_n_layers is present")?; + // The schedule below tests i % every; a checkpoint config is an + // external startup input, so refuse the zero before it panics. + ensure!( + every > 0, + "global_attn_every_n_layers must be positive, got 0" + ); + provenance.insert( + "schedule".into(), + FieldOrigin::Derived(format!( + "global_attn_every_n_layers={every}, full at layer index % {every} == 0" + )), + ); + (0..uint(config, "num_hidden_layers")?) + .map(|i| { + if i % every == 0 { + AttentionKind::Full + } else { + AttentionKind::Sliding + } + }) + .collect() + }; + // ModernBERT configs carry per-kind thetas; a stock HF config may + // carry only the generic rope_theta, which then fills both kinds. + let theta = |field: &str, + json_key: &str, + generic_key: &str, + provenance: &mut BTreeMap| { + let key = if config.get(json_key).is_some() { + json_key + } else if config.get(generic_key).is_some() { + generic_key + } else { + provenance.insert( + field.into(), + FieldOrigin::Default("HF ModernBERT rope_theta default"), + ); + return Ok(10000.0); + }; + provenance.insert(field.into(), FieldOrigin::Source(format!("{source}#{key}"))); + float(config, key) + }; + let theta_full = theta( + "rope_full", + "global_rope_theta", + "rope_theta", + &mut provenance, + )?; + let theta_sliding = theta( + "rope_sliding", + "local_rope_theta", + "rope_theta", + &mut provenance, + )?; + let scaling = match config.get("rope_scaling") { + None => None, + Some(entry) => { + ensure!( + entry["rope_type"].as_str() == Some("yarn"), + "rope_scaling type {:?} is not yarn", + entry["rope_type"] + ); + ensure!( + entry + .get("truncate") + .map(|v| v == &Json::Bool(true)) + .unwrap_or(true), + "rope_scaling truncate must be true" + ); + let yarn = YarnScaling { + factor: float(entry, "factor")?, + original_max_positions: uint(entry, "original_max_position_embeddings")?, + beta_fast: match entry.get("beta_fast") { + None => 32.0, + Some(_) => float(entry, "beta_fast")?, + }, + beta_slow: match entry.get("beta_slow") { + None => 1.0, + Some(_) => float(entry, "beta_slow")?, + }, + }; + provenance.insert( + "rope_full.scaling".into(), + FieldOrigin::Source(format!("{source}#rope_scaling")), + ); + provenance.insert( + "rope_sliding.scaling".into(), + FieldOrigin::Source(format!("{source}#rope_scaling")), + ); + Some(yarn) + } + }; + if let Some(position_type) = config + .get("position_embedding_type") + .and_then(|v| v.as_str()) + { + ensure!( + position_type == "sans_pos" || position_type == "rope", + "position_embedding_type {position_type:?} is not rotary-only" + ); + } + let eps_key = if config.get("norm_eps").is_some() { + "norm_eps" + } else if config.get("layer_norm_eps").is_some() { + "layer_norm_eps" + } else { + anyhow::bail!("neither norm_eps nor layer_norm_eps is present") + }; + let norm_eps = float(config, eps_key)? as f32; + let hidden = uint(config, "hidden_size")?; + let heads = uint(config, "num_attention_heads")?; + let head_dim = derive_head_dim(hidden, heads, config, &mut provenance)?; + let type_embedding_rows = match config.get("type_vocab_size") { + None => None, + Some(_) => { + let rows = uint(config, "type_vocab_size")?; + provenance.insert( + "type_embedding_rows".into(), + FieldOrigin::Source(format!("{source}#type_vocab_size")), + ); + (rows > 0).then_some(rows) + } + }; + for key in [ + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "vocab_size", + "local_window", + ] { + let json_key = match key { + "local_window" => "local_attention", + _ => key, + }; + provenance.insert( + key.into(), + FieldOrigin::Source(format!("{source}#{json_key}")), + ); + } + provenance.insert( + "norm_eps".into(), + FieldOrigin::Source(format!("{source}#{eps_key}")), + ); + provenance.insert( + "max_positions".into(), + FieldOrigin::Source(format!("{source}#max_position_embeddings")), + ); + let structure = ModernBertStructure { + hidden_size: hidden, + intermediate_size: uint(config, "intermediate_size")?, + num_hidden_layers: uint(config, "num_hidden_layers")?, + activation: Activation::Gelu, + norm_eps, + num_attention_heads: heads, + head_dim, + schedule, + local_window: uint(config, "local_attention")?, + rope_full: RopeSpec { + theta: theta_full, + scaling, + }, + rope_sliding: RopeSpec { + theta: theta_sliding, + scaling, + }, + vocab_size: uint(config, "vocab_size")?, + type_embedding_rows, + max_positions: uint(config, "max_position_embeddings")?, + provenance, + }; + structure.validate()?; + let meta = SourceMeta { + config_files: vec![source.into()], + model_type: "modernbert".into(), + unrecognized_keys: unrecognized( + config, + &[ + "model_type", + "hidden_activation", + "attention_bias", + "mlp_bias", + "norm_bias", + "hidden_size", + "intermediate_size", + "num_attention_heads", + "num_hidden_layers", + "vocab_size", + "norm_eps", + "layer_norm_eps", + "local_attention", + "layer_types", + "global_attn_every_n_layers", + "global_rope_theta", + "local_rope_theta", + "rope_theta", + "rope_scaling", + "max_position_embeddings", + "head_dim", + "position_embedding_type", + "type_vocab_size", + ], + ), + }; + Ok((structure, meta)) +} diff --git a/src/models/modernbert/src/gate.rs b/src/models/modernbert/src/gate.rs new file mode 100644 index 00000000..f6aa9fc8 --- /dev/null +++ b/src/models/modernbert/src/gate.rs @@ -0,0 +1,244 @@ +use crate::structure::{ModernBertStructure, SCHEMA_VERSION}; +use anyhow::{Result, ensure}; +use sha2::{Digest, Sha256}; + +/// What one CUDA bundle can execute: the geometry its kernels were +/// generated for and the request envelope it admits. A new bundle +/// registers new limits; the gate code does not change. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BackendLimits { + /// Hidden sizes the bundle's kernels were instantiated for. + pub hidden_sizes: Vec, + /// Batch bound; a request batch must be a power of two up to it. + pub batch_max: usize, + pub len_min: usize, + pub len_max: usize, + pub len_step: usize, + /// Whether the bundle carries YaRN-scaled rotary tables. + pub yarn_scaling: bool, +} + +impl BackendLimits { + pub fn check_request(&self, batch: usize, len: usize) -> Result<()> { + ensure!( + batch.is_power_of_two() && batch <= self.batch_max, + "batch {batch} exceeds the power-of-two bound {}", + self.batch_max + ); + ensure!( + (self.len_min..=self.len_max).contains(&len), + "packed length {len} is outside {}..={}", + self.len_min, + self.len_max + ); + ensure!( + len.is_multiple_of(self.len_step), + "packed length {len} is not a multiple of {}", + self.len_step + ); + Ok(()) + } + + pub fn check_structure(&self, s: &ModernBertStructure) -> Result<()> { + ensure!( + self.hidden_sizes.contains(&s.hidden_size), + "hidden_size {} is not instantiated in this bundle ({:?})", + s.hidden_size, + self.hidden_sizes + ); + ensure!( + self.yarn_scaling + || (s.rope_full.scaling.is_none() && s.rope_sliding.scaling.is_none()), + "rope scaling is not available in this bundle" + ); + Ok(()) + } +} + +/// A named, frozen structure-and-limits pair. Admission checks family +/// invariants, backend limits, and equality with the canonical +/// structure, and returns the identity every cache key must carry. +#[derive(Debug, Clone)] +pub struct RegisteredProfile { + pub name: &'static str, + pub structure: ModernBertStructure, + pub limits: BackendLimits, + pub backend: &'static str, + pub abi_version: usize, +} + +impl RegisteredProfile { + pub fn admit(&self, candidate: &ModernBertStructure) -> Result { + candidate.validate()?; + self.limits.check_structure(candidate)?; + let expected = &self.structure; + let field = |name: &str, actual: String, wanted: String| { + format!("profile {}: {name} {actual}, expected {wanted}", self.name) + }; + ensure!( + candidate.hidden_size == expected.hidden_size, + "{}", + field( + "hidden_size", + candidate.hidden_size.to_string(), + expected.hidden_size.to_string() + ) + ); + ensure!( + candidate.intermediate_size == expected.intermediate_size, + "{}", + field( + "intermediate_size", + candidate.intermediate_size.to_string(), + expected.intermediate_size.to_string() + ) + ); + ensure!( + candidate.num_hidden_layers == expected.num_hidden_layers, + "{}", + field( + "num_hidden_layers", + candidate.num_hidden_layers.to_string(), + expected.num_hidden_layers.to_string() + ) + ); + ensure!( + candidate.num_attention_heads == expected.num_attention_heads, + "{}", + field( + "num_attention_heads", + candidate.num_attention_heads.to_string(), + expected.num_attention_heads.to_string() + ) + ); + ensure!( + candidate.head_dim == expected.head_dim, + "{}", + field( + "head_dim", + candidate.head_dim.to_string(), + expected.head_dim.to_string() + ) + ); + ensure!( + candidate.schedule == expected.schedule, + "{}", + field( + "schedule", + render_schedule(candidate), + render_schedule(expected) + ) + ); + ensure!( + candidate.local_window == expected.local_window, + "{}", + field( + "local_window", + candidate.local_window.to_string(), + expected.local_window.to_string() + ) + ); + ensure!( + candidate.norm_eps.to_bits() == expected.norm_eps.to_bits(), + "{}", + field( + "norm_eps", + candidate.norm_eps.to_string(), + expected.norm_eps.to_string() + ) + ); + ensure!( + candidate.activation == expected.activation, + "{}", + field( + "activation", + format!("{:?}", candidate.activation), + format!("{:?}", expected.activation) + ) + ); + ensure!( + candidate.rope_full == expected.rope_full + && candidate.rope_sliding == expected.rope_sliding, + "{}", + field( + "rope", + format!("{:?}/{:?}", candidate.rope_full, candidate.rope_sliding), + format!("{:?}/{:?}", expected.rope_full, expected.rope_sliding) + ) + ); + ensure!( + candidate.vocab_size == expected.vocab_size, + "{}", + field( + "vocab_size", + candidate.vocab_size.to_string(), + expected.vocab_size.to_string() + ) + ); + ensure!( + candidate.type_embedding_rows == expected.type_embedding_rows, + "{}", + field( + "type_embedding_rows", + format!("{:?}", candidate.type_embedding_rows), + format!("{:?}", expected.type_embedding_rows) + ) + ); + ensure!( + candidate.max_positions == expected.max_positions, + "{}", + field( + "max_positions", + candidate.max_positions.to_string(), + expected.max_positions.to_string() + ) + ); + Ok(structure_identity( + candidate, + self.backend, + self.abi_version, + )) + } +} + +fn render_schedule(s: &ModernBertStructure) -> String { + s.schedule + .iter() + .map(|k| k.as_str()) + .collect::>() + .join(",") +} + +/// Stable identity of one structure executing on one backend. Any change +/// to the digest inputs means previously captured runtime state (graphs, +/// caches, prefix state) must not be reused. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StructureIdentity { + pub schema_version: usize, + pub family: &'static str, + pub backend: String, + pub abi_version: usize, + pub digest: String, +} + +pub fn structure_identity( + s: &ModernBertStructure, + backend: &str, + abi_version: usize, +) -> StructureIdentity { + let body = serde_json::json!({ + "schema_version": SCHEMA_VERSION, + "family": "modernbert", + "backend": backend, + "abi_version": abi_version, + "structure": s.canonical_json(), + }) + .to_string(); + StructureIdentity { + schema_version: SCHEMA_VERSION, + family: "modernbert", + backend: backend.to_string(), + abi_version, + digest: format!("{:x}", Sha256::digest(body.as_bytes())), + } +} diff --git a/src/models/modernbert/src/geometry.rs b/src/models/modernbert/src/geometry.rs new file mode 100644 index 00000000..0e6a9cf9 --- /dev/null +++ b/src/models/modernbert/src/geometry.rs @@ -0,0 +1,108 @@ +use crate::structure::ModernBertStructure; + +/// Every tensor shape below is derived, never read from a checkpoint: +/// the structure is the only geometry input. +pub fn embedding_shape(s: &ModernBertStructure) -> Vec { + vec![s.vocab_size, s.hidden_size] +} + +pub fn qkv_shape(s: &ModernBertStructure) -> Vec { + vec![3 * s.hidden_size, s.hidden_size] +} + +pub fn attn_out_shape(s: &ModernBertStructure) -> Vec { + vec![s.hidden_size, s.hidden_size] +} + +/// The fused gated-MLP up projection; `2 * intermediate` is derived +/// geometry, not a semantic field of its own. +pub fn mlp_in_shape(s: &ModernBertStructure) -> Vec { + vec![2 * s.intermediate_size, s.hidden_size] +} + +pub fn mlp_out_shape(s: &ModernBertStructure) -> Vec { + vec![s.hidden_size, s.intermediate_size] +} + +/// Rotary dimensions per head (rotate-half pairing needs the even half). +pub fn rotary_half(s: &ModernBertStructure) -> usize { + s.head_dim / 2 +} + +/// One cos/sin table covers `max_positions` rows of `rotary_half` f32 +/// columns; a bundle serving shorter requests sizes tables to its own +/// length bound instead. +pub fn rope_table_bytes(s: &ModernBertStructure) -> usize { + s.max_positions * rotary_half(s) * 4 +} + +/// The canonical encoder tensor inventory under family naming. Layer +/// zero's attention norm is the embedding norm (the embedding kernel +/// applies it); every layer owns its post-attention norm, and the last +/// layer's post-MLP norm is the final norm. +pub fn encoder_tensors(s: &ModernBertStructure) -> Vec<(String, Vec)> { + let h = s.hidden_size; + let mut tensors = vec![ + ( + "encoder.embeddings.tok_embeddings.weight".to_string(), + embedding_shape(s), + ), + ("encoder.embeddings.norm.weight".to_string(), vec![h]), + ("encoder.final_norm.weight".to_string(), vec![h]), + ]; + for i in 0..s.num_hidden_layers { + let p = format!("encoder.layers.{i}"); + if i > 0 { + tensors.push((format!("{p}.attn_norm.weight"), vec![h])); + } + tensors.push((format!("{p}.mlp_norm.weight"), vec![h])); + tensors.push((format!("{p}.attn.Wqkv.weight"), qkv_shape(s))); + tensors.push((format!("{p}.attn.Wo.weight"), attn_out_shape(s))); + tensors.push((format!("{p}.mlp.Wi.weight"), mlp_in_shape(s))); + tensors.push((format!("{p}.mlp.Wo.weight"), mlp_out_shape(s))); + } + if let Some(rows) = s.type_embedding_rows { + tensors.push(("type_emb.weight".to_string(), vec![rows, h])); + } + tensors +} + +/// Encoder scratch for one request of `batch` rows packed to `len` +/// tokens, in canonical allocation order. Task-side buffers are sized +/// by the consuming model crate and appended after these. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EncoderScratch { + pub ids: usize, + pub lens: usize, + pub x: usize, + pub y: usize, + pub qkv: usize, + pub o: usize, + pub gate: usize, +} + +impl EncoderScratch { + pub fn new(batch: usize, len: usize, s: &ModernBertStructure) -> Self { + let m = batch * len; + Self { + ids: m * 8, + lens: batch * 4, + x: m * s.hidden_size * 4, + y: m * s.hidden_size * 2, + qkv: m * 3 * s.hidden_size * 2, + o: m * s.hidden_size * 2, + gate: m * s.intermediate_size * 2, + } + } + + pub fn sizes(&self) -> Vec { + [ + self.ids, self.lens, self.x, self.y, self.qkv, self.o, self.gate, + ] + .into() + } + + pub fn total(&self) -> usize { + self.sizes().iter().sum() + } +} diff --git a/src/models/modernbert/src/lib.rs b/src/models/modernbert/src/lib.rs new file mode 100644 index 00000000..48daa370 --- /dev/null +++ b/src/models/modernbert/src/lib.rs @@ -0,0 +1,23 @@ +//! The ModernBERT encoder family layer: structure parameters parsed from +//! checkpoint configs, the geometry derived from them, the capability +//! profiles deciding whether a CUDA bundle can execute a structure, and +//! the named-kernel execution plan for the encoder stack. +//! +//! This crate owns family data only. It never touches CUDA, device +//! memory, or weight files; loading weights, resolving kernel symbols, +//! and every task head stay with the model crates that consume a +//! profile. + +pub mod adapter; +pub mod gate; +pub mod geometry; +pub mod plan; +pub mod structure; + +pub use adapter::{ModernBertConfigAdapter, SourceMeta}; +pub use gate::{BackendLimits, RegisteredProfile, StructureIdentity}; +pub use plan::{EncoderPlanOptions, PlanAnchor, PlanArg, PlanStep, ScratchSlot}; +pub use structure::{ + Activation, AttentionKind, FieldOrigin, ModernBertStructure, RopeSpec, SCHEMA_VERSION, + YarnScaling, +}; diff --git a/src/models/modernbert/src/plan.rs b/src/models/modernbert/src/plan.rs new file mode 100644 index 00000000..771767fa --- /dev/null +++ b/src/models/modernbert/src/plan.rs @@ -0,0 +1,206 @@ +use crate::structure::ModernBertStructure; + +/// Scratch slots an encoder plan addresses. Allocating them is the +/// consumer's job; `ids`, `lens`, `x`, `y`, `qkv`, `o` and `gate` map to +/// the encoder scratch of [`crate::geometry::EncoderScratch`], and +/// `types` holds the per-row type ids consumed by the type-embedding +/// kernel when the structure carries one. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ScratchSlot { + Ids, + Lens, + Types, + X, + Y, + Qkv, + O, + Gate, +} + +impl ScratchSlot { + pub fn name(self) -> &'static str { + match self { + Self::Ids => "ids", + Self::Lens => "lens", + Self::Types => "types", + Self::X => "x", + Self::Y => "y", + Self::Qkv => "qkv", + Self::O => "o", + Self::Gate => "g", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PlanArg { + /// A weight-buffer name: checkpoint tensors under canonical family + /// naming, or `zeros.` synthetic zero-bias buffers the consumer + /// uploads alongside them. + Weight(String), + Scratch(ScratchSlot), +} + +/// Sequence anchors a consumer can hook instrumentation to; they carry +/// no execution of their own. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PlanAnchor { + AfterEmbed, + LayerEnd(usize), + AfterEncoder, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PlanStep { + Launch { kernel: String, args: Vec }, + Anchor(PlanAnchor), +} + +impl PlanStep { + /// Canonical one-line rendering used by plan fixtures. + pub fn render(&self) -> String { + match self { + Self::Launch { kernel, args } => { + let mut line = kernel.clone(); + for arg in args { + line.push(' '); + line.push_str(&match arg { + PlanArg::Weight(w) => w.clone(), + PlanArg::Scratch(s) => s.name().to_string(), + }); + } + line + } + Self::Anchor(a) => format!("anchor {a:?}"), + } + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct EncoderPlanOptions { + /// Select the reference-rotation kernel variant a bundle may carry + /// beside the shipped tables. + pub original_rope: bool, +} + +/// The named-kernel encoder plan: embedding, per-layer attention and +/// gated MLP, the final norm fused into the last layer's post-MLP norm, +/// and the optional type-embedding add. Kernel names and argument roles +/// only — resolving symbols to loaded kernels and names to device +/// pointers stays with the consuming model crate. +pub fn encoder_plan(s: &ModernBertStructure, opts: &EncoderPlanOptions) -> Vec { + let h = s.hidden_size; + let zero_bias = format!("zeros.{h}"); + let weight = |name: String| PlanArg::Weight(name); + let scratch = PlanArg::Scratch; + let mut steps = Vec::new(); + steps.push(PlanStep::Launch { + kernel: "embed".into(), + args: vec![ + scratch(ScratchSlot::Ids), + weight("encoder.embeddings.tok_embeddings.weight".into()), + weight("encoder.embeddings.norm.weight".into()), + scratch(ScratchSlot::X), + scratch(ScratchSlot::Y), + ], + }); + steps.push(PlanStep::Anchor(PlanAnchor::AfterEmbed)); + for (i, kind) in s.schedule.iter().enumerate() { + let p = format!("encoder.layers.{i}"); + let suffix = kind.kernel_suffix(); + steps.push(PlanStep::Launch { + kernel: "qkv".into(), + args: vec![ + scratch(ScratchSlot::Y), + weight(format!("{p}.attn.Wqkv.weight")), + weight(format!("zeros.{}", 3 * h)), + scratch(ScratchSlot::Qkv), + ], + }); + steps.push(PlanStep::Launch { + kernel: if opts.original_rope { + "rope_original".into() + } else { + "rope".into() + }, + args: vec![ + scratch(ScratchSlot::Qkv), + weight(format!("rope_{suffix}_cos")), + weight(format!("rope_{suffix}_sin")), + ], + }); + steps.push(PlanStep::Launch { + kernel: format!("attn_{suffix}"), + args: vec![ + scratch(ScratchSlot::Qkv), + scratch(ScratchSlot::Lens), + scratch(ScratchSlot::O), + ], + }); + steps.push(PlanStep::Launch { + kernel: "out".into(), + args: vec![ + scratch(ScratchSlot::O), + weight(format!("{p}.attn.Wo.weight")), + weight(zero_bias.clone()), + scratch(ScratchSlot::Y), + ], + }); + steps.push(PlanStep::Launch { + kernel: "addln".into(), + args: vec![ + scratch(ScratchSlot::X), + scratch(ScratchSlot::Y), + weight(format!("{p}.mlp_norm.weight")), + weight(zero_bias.clone()), + scratch(ScratchSlot::Y), + ], + }); + steps.push(PlanStep::Launch { + kernel: "geglu".into(), + args: vec![ + scratch(ScratchSlot::Y), + weight(format!("{p}.mlp.Wi.weight")), + scratch(ScratchSlot::Gate), + ], + }); + steps.push(PlanStep::Launch { + kernel: "down".into(), + args: vec![ + scratch(ScratchSlot::Gate), + weight(format!("{p}.mlp.Wo.weight")), + weight(zero_bias.clone()), + scratch(ScratchSlot::Y), + ], + }); + let next_norm = if i + 1 < s.schedule.len() { + format!("encoder.layers.{}.attn_norm.weight", i + 1) + } else { + "encoder.final_norm.weight".into() + }; + steps.push(PlanStep::Launch { + kernel: "addln".into(), + args: vec![ + scratch(ScratchSlot::X), + scratch(ScratchSlot::Y), + weight(next_norm), + weight(zero_bias.clone()), + scratch(ScratchSlot::Y), + ], + }); + steps.push(PlanStep::Anchor(PlanAnchor::LayerEnd(i))); + } + steps.push(PlanStep::Anchor(PlanAnchor::AfterEncoder)); + if s.type_embedding_rows.is_some() { + steps.push(PlanStep::Launch { + kernel: "type".into(), + args: vec![ + scratch(ScratchSlot::Y), + weight("type_emb.weight".into()), + scratch(ScratchSlot::Types), + scratch(ScratchSlot::X), + ], + }); + } + steps +} diff --git a/src/models/modernbert/src/structure.rs b/src/models/modernbert/src/structure.rs new file mode 100644 index 00000000..e02d7828 --- /dev/null +++ b/src/models/modernbert/src/structure.rs @@ -0,0 +1,201 @@ +use anyhow::{Result, ensure}; +use std::collections::BTreeMap; + +/// Bumped when the canonical serialization below changes meaning. +pub const SCHEMA_VERSION: usize = 1; + +/// One encoder layer's attention mixer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AttentionKind { + Full, + Sliding, +} + +impl AttentionKind { + /// The name checkpoint configs spell in `layer_types`. + pub fn as_str(self) -> &'static str { + match self { + Self::Full => "full_attention", + Self::Sliding => "sliding_attention", + } + } + + /// The suffix the family CUDA kernels carry per attention kind. + pub fn kernel_suffix(self) -> &'static str { + match self { + Self::Full => "full", + Self::Sliding => "local", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "full_attention" => Ok(Self::Full), + "sliding_attention" => Ok(Self::Sliding), + other => anyhow::bail!("unknown attention kind {other:?}"), + } + } +} + +/// The only activation the family kernels implement. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Activation { + Gelu, +} + +/// YaRN frequency scaling applied on top of a base theta. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct YarnScaling { + pub factor: f64, + pub original_max_positions: usize, + pub beta_fast: f64, + pub beta_slow: f64, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RopeSpec { + pub theta: f64, + pub scaling: Option, +} + +/// Where one structure field's value came from. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum FieldOrigin { + Source(String), + Default(&'static str), + Derived(String), +} + +/// The ModernBERT encoder structure: family semantics and math geometry +/// with per-field provenance. Weight names, kernels, and buffers are not +/// part of the structure; they are derived in [`crate::geometry`] and +/// [`crate::plan`]. +#[derive(Debug, Clone, PartialEq)] +pub struct ModernBertStructure { + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub activation: Activation, + pub norm_eps: f32, + pub num_attention_heads: usize, + pub head_dim: usize, + pub schedule: Vec, + pub local_window: usize, + pub rope_full: RopeSpec, + pub rope_sliding: RopeSpec, + pub vocab_size: usize, + pub type_embedding_rows: Option, + pub max_positions: usize, + pub provenance: BTreeMap, +} + +impl ModernBertStructure { + /// Family invariants. Every failure names the field and its value. + pub fn validate(&self) -> Result<()> { + ensure!( + self.num_hidden_layers > 0, + "num_hidden_layers is {}", + self.num_hidden_layers + ); + ensure!( + self.num_attention_heads > 0 + && self.num_attention_heads * self.head_dim == self.hidden_size, + "num_attention_heads {} with head_dim {} does not form hidden_size {}", + self.num_attention_heads, + self.head_dim, + self.hidden_size + ); + ensure!( + self.head_dim.is_multiple_of(2), + "head_dim {} has no rotary half", + self.head_dim + ); + ensure!( + self.schedule.len() == self.num_hidden_layers, + "schedule holds {} entries for {} layers", + self.schedule.len(), + self.num_hidden_layers + ); + ensure!( + self.intermediate_size > 0, + "intermediate_size is {}", + self.intermediate_size + ); + ensure!( + self.local_window > 0, + "local_attention is {}", + self.local_window + ); + ensure!(self.vocab_size > 0, "vocab_size is {}", self.vocab_size); + ensure!( + self.norm_eps > 0.0 && self.norm_eps.is_finite(), + "norm_eps is {}", + self.norm_eps + ); + ensure!( + self.max_positions > 0, + "max_position_embeddings is {}", + self.max_positions + ); + for (name, rope) in [("full", &self.rope_full), ("sliding", &self.rope_sliding)] { + ensure!( + rope.theta.is_finite() && rope.theta > 0.0, + "{name} rope_theta is {}", + rope.theta + ); + if let Some(YarnScaling { + factor, + original_max_positions, + beta_fast, + beta_slow, + }) = rope.scaling + { + ensure!(factor > 1.0, "yarn factor is {factor}"); + ensure!( + original_max_positions > 0, + "yarn original_max_position_embeddings is {original_max_positions}" + ); + ensure!( + beta_slow > 0.0 && beta_fast > beta_slow, + "yarn beta_fast {beta_fast} / beta_slow {beta_slow}" + ); + } + } + Ok(()) + } + + /// Canonical, sorted-key JSON body (provenance excluded) used for + /// identity digests. Same structure always serializes identically. + pub fn canonical_json(&self) -> serde_json::Value { + let rope = |r: &RopeSpec| { + let mut value = serde_json::json!({ "theta": r.theta }); + if let Some(y) = &r.scaling { + value["scaling"] = serde_json::json!({ + "factor": y.factor, + "original_max_positions": y.original_max_positions, + "beta_fast": y.beta_fast, + "beta_slow": y.beta_slow, + }); + } + value + }; + serde_json::json!({ + "schema_version": SCHEMA_VERSION, + "family": "modernbert", + "activation": "gelu", + "hidden_size": self.hidden_size, + "intermediate_size": self.intermediate_size, + "num_hidden_layers": self.num_hidden_layers, + "norm_eps": self.norm_eps, + "num_attention_heads": self.num_attention_heads, + "head_dim": self.head_dim, + "schedule": self.schedule.iter().map(|k| k.as_str()).collect::>(), + "local_window": self.local_window, + "rope_full": rope(&self.rope_full), + "rope_sliding": rope(&self.rope_sliding), + "vocab_size": self.vocab_size, + "type_embedding_rows": self.type_embedding_rows, + "max_positions": self.max_positions, + }) + } +} diff --git a/src/models/qwen3_5/native/Cargo.toml b/src/models/qwen3_5/native/Cargo.toml index 7db6a25f..dd09a36b 100644 --- a/src/models/qwen3_5/native/Cargo.toml +++ b/src/models/qwen3_5/native/Cargo.toml @@ -13,6 +13,7 @@ image = { version = "0.25", default-features = false, features = ["png", "jpeg"] libloading = "0.8" memmap2 = "0.9.9" safetensors = "0.8.0" +sha2 = "0.10" serde = { version = "1", features = ["derive"] } serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order", "raw_value"] } @@ -24,6 +25,13 @@ path = "../../../../tests/qwen3_5/config.rs" name = "kernels" path = "../../../../tests/qwen3_5/kernels.rs" +[dev-dependencies] +tempfile = "3" + +[[test]] +name = "structure" +path = "../../../../tests/qwen3_5/structure.rs" + [[test]] name = "graph" path = "../../../../tests/qwen3_5/graph.rs" diff --git a/src/models/qwen3_5/native/src/cuda.rs b/src/models/qwen3_5/native/src/cuda.rs index 612d49da..4a42f85a 100644 --- a/src/models/qwen3_5/native/src/cuda.rs +++ b/src/models/qwen3_5/native/src/cuda.rs @@ -13,6 +13,10 @@ pub use reference::Reference; /// `CS1_ABI_VERSION` in ops.h. const ABI_VERSION: u32 = 9; +// The structure side records the same version in identity records; if +// one bumps without the other, runtime state keyed by structure +// identity would silently disagree with the loaded bundle. +const _: () = assert!(ABI_VERSION == crate::structure::ABI_VERSION as u32); pub const LIBRARY: &str = "libqwen3_5_cuda.so"; /// A `cudaStream_t`. diff --git a/src/models/qwen3_5/native/src/lib.rs b/src/models/qwen3_5/native/src/lib.rs index 14fceae4..3244389b 100644 --- a/src/models/qwen3_5/native/src/lib.rs +++ b/src/models/qwen3_5/native/src/lib.rs @@ -5,6 +5,7 @@ pub mod json; pub mod model; pub mod prefix; pub mod prefix_cache; +pub mod structure; pub mod inputs; diff --git a/src/models/qwen3_5/native/src/model.rs b/src/models/qwen3_5/native/src/model.rs index 05f8c6cb..44a803a7 100644 --- a/src/models/qwen3_5/native/src/model.rs +++ b/src/models/qwen3_5/native/src/model.rs @@ -93,135 +93,38 @@ pub struct Config { impl Config { pub fn load(dir: &Path) -> Result { - let path = dir.join("config.json"); - let root: Json = serde_json::from_str( - &std::fs::read_to_string(&path).with_context(|| format!("{}", path.display()))?, - )?; - let c = root.get("text_config").unwrap_or(&root); - let int = |k: &str| { - c[k].as_u64() - .map(|v| v as usize) - .with_context(|| format!("config.json: `{k}` is missing")) - }; - let rope = &c["rope_parameters"]; - let partial = rope["partial_rotary_factor"] - .as_f64() - .or(c["partial_rotary_factor"].as_f64()) - .unwrap_or(1.0); - let head_dim = int("head_dim")?; - let full_attention = c["layer_types"] - .as_array() - .context("config.json: `layer_types` is missing")? - .iter() - .map(|t| match t.as_str() { - Some("full_attention") => Ok(true), - Some("linear_attention") => Ok(false), - other => bail!("unknown layer type {other:?}"), - }) - .collect::>>()?; - let sections = rope - .get("mrope_section") - .cloned() - .unwrap_or(serde_json::json!([11, 11, 10])); - let sections = sections - .as_array() - .context("mrope_section must be an array")? - .iter() - .map(|v| { - v.as_u64() - .and_then(|n| usize::try_from(n).ok()) - .context("mrope_section must contain integers") - }) - .collect::>>()?; - let mrope_section: [usize; 3] = sections - .try_into() - .map_err(|_| anyhow::anyhow!("mrope_section must have three entries"))?; - let image_token_id = root - .get("image_token_id") - .map(|v| { - v.as_u64() - .and_then(|id| u32::try_from(id).ok()) - .context("image_token_id must be a u32") - }) - .transpose()?; - let cfg = Config { - hidden: int("hidden_size")?, - intermediate: int("intermediate_size")?, - eps: c["rms_norm_eps"].as_f64().context("rms_norm_eps")? as f32, - heads: int("num_attention_heads")?, - kv_heads: int("num_key_value_heads")?, - head_dim, - rotary_half: (head_dim as f64 * partial) as usize / 2, - rope_theta: rope["rope_theta"] - .as_f64() - .or(c["rope_theta"].as_f64()) - .context("rope_theta")?, - mrope_section, - max_positions: int("max_position_embeddings")?, - image_token_id, - lin_k_heads: int("linear_num_key_heads")?, - lin_v_heads: int("linear_num_value_heads")?, - lin_k_dim: int("linear_key_head_dim")?, - lin_v_dim: int("linear_value_head_dim")?, - full_attention, - }; - // What the kernels implement. - ensure!( - cfg.full_attention.len() == int("num_hidden_layers")?, - "layer_types does not match num_hidden_layers" - ); - ensure!(c["hidden_act"] == "silu", "hidden_act is not silu"); - ensure!( - c["attn_output_gate"].as_bool().unwrap_or(true), - "attention without the output gate" - ); - ensure!( - c["attention_bias"].as_bool() != Some(true), - "attention with bias" - ); - ensure!( - rope["rope_type"].as_str().unwrap_or("default") == "default", - "rope type {}", - rope["rope_type"] - ); - ensure!(int("linear_conv_kernel_dim")? == 4, "conv kernel is not 4"); - ensure!( - cfg.head_dim == 256 && cfg.lin_k_dim == 128 && cfg.lin_v_dim == 128, - "head dims {} / {} / {}", - cfg.head_dim, - cfg.lin_k_dim, - cfg.lin_v_dim - ); - ensure!(cfg.rotary_half == 32, "{} rotary dims", 2 * cfg.rotary_half); - ensure!( - cfg.rope_theta.is_finite() && cfg.rope_theta > 0.0, - "invalid rope_theta" - ); - ensure!( - rope["mrope_interleaved"].as_bool() != Some(false), - "non-interleaved mrope is unsupported" - ); - ensure!( - cfg.mrope_section + let structure = crate::structure::read_structure(dir)?; + crate::structure::check_abi(&structure)?; + Ok(Self::from_structure(structure)) + } + + /// Build the runtime config from an already-validated structure; the + /// seam a consumer with its own bundle layout (for example a Vela + /// hybrid member whose Qwen fields live at `backbone_config`) uses + /// instead of reading a checkpoint directory. + pub fn from_structure(s: crate::structure::QwenHybridStructure) -> Self { + Self { + hidden: s.hidden, + intermediate: s.intermediate, + eps: s.eps, + full_attention: s + .layer_kinds .iter() - .try_fold(0usize, |sum, &x| sum.checked_add(x)) - == Some(cfg.rotary_half), - "mrope sections must sum to rotary_half" - ); - ensure!( - cfg.max_positions > 0 && cfg.max_positions <= (1 << 24), - "unsupported max_position_embeddings" - ); - ensure!( - cfg.kv_heads > 0 && cfg.heads.is_multiple_of(cfg.kv_heads), - "attention heads" - ); - ensure!( - cfg.lin_k_heads > 0 && cfg.lin_v_heads.is_multiple_of(cfg.lin_k_heads), - "linear attention heads" - ); - ensure!(cfg.hidden.is_multiple_of(8), "hidden size"); - Ok(cfg) + .map(|k| *k == crate::structure::LayerKind::FullAttention) + .collect(), + heads: s.heads, + kv_heads: s.kv_heads, + head_dim: s.head_dim, + rotary_half: s.rotary_half, + rope_theta: s.rope_theta, + mrope_section: s.mrope_section, + max_positions: s.max_positions, + image_token_id: s.image_token_id, + lin_k_heads: s.lin_k_heads, + lin_v_heads: s.lin_v_heads, + lin_k_dim: s.lin_k_dim, + lin_v_dim: s.lin_v_dim, + } } fn key_dim(&self) -> usize { diff --git a/src/models/qwen3_5/native/src/structure.rs b/src/models/qwen3_5/native/src/structure.rs new file mode 100644 index 00000000..c2be71f3 --- /dev/null +++ b/src/models/qwen3_5/native/src/structure.rs @@ -0,0 +1,428 @@ +//! The Qwen3.5 hybrid family layer: the source adapter that reads a +//! checkpoint config, the structure it produces, the frozen ABI-9 +//! capability set the CUDA kernels implement, and the prefix-state +//! layout derived from a structure. Parsing and capability live apart +//! so both are checkable without a device or a checkpoint. + +use anyhow::{Context, Result, bail, ensure}; +use serde_json::Value as Json; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; +use std::path::Path; + +pub const SCHEMA_VERSION: usize = 1; +pub const FAMILY: &str = "qwen3_5_hybrid"; +pub const BACKEND: &str = "qwen3_5_cuda"; +pub const ABI_VERSION: usize = 9; + +const BF16: usize = 2; +const F32: usize = 4; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LayerKind { + FullAttention, + GatedDeltaNet, +} + +/// The hybrid structure: family semantics and math geometry with +/// per-field provenance. The multimodal behavior around it (image +/// token, M-RoPE positions) stays opaque and is carried, not +/// parameterized. +#[derive(Debug, Clone, PartialEq)] +pub struct QwenHybridStructure { + pub hidden: usize, + pub intermediate: usize, + pub eps: f32, + /// Ordered per-layer mixer kinds; index 0 is the first layer. + pub layer_kinds: Vec, + pub heads: usize, + pub kv_heads: usize, + pub head_dim: usize, + /// Half the number of rotary dims (rotate_half pairs i with i + half). + pub rotary_half: usize, + pub rope_theta: f64, + pub mrope_section: [usize; 3], + pub max_positions: usize, + pub image_token_id: Option, + pub lin_k_heads: usize, + pub lin_v_heads: usize, + pub lin_k_dim: usize, + pub lin_v_dim: usize, + /// Source-declared family behavior the kernels pin; the gate + /// rejects anything but the frozen combination. + pub hidden_act: Option, + pub attn_output_gate: bool, + pub attention_bias: bool, + pub rope_type: String, + pub conv_kernel_dim: usize, + pub mrope_interleaved: bool, + pub provenance: BTreeMap, +} + +impl QwenHybridStructure { + pub fn full_attention_layers(&self) -> usize { + self.layer_kinds + .iter() + .filter(|k| **k == LayerKind::FullAttention) + .count() + } + + pub fn gdn_layers(&self) -> usize { + self.layer_kinds.len() - self.full_attention_layers() + } + + pub fn key_dim(&self) -> usize { + self.lin_k_heads * self.lin_k_dim + } + + pub fn value_dim(&self) -> usize { + self.lin_v_heads * self.lin_v_dim + } + + /// Width of the fused Gated DeltaNet input projection (conv q/k/v, + /// gate/value and beta rows stacked). + pub fn gdn_in_width(&self) -> usize { + 2 * self.key_dim() + self.value_dim() + self.value_dim() + 2 * self.lin_v_heads + } + + fn canonical_json(&self) -> Json { + let num = |v: usize| Json::String(v.to_string()); + let mut object = serde_json::Map::new(); + object.insert("schema_version".into(), num(SCHEMA_VERSION)); + object.insert("family".into(), Json::String(FAMILY.into())); + object.insert("hidden".into(), num(self.hidden)); + object.insert("intermediate".into(), num(self.intermediate)); + object.insert("eps".into(), Json::String(self.eps.to_string())); + object.insert( + "layer_kinds".into(), + Json::Array( + self.layer_kinds + .iter() + .map(|k| { + Json::String( + match k { + LayerKind::FullAttention => "full_attention", + LayerKind::GatedDeltaNet => "linear_attention", + } + .into(), + ) + }) + .collect(), + ), + ); + object.insert("heads".into(), num(self.heads)); + object.insert("kv_heads".into(), num(self.kv_heads)); + object.insert("head_dim".into(), num(self.head_dim)); + object.insert("rotary_half".into(), num(self.rotary_half)); + object.insert( + "rope_theta".into(), + Json::String(self.rope_theta.to_string()), + ); + object.insert( + "mrope_section".into(), + Json::Array( + self.mrope_section + .iter() + .map(|s| Json::String(s.to_string())) + .collect(), + ), + ); + object.insert("max_positions".into(), num(self.max_positions)); + object.insert( + "image_token_id".into(), + self.image_token_id + .map(|id| Json::String(id.to_string())) + .unwrap_or(Json::Null), + ); + object.insert("lin_k_heads".into(), num(self.lin_k_heads)); + object.insert("lin_v_heads".into(), num(self.lin_v_heads)); + object.insert("lin_k_dim".into(), num(self.lin_k_dim)); + object.insert("lin_v_dim".into(), num(self.lin_v_dim)); + object.insert( + "hidden_act".into(), + self.hidden_act + .clone() + .map(Json::String) + .unwrap_or(Json::Null), + ); + object.insert( + "attn_output_gate".into(), + Json::String(self.attn_output_gate.to_string()), + ); + object.insert( + "attention_bias".into(), + Json::String(self.attention_bias.to_string()), + ); + object.insert("rope_type".into(), Json::String(self.rope_type.clone())); + object.insert("conv_kernel_dim".into(), num(self.conv_kernel_dim)); + object.insert( + "mrope_interleaved".into(), + Json::String(self.mrope_interleaved.to_string()), + ); + Json::Object(object) + } +} + +/// The source adapter: reads `config.json` from a checkpoint directory, +/// preferring `text_config`, exactly as the loader always has. No +/// capability judgment happens here. +pub fn read_structure(dir: &Path) -> Result { + let path = dir.join("config.json"); + let root: Json = serde_json::from_str( + &std::fs::read_to_string(&path).with_context(|| format!("{}", path.display()))?, + )?; + let c = root.get("text_config").unwrap_or(&root); + let int = |k: &str| { + c[k].as_u64() + .map(|v| v as usize) + .with_context(|| format!("config.json: `{k}` is missing")) + }; + let rope = &c["rope_parameters"]; + let partial = rope["partial_rotary_factor"] + .as_f64() + .or(c["partial_rotary_factor"].as_f64()) + .unwrap_or(1.0); + let head_dim = int("head_dim")?; + let layer_kinds = c["layer_types"] + .as_array() + .context("config.json: `layer_types` is missing")? + .iter() + .map(|t| match t.as_str() { + Some("full_attention") => Ok(LayerKind::FullAttention), + Some("linear_attention") => Ok(LayerKind::GatedDeltaNet), + other => bail!("unknown layer type {other:?}"), + }) + .collect::>>()?; + let sections = rope + .get("mrope_section") + .cloned() + .unwrap_or(serde_json::json!([11, 11, 10])); + let sections = sections + .as_array() + .context("mrope_section must be an array")? + .iter() + .map(|v| { + v.as_u64() + .and_then(|n| usize::try_from(n).ok()) + .context("mrope_section must contain integers") + }) + .collect::>>()?; + let mrope_section: [usize; 3] = sections + .try_into() + .map_err(|_| anyhow::anyhow!("mrope_section must have three entries"))?; + let image_token_id = root + .get("image_token_id") + .map(|v| { + v.as_u64() + .and_then(|id| u32::try_from(id).ok()) + .context("image_token_id must be a u32") + }) + .transpose()?; + let num_hidden_layers = int("num_hidden_layers")?; + let mut provenance = BTreeMap::new(); + for field in [ + "hidden", + "intermediate", + "eps", + "layer_kinds", + "heads", + "kv_heads", + "head_dim", + "rotary_half", + "rope_theta", + "mrope_section", + "max_positions", + "lin_k_heads", + "lin_v_heads", + "lin_k_dim", + "lin_v_dim", + ] { + let json_key = match field { + "hidden" => "hidden_size", + "intermediate" => "intermediate_size", + "eps" => "rms_norm_eps", + "layer_kinds" => "layer_types", + "heads" => "num_attention_heads", + "kv_heads" => "num_key_value_heads", + "max_positions" => "max_position_embeddings", + "lin_k_heads" => "linear_num_key_heads", + "lin_v_heads" => "linear_num_value_heads", + "lin_k_dim" => "linear_key_head_dim", + "lin_v_dim" => "linear_value_head_dim", + "rotary_half" => "rope_parameters.partial_rotary_factor", + _ => field, + }; + provenance.insert( + field.to_string(), + format!( + "config.json#{}{json_key}", + if root.get("text_config").is_some() { + "text_config." + } else { + "" + } + ), + ); + } + let hidden_act = c["hidden_act"].as_str().map(str::to_string); + let attn_output_gate = c["attn_output_gate"].as_bool().unwrap_or(true); + let attention_bias = c["attention_bias"].as_bool().unwrap_or(false); + let rope_type = rope["rope_type"].as_str().unwrap_or("default").to_string(); + let mrope_interleaved = rope["mrope_interleaved"].as_bool().unwrap_or(true); + // The layer-count check keeps its base position: after every field + // extraction, before the gate's other checks. linear_conv_kernel_dim + // is read after it, as in the base loader where the gate extracted + // it at its own step. + ensure!( + layer_kinds.len() == num_hidden_layers, + "layer_types does not match num_hidden_layers" + ); + let conv_kernel_dim = int("linear_conv_kernel_dim")?; + Ok(QwenHybridStructure { + hidden_act, + attn_output_gate, + attention_bias, + rope_type, + conv_kernel_dim, + mrope_interleaved, + hidden: int("hidden_size")?, + intermediate: int("intermediate_size")?, + eps: c["rms_norm_eps"].as_f64().context("rms_norm_eps")? as f32, + heads: int("num_attention_heads")?, + kv_heads: int("num_key_value_heads")?, + head_dim, + rotary_half: (head_dim as f64 * partial) as usize / 2, + rope_theta: rope["rope_theta"] + .as_f64() + .or(c["rope_theta"].as_f64()) + .context("rope_theta")?, + mrope_section, + max_positions: int("max_position_embeddings")?, + image_token_id, + lin_k_heads: int("linear_num_key_heads")?, + lin_v_heads: int("linear_num_value_heads")?, + lin_k_dim: int("linear_key_head_dim")?, + lin_v_dim: int("linear_value_head_dim")?, + layer_kinds, + provenance, + }) +} + +/// The frozen ABI-9 kernel capability set: exactly what the CUDA +/// kernels implement, no more. Widening anything here is a kernel/ABI +/// decision, not a loader change. Every rejection reports the actual +/// value. +pub fn check_abi(s: &QwenHybridStructure) -> Result<()> { + ensure!( + s.hidden_act.as_deref() == Some("silu"), + "hidden_act is not silu" + ); + ensure!(s.attn_output_gate, "attention without the output gate"); + ensure!(!s.attention_bias, "attention with bias"); + ensure!(s.rope_type == "default", "rope type {:?}", s.rope_type); + ensure!(s.conv_kernel_dim == 4, "conv kernel is not 4"); + ensure!( + s.head_dim == 256 && s.lin_k_dim == 128 && s.lin_v_dim == 128, + "head dims {} / {} / {}", + s.head_dim, + s.lin_k_dim, + s.lin_v_dim + ); + ensure!(s.rotary_half == 32, "{} rotary dims", 2 * s.rotary_half); + ensure!( + s.rope_theta.is_finite() && s.rope_theta > 0.0, + "invalid rope_theta" + ); + ensure!(s.mrope_interleaved, "non-interleaved mrope is unsupported"); + ensure!( + s.mrope_section + .iter() + .try_fold(0usize, |sum, &x| sum.checked_add(x)) + == Some(s.rotary_half), + "mrope sections must sum to rotary_half" + ); + ensure!( + s.max_positions > 0 && s.max_positions <= (1 << 24), + "unsupported max_position_embeddings" + ); + ensure!( + s.kv_heads > 0 && s.heads.is_multiple_of(s.kv_heads), + "attention heads" + ); + ensure!( + s.lin_k_heads > 0 && s.lin_v_heads.is_multiple_of(s.lin_k_heads), + "linear attention heads" + ); + ensure!(s.hidden.is_multiple_of(8), "hidden size"); + Ok(()) +} + +/// Per-layer prefix-state byte layout for one token length, derived +/// from the structure: per full-attention layer the post-prep K and +/// raw V rows, per Gated DeltaNet layer the float32 recurrent state +/// and the three pre-conv projection columns. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PrefixLayout { + pub len: usize, + pub full_layers: usize, + pub gdn_layers: usize, + /// One full-attention layer's K-plus-V bytes (`2 * len * kv row`). + pub full_kv_bytes: usize, + /// One GDN layer's float32 `[H, K, V]` state bytes. + pub gdn_state_bytes: usize, + /// One GDN layer's three conv-tail columns. + pub conv_tail_bytes: usize, +} + +impl PrefixLayout { + pub fn of(s: &QwenHybridStructure, len: usize) -> Self { + let kvrow = s.kv_heads * s.head_dim * BF16; + let state_floats = s.lin_v_heads * s.lin_k_dim * s.lin_v_dim; + Self { + len, + full_layers: s.full_attention_layers(), + gdn_layers: s.gdn_layers(), + full_kv_bytes: 2 * len * kvrow, + gdn_state_bytes: state_floats * F32, + conv_tail_bytes: 3 * s.gdn_in_width() * BF16, + } + } + + /// Total device bytes of one prefix of `len` tokens. + pub fn total_bytes(&self) -> usize { + self.full_layers * self.full_kv_bytes + + self.gdn_layers * (self.gdn_state_bytes + self.conv_tail_bytes) + } +} + +/// Stable identity of one structure executing on the ABI-9 backend. +/// Any prefix-cache record must carry it; a change means the state +/// must be re-captured, never reused. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QwenIdentity { + pub schema_version: usize, + pub family: &'static str, + pub backend: &'static str, + pub abi_version: usize, + pub digest: String, +} + +pub fn structure_identity(s: &QwenHybridStructure) -> QwenIdentity { + let mut body = serde_json::Map::new(); + body.insert( + "schema_version".into(), + Json::String(SCHEMA_VERSION.to_string()), + ); + body.insert("family".into(), Json::String(FAMILY.into())); + body.insert("backend".into(), Json::String(BACKEND.into())); + body.insert("abi_version".into(), Json::String(ABI_VERSION.to_string())); + body.insert("structure".into(), s.canonical_json()); + let body = Json::Object(body).to_string(); + QwenIdentity { + schema_version: SCHEMA_VERSION, + family: FAMILY, + backend: BACKEND, + abi_version: ABI_VERSION, + digest: format!("{:x}", Sha256::digest(body.as_bytes())), + } +} diff --git a/src/models/vela/README.md b/src/models/vela/README.md index 74f2c4bb..b40d7950 100644 --- a/src/models/vela/README.md +++ b/src/models/vela/README.md @@ -1,5 +1,9 @@ # Vela 2.0 reference contracts (0.3B, 0.8B hybrid) +This directory contains the System1-Omni integration slices for +[Vela 2.0 0.3B](https://huggingface.co/vllm-sr/Vela-2.0-0.3B) +# Vela 2.0 reference contracts (0.3B, 0.8B hybrid) + This directory contains the System1-Omni integration slices for [Vela 2.0 0.3B](https://huggingface.co/vllm-sr/Vela-2.0-0.3B) (`vela2-unified`) and @@ -10,7 +14,11 @@ question kinds (`choice`, `noul`, `score`, `span`, `set`) on `/v1/systemone`, one pinned member per worker process. The 0.3B worker's contract freezes the text/token_ids input, the position-level final-hidden readout contract, the error/length -boundaries and the parity baseline. +boundaries and the parity baseline for later ModernBERT-profile work; +the CPU half of that native route — the 0.3B ModernBERT profile +registration, the checkpoint census and the YaRN tables — lives in +[`native/`](native/README.md), whose `cuda`-gated executor runs the +encoder on the regenerated bundle. The shared Vela 2.0 wire module is deliberately shaped for the family: the envelope admission (`protocol`), the pinned-runtime error mapping @@ -31,7 +39,7 @@ implemented. | Runtime | the checkpoint's own flat files: `vela2_inference.py`, `modeling_vela2.py`, `configuration_vela2.py`, `vela2_serve.py` | the same four modules, decoder-specific | | Calibration | the package's own `calibration.json` (never cross-copied between members — shipped thresholds differ) | same rule; the side-loaded broad head needs its `broad_head` block | | License | Apache-2.0; the inherited tokenizer material carries Gemma Terms of Use | Apache-2.0; Qwen3.5-0.8B lineage notices ship in the package | -| Runtime dependencies | [recipe/vela/requirements.txt](../../recipe/vela/requirements.txt) (torch, transformers, safetensors, tokenizers, numpy) | same | +| Runtime dependencies | [recipe/vela/requirements.txt](../../../recipe/vela/requirements.txt) (torch, transformers, safetensors, tokenizers, numpy) | same | The 0.3B package ships no manifest, so its gate checks the config's structural identity (model type, architecture, `max_length` 8,192, the @@ -224,4 +232,4 @@ the shared worker-contract issue, and nothing is posted upstream. flash-linear-attention) is the oracle setting; GPU bf16-autocast is a deployment option outside these slices' evidence. - The serving procedure, hardware requirements and the recorded - evidence pointer are [recipe/vela/reference.md](../../recipe/vela/reference.md). + evidence pointer are [recipe/vela/reference.md](../../../recipe/vela/reference.md). diff --git a/src/models/vela/engine.py b/src/models/vela/engine.py index ac2f316a..f9dce359 100644 --- a/src/models/vela/engine.py +++ b/src/models/vela/engine.py @@ -208,10 +208,16 @@ def verify_no_foreign_runtime(root: Path) -> None: for name in RUNTIME_MODULES: loaded = sys.modules.get(name) - if loaded is not None and Path(loaded.__file__).resolve().parent != Path(root).resolve(): + if loaded is None: + continue + origin = getattr(loaded, "__file__", None) + if origin is None or Path(origin).resolve().parent != Path(root).resolve(): + # A module without a file is a namespace package or a + # builtin; either way it is not this package's flat module. raise ValueError( - f"a {name} module from another package directory is already " - "imported in this process; serve one package per process" + f"a {name} module from another package directory (or a namespace " + "package) is already imported in this process; " + "serve one package per process" ) diff --git a/src/models/vela/native/Cargo.toml b/src/models/vela/native/Cargo.toml new file mode 100644 index 00000000..6f0b03ca --- /dev/null +++ b/src/models/vela/native/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "omni-vela-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "Vela 2.0 0.3B ModernBERT profile on the family core" + +[dependencies] +anyhow = "1" +half = { version = "2", optional = true } +memmap2 = "0.9" +omni-cuda = { path = "../../../backends/cuda", optional = true } +omni-modernbert = { path = "../../modernbert" } +safetensors = "0.8.0" +serde_json = "1" + +[features] +cuda = ["dep:omni-cuda", "dep:half"] + +[dev-dependencies] +tempfile = "3" + +[[test]] +name = "vela_profile" +path = "../../../../tests/vela/native/profile.rs" + +[[test]] +name = "vela_weights" +path = "../../../../tests/vela/native/weights.rs" + +[[test]] +name = "vela_rope" +path = "../../../../tests/vela/native/rope.rs" + +[[test]] +name = "vela_executor" +path = "../../../../tests/vela/native/executor.rs" +required-features = ["cuda"] diff --git a/src/models/vela/native/README.md b/src/models/vela/native/README.md new file mode 100644 index 00000000..1556fa9e --- /dev/null +++ b/src/models/vela/native/README.md @@ -0,0 +1,23 @@ +# Vela 2.0 0.3B native profile + +`omni-vela-native` registers the Vela 2.0 0.3B encoder as the second +ModernBERT family profile and owns its CPU-verifiable half: bundle +config parsing (`config`), the registered profile and its bundle +limits (`profile`), the YaRN-scaled rotary tables (`rope`), and the +model-local checkpoint inventory with strict f32 accounting (`weights`). + +Weight upload and encoder execution live behind the crate's `cuda` +feature (`executor`): the admitted structure runs as the family core +plan against the regenerated shared bundle, whose ABI-2 kernels serve +both family profiles — the Laya 1024-hidden geometry under its +unchanged symbols and this profile's 768-hidden geometry under +`_h768`-suffixed names. Everything CPU-side is validated against the +pinned checkpoint's own files; see the tests under +[`tests/vela/native/`](../../../../tests/vela/native/) and the pinned +census and rotary-reference fixtures they consume. + +The serving path today is the Python reference worker documented in +[`src/models/vela`](../README.md) and +[`recipe/vela/reference.md`](../../../../recipe/vela/reference.md); +its frozen parity fixtures are the comparison baseline this native +route will be judged against. diff --git a/src/models/vela/native/src/config.rs b/src/models/vela/native/src/config.rs new file mode 100644 index 00000000..21361caa --- /dev/null +++ b/src/models/vela/native/src/config.rs @@ -0,0 +1,30 @@ +use anyhow::{Context, Result, ensure}; +use omni_modernbert::structure::ModernBertStructure; +use omni_modernbert::{ModernBertConfigAdapter, SourceMeta}; +use serde_json::Value as Json; +use std::path::Path; + +/// Reads the pinned bundle layout: one `config.json` whose root is the +/// `vela2-unified` model description with the ModernBERT fields under +/// `encoder_config`. The root's recipe fields (`max_length`, markers, +/// calibration pointers) are worker-side and never structure input. +pub fn read_structure(dir: &Path) -> Result<(ModernBertStructure, SourceMeta)> { + let path = dir.join("config.json"); + let bytes = std::fs::read(&path).with_context(|| format!("{}", path.display()))?; + let root: Json = + serde_json::from_slice(&bytes).with_context(|| format!("parse {}", path.display()))?; + ensure!( + root["model_type"].as_str() == Some("vela2-unified"), + "model_type {:?} is not the vela2-unified 0.3B bundle", + root["model_type"] + ); + let encoder = root + .get("encoder_config") + .context("encoder_config is missing")?; + ensure!( + encoder["model_type"].as_str() == Some("modernbert"), + "encoder_config model_type {:?} is not modernbert", + encoder["model_type"] + ); + ModernBertConfigAdapter::from_value(encoder, "config.json#encoder_config") +} diff --git a/src/models/vela/native/src/executor.rs b/src/models/vela/native/src/executor.rs new file mode 100644 index 00000000..606bb796 --- /dev/null +++ b/src/models/vela/native/src/executor.rs @@ -0,0 +1,250 @@ +//! Encoder execution for the admitted 0.3B profile against the +//! regenerated CUDA bundle. Weight upload, the YaRN rotary tables and +//! the workspace follow the family conventions; the executed step +//! sequence is the family core plan with this profile's `_h768` +//! kernel symbols. The final hidden of the last layer lives in the +//! bundle's bf16 working precision after the final norm; readouts +//! gather requested positions from it. + +use anyhow::{Context, Result, bail, ensure}; +use half::bf16; +use omni_cuda::{Buffer, Cuda, Kernel, Ptr}; +use omni_modernbert::geometry::EncoderScratch; +use omni_modernbert::plan::{self, EncoderPlanOptions, PlanArg, PlanStep, ScratchSlot}; +use omni_modernbert::structure::ModernBertStructure; +use std::{collections::HashMap, path::Path}; + +use crate::config::read_structure; +use crate::profile; +use crate::rope::cos_sin_tables; +use crate::weights::Weights; + +/// Plan kernel name → bundle symbol for the 768-hidden profile. +fn symbol(kernel: &str) -> &str { + match kernel { + "embed" => "embed_h768", + "qkv" => "qkv_h768", + "rope" => "rope_h768", + "attn_full" => "attn_full_h768", + "attn_local" => "attn_local_h768", + "out" => "out_h768", + "addln" => "addln_h768", + "geglu" => "geglu_h768", + "down" => "down_h768", + other => other, + } +} + +#[derive(Clone, Copy)] +enum Arg { + Weight(Ptr), + Scratch(ScratchSlot), +} + +struct Step { + kernel: Kernel, + args: [Arg; 5], + count: usize, +} + +struct Scratch { + ids: Buffer, + lens: Buffer, + x: Buffer, + y: Buffer, + qkv: Buffer, + o: Buffer, + gate: Buffer, +} + +impl Scratch { + fn new(cuda: &Cuda, structure: &ModernBertStructure, batch: usize, len: usize) -> Result { + let sizes = EncoderScratch::new(batch, len, structure); + let mut values = sizes.sizes().into_iter(); + let mut alloc = || { + cuda.alloc(values.next().expect("scratch layout")) + .context("allocate encoder scratch") + }; + Ok(Self { + ids: alloc()?, + lens: alloc()?, + x: alloc()?, + y: alloc()?, + qkv: alloc()?, + o: alloc()?, + gate: alloc()?, + }) + } + + fn buffer(&self, slot: ScratchSlot) -> Result<&Buffer> { + match slot { + ScratchSlot::Ids => Ok(&self.ids), + ScratchSlot::Lens => Ok(&self.lens), + // No admitted profile carries a type embedding; a plan + // asking for that slot is a plan/profile mismatch to fail + // loudly, not a mapping onto the token-id buffer. + ScratchSlot::Types => { + bail!("type-embedding scratch is not supported by the admitted profile") + } + ScratchSlot::X => Ok(&self.x), + ScratchSlot::Y => Ok(&self.y), + ScratchSlot::Qkv => Ok(&self.qkv), + ScratchSlot::O => Ok(&self.o), + ScratchSlot::Gate => Ok(&self.gate), + } + } +} + +pub struct Encoder { + pub structure: ModernBertStructure, + cuda: Cuda, + /// Owns every uploaded tensor; the plan holds raw pointers into it. + _device_tensors: HashMap, + plan: Vec, +} + +fn bytes16(values: &[u16]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +fn bytes32(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() +} + +impl Encoder { + /// Load the pinned bundle (config.json + model.safetensors) and the + /// regenerated CUDA library, admitting the structure first. + /// + /// # Safety + /// The library must be a trusted build artifact supplied by the + /// operator; it is loaded into the process. + pub unsafe fn load(checkpoint: &Path, library: &Path) -> Result { + let (structure, _meta) = read_structure(checkpoint)?; + profile().admit(&structure)?; + let limits = &profile().limits; + let cuda = unsafe { Cuda::load(library)? }; + let source = Weights::open(&checkpoint.join("model.safetensors"))?; + source.validate(&structure)?; + let mut weights = HashMap::new(); + for (name, shape) in omni_modernbert::geometry::encoder_tensors(&structure) { + let values = source.f32(&name, &shape)?; + let data = match crate::weights::storage_dtype(&name) { + "f32" => bytes32(&values), + "f16" => bytes16( + &values + .iter() + .map(|v| half::f16::from_f32(*v).to_bits()) + .collect::>(), + ), + _ => bytes16( + &values + .iter() + .map(|v| bf16::from_f32(*v).to_bits()) + .collect::>(), + ), + }; + let buffer = cuda + .upload(&data) + .with_context(|| format!("upload {name}"))?; + weights.insert(name, buffer); + } + for width in [structure.hidden_size, 3 * structure.hidden_size] { + let name = format!("zeros.{width}"); + weights.insert(name.clone(), cuda.upload(&vec![0; width * 4])?); + } + // Both attention kinds share one theta and one yarn scaling, so + // one table set serves the plan's per-kind names. + let len = limits.len_max.min(structure.max_positions); + let (cos, sin, _tables) = cos_sin_tables(&structure, len)?; + for kind in ["full", "local"] { + weights.insert(format!("rope_{kind}_cos"), cuda.upload(&bytes32(&cos))?); + weights.insert(format!("rope_{kind}_sin"), cuda.upload(&bytes32(&sin))?); + } + let mut plan = Vec::new(); + for step in plan::encoder_plan(&structure, &EncoderPlanOptions::default()) { + let PlanStep::Launch { kernel, args } = step else { + continue; + }; + ensure!(args.len() <= 5, "encoder argument count"); + let mut resolved = [Arg::Weight(std::ptr::null_mut()); 5]; + for (slot, arg) in resolved.iter_mut().zip(&args) { + *slot = match arg { + PlanArg::Weight(name) => Arg::Weight(weights[name].ptr()), + PlanArg::Scratch(s) => Arg::Scratch(*s), + }; + } + plan.push(Step { + kernel: cuda.resolve(symbol(&kernel))?, + args: resolved, + count: args.len(), + }); + } + Ok(Self { + structure, + cuda, + _device_tensors: weights, + plan, + }) + } + + /// Run one sequence and return the final hidden at the requested + /// positions, decoded from the bundle's bf16 working precision. + pub fn encode(&self, ids: &[i64], positions: &[usize]) -> Result>> { + let limits = &profile().limits; + ensure!(!ids.is_empty(), "empty sequence"); + // The embedding kernel indexes the token table unchecked, so an + // id outside [0, vocab_size) would fault on device; refuse it + // here before anything is allocated or launched. + ensure!( + ids.iter() + .all(|id| (0..self.structure.vocab_size as i64).contains(id)), + "token id outside the vocabulary" + ); + let len = ids.len().next_multiple_of(limits.len_step); + limits.check_request(1, len)?; + ensure!( + positions.iter().all(|&p| p < ids.len()), + "position outside the sequence" + ); + let scratch = Scratch::new(&self.cuda, &self.structure, 1, len)?; + let mut staging = vec![0i64; len]; + staging[..ids.len()].copy_from_slice(ids); + scratch.ids.write( + &staging + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + scratch.lens.write( + &vec![ids.len() as i32; 1][..] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + )?; + for step in &self.plan { + let pointers: Vec = step.args[..step.count] + .iter() + .map(|arg| match arg { + Arg::Weight(p) => Ok(*p), + Arg::Scratch(slot) => Ok(scratch.buffer(*slot)?.ptr()), + }) + .collect::>>()?; + // SAFETY: buffers are sized by the family geometry for this + // admitted structure and stay alive through the sync below. + unsafe { step.kernel.launch(&pointers, 1, len)? }; + } + self.cuda.sync()?; + let hidden = scratch.y.read(len * self.structure.hidden_size * 2)?; + let decoded: Vec = hidden + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_bits(u16::from_le_bytes([b[0], b[1]])).to_f32()) + .collect(); + let width = self.structure.hidden_size; + Ok(positions + .iter() + .map(|&p| decoded[p * width..(p + 1) * width].to_vec()) + .collect()) + } +} diff --git a/src/models/vela/native/src/lib.rs b/src/models/vela/native/src/lib.rs new file mode 100644 index 00000000..cf85fbf5 --- /dev/null +++ b/src/models/vela/native/src/lib.rs @@ -0,0 +1,18 @@ +//! The Vela 2.0 0.3B model as a ModernBERT profile on the family core: +//! bundle config parsing, profile admission, the model-local checkpoint +//! inventory, and the YaRN-scaled rotary tables, all CPU-verifiable +//! against the pinned checkpoint's own files. The `cuda` feature adds +//! weight upload and encoder execution against the regenerated bundle. + +pub mod config; +pub mod profile; +pub mod rope; +pub mod weights; + +#[cfg(feature = "cuda")] +pub mod executor; + +pub use config::read_structure; +pub use profile::profile; +pub use rope::{YarnTables, cos_sin_tables, yarn}; +pub use weights::{Weights, head_tensors, storage_dtype}; diff --git a/src/models/vela/native/src/profile.rs b/src/models/vela/native/src/profile.rs new file mode 100644 index 00000000..bf9309e8 --- /dev/null +++ b/src/models/vela/native/src/profile.rs @@ -0,0 +1,88 @@ +use omni_modernbert::gate::{BackendLimits, RegisteredProfile}; +use omni_modernbert::structure::{ + Activation, AttentionKind, FieldOrigin, ModernBertStructure, RopeSpec, YarnScaling, +}; + +/// The registered `vela-2.0-0.3b` profile: the pinned checkpoint's +/// structure (read from `encoder_config`, cross-checked against the +/// recorded tensor census) against the limits of the regenerated CUDA +/// bundle this profile targets — hidden 768 instantiations, YaRN rotary +/// tables, requests packed to 8 192 tokens. The bundle's own ABI is +/// checked when it is loaded. +pub fn profile() -> RegisteredProfile { + RegisteredProfile { + name: "vela-2.0-0.3b", + structure: canonical_structure(), + limits: BackendLimits { + hidden_sizes: vec![768], + batch_max: 4, + len_min: 16, + len_max: 8192, + len_step: 16, + yarn_scaling: true, + }, + backend: "laya_cuda", + abi_version: 2, + } +} + +fn canonical_structure() -> ModernBertStructure { + let mut provenance = std::collections::BTreeMap::new(); + for field in [ + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "norm_eps", + "num_attention_heads", + "head_dim", + "schedule", + "local_window", + "rope_full", + "rope_sliding", + "vocab_size", + "type_embedding_rows", + "max_positions", + ] { + provenance.insert( + field.to_string(), + FieldOrigin::Source("registered profile vela-2.0-0.3b".into()), + ); + } + let yarn = YarnScaling { + factor: 4.0, + original_max_positions: 8192, + beta_fast: 32.0, + beta_slow: 1.0, + }; + ModernBertStructure { + hidden_size: 768, + intermediate_size: 1152, + num_hidden_layers: 22, + activation: Activation::Gelu, + norm_eps: 1e-5, + num_attention_heads: 12, + head_dim: 64, + schedule: (0..22) + .map(|i| { + if i % 3 == 0 { + AttentionKind::Full + } else { + AttentionKind::Sliding + } + }) + .collect(), + local_window: 128, + rope_full: RopeSpec { + theta: 160000.0, + scaling: Some(yarn), + }, + rope_sliding: RopeSpec { + theta: 160000.0, + scaling: Some(yarn), + }, + vocab_size: 256008, + type_embedding_rows: None, + max_positions: 32768, + provenance, + } +} diff --git a/src/models/vela/native/src/rope.rs b/src/models/vela/native/src/rope.rs new file mode 100644 index 00000000..24f3d33c --- /dev/null +++ b/src/models/vela/native/src/rope.rs @@ -0,0 +1,104 @@ +use anyhow::{Context, Result, ensure}; +use omni_modernbert::structure::{ModernBertStructure, RopeSpec, YarnScaling}; + +/// YaRN-derived rotary parameters: the interpolated inverse frequencies +/// and the attention scaling the reference applies beside them. Both +/// attention kinds share one set (the pinned config states one theta +/// and one rope_scaling). +#[derive(Debug, Clone, PartialEq)] +pub struct YarnTables { + /// `rotary_half` entries; row `p` of a cos/sin table is + /// `p * inv_freq[j]` passed through cos/sin. + pub inv_freq: Vec, + /// `1 + 0.1 * ln(factor)` in the reference's default-mscale path. + pub attention_scaling: f64, +} + +/// Computes the YaRN inverse frequencies exactly as the reference +/// runtime does (transformers' yarn initialization): extrapolation and +/// interpolation frequencies blended by a linear ramp whose bounds come +/// from the correction-dimension formula, floored/ceiled under +/// `truncate`, and the default attention scaling. Internals run in +/// f64, outputs narrow to f32. +pub fn yarn(rope: &RopeSpec, rotary_half: usize) -> Result { + let YarnScaling { + factor, + original_max_positions, + beta_fast, + beta_slow, + } = rope.scaling.as_ref().context("rope has no yarn scaling")?; + let base = rope.theta; + let half = rotary_half as f64; + let correction_dim = |rotations: f64| { + (2.0 * half * (*original_max_positions as f64 / (rotations * std::f64::consts::TAU)).ln()) + / (2.0 * base.ln()) + }; + let low = correction_dim(*beta_fast).floor().max(0.0); + // The reference clamps the ramp's upper bound to the head dimension + // (2 * rotary_half) - 1, not to the ramp array's last index: with + // `truncate`, a bound above the array saturates the mask instead of + // rescaling the in-band ramp. Clamping to rotary_half - 1 would + // shrink (high - low) and change every in-band blend. + let high = correction_dim(*beta_slow) + .ceil() + .min(2.0 * rotary_half as f64 - 1.0); + ensure!(high > low, "degenerate yarn ramp bounds {low}..{high}"); + let ramp = |j: f64| ((j - low) / (high - low)).clamp(0.0, 1.0); + let inv_freq = (0..rotary_half) + .map(|j| { + let j = j as f64; + let pos_freq = base.powf(j / half); + let interpolation = 1.0 / (factor * pos_freq); + let extrapolation = 1.0 / pos_freq; + let ramp = ramp(j); + (interpolation * ramp + extrapolation * (1.0 - ramp)) as f32 + }) + .collect(); + Ok(YarnTables { + inv_freq, + attention_scaling: 1.0 + 0.1 * factor.ln(), + }) +} + +/// Full cos/sin tables for positions `[0, len)` in f32, row-major +/// `[len, rotary_half]` each — the layout the family rope kernel +/// consumes. The reference runtime multiplies the YaRN attention +/// scaling into its tables, so these do too; the kernel's own +/// attention scale stays unchanged and must not apply it twice. +pub fn cos_sin_tables( + structure: &ModernBertStructure, + len: usize, +) -> Result<(Vec, Vec, YarnTables)> { + ensure!( + len > 0 && len <= structure.max_positions, + "table length {len}" + ); + let rope = &structure.rope_full; + ensure!( + structure.rope_full == structure.rope_sliding, + "the two attention kinds must share one rope" + ); + let tables = yarn(rope, omni_modernbert::geometry::rotary_half(structure))?; + let half = tables.inv_freq.len(); + // Fail closed on a corrupt structure instead of attempting an + // oversized allocation: `max_positions` is untrusted input until a + // profile admits it. + let elements = len + .checked_mul(half) + .with_context(|| format!("table length {len} overflows the table size"))?; + let mut cos = Vec::new(); + let mut sin = Vec::new(); + cos.try_reserve(elements) + .with_context(|| format!("cannot allocate a {len}x{half} cosine table"))?; + sin.try_reserve(elements) + .with_context(|| format!("cannot allocate a {len}x{half} sine table"))?; + let scaling = tables.attention_scaling as f32; + for pos in 0..len { + for &f in &tables.inv_freq { + let angle = pos as f32 * f; + cos.push(angle.cos() * scaling); + sin.push(angle.sin() * scaling); + } + } + Ok((cos, sin, tables)) +} diff --git a/src/models/vela/native/src/weights.rs b/src/models/vela/native/src/weights.rs new file mode 100644 index 00000000..554f379d --- /dev/null +++ b/src/models/vela/native/src/weights.rs @@ -0,0 +1,144 @@ +use anyhow::{Context, Result, ensure}; +use memmap2::Mmap; +use omni_modernbert::structure::ModernBertStructure; +use safetensors::{Dtype, SafeTensors}; +use std::{collections::BTreeMap, fs::File, path::Path}; + +/// Storage class per tensor name, mirroring the Laya baseline's +/// serving classes: the embedding uploads as f16, every norm, the two +/// calibration scalars and the head projection biases stay f32, and +/// only the GEMM-class projections narrow to bf16 at upload. The +/// pinned checkpoint stores everything f32; the narrowing is a +/// declared numerics choice carried in the profile identity — a +/// learned scalar (`log_tau`) or bias is not GEMM-class, and rounding +/// it to bf16 would surface as unexplained drift in the later parity +/// harness. +pub fn storage_dtype(name: &str) -> &'static str { + if name == "encoder.embeddings.tok_embeddings.weight" { + "f16" + } else if is_f32_class(name) { + "f32" + } else { + "bf16" + } +} + +/// Norms, the two calibration scalars and the head projection biases +/// stay f32 — the Laya baseline keeps its learned temperature and head +/// biases f32 too; a learned scalar or bias is not GEMM-class, and +/// rounding it to bf16 would surface as unexplained drift in the later +/// parity harness. +fn is_f32_class(name: &str) -> bool { + name.contains("norm") + || name == "log_tau" + || name == "log_tau_span" + || (!name.starts_with("encoder.") && name.ends_with(".bias")) +} + +/// The head-side tensors of the pinned checkpoint, recorded from its +/// safetensors header. They belong to the task heads (the readout +/// projections, the classifier MLP, the calibration scalars) and load +/// in that slice; the checkpoint accounting here covers them so every +/// tensor in the file is named exactly once. +pub fn head_tensors() -> Vec<(String, Vec)> { + let h = 768; + let mut tensors = vec![ + ("cls_mlp.0.weight".to_string(), vec![1536, h]), + ("cls_mlp.0.bias".to_string(), vec![1536]), + ("cls_mlp.3.weight".to_string(), vec![1, 1536]), + ("cls_mlp.3.bias".to_string(), vec![1]), + ("norm.weight".to_string(), vec![h]), + ("norm.bias".to_string(), vec![h]), + ]; + for name in ["w_e", "w_o", "w_p", "w_q", "w_t"] { + tensors.push((format!("{name}.weight"), vec![256, h])); + tensors.push((format!("{name}.bias"), vec![256])); + } + tensors.push(("log_tau".to_string(), vec![])); + tensors.push(("log_tau_span".to_string(), vec![])); + tensors +} + +pub struct Weights { + data: Mmap, +} + +impl Weights { + /// The checkpoint must remain immutable while the mapping exists. + pub fn open(path: &Path) -> Result { + let file = + File::open(path).with_context(|| format!("open weights file {}", path.display()))?; + // SAFETY: model files are read-only inputs; no mutable mapping is created. + let data = unsafe { Mmap::map(&file) } + .with_context(|| format!("map weights file {}", path.display()))?; + SafeTensors::deserialize(&data).with_context(|| format!("parse {}", path.display()))?; + Ok(Self { data }) + } + + /// Validate the whole checkpoint against the admitted structure: + /// every encoder tensor must appear with the derived shape and f32 + /// storage, the head inventory must cover the rest exactly, and any + /// missing, extra or duplicate name fails listing both sides. + pub fn validate(&self, structure: &ModernBertStructure) -> Result { + let tensors = SafeTensors::deserialize(&self.data)?; + let mut expected: BTreeMap> = BTreeMap::new(); + for (name, shape) in omni_modernbert::geometry::encoder_tensors(structure) { + ensure!( + expected.insert(name.clone(), shape).is_none(), + "duplicate {name}" + ); + } + for (name, shape) in head_tensors() { + ensure!( + expected.insert(name.clone(), shape).is_none(), + "duplicate {name}" + ); + } + let actual: std::collections::BTreeSet = + tensors.names().into_iter().map(String::from).collect(); + let expected_names: std::collections::BTreeSet = expected.keys().cloned().collect(); + if actual != expected_names { + let missing: Vec<_> = expected_names.difference(&actual).collect(); + let extra: Vec<_> = actual.difference(&expected_names).collect(); + anyhow::bail!( + "checkpoint inventory does not match the profile: missing {missing:?}, unexpected {extra:?}" + ); + } + for (name, shape) in &expected { + let view = tensors.tensor(name)?; + ensure!( + view.shape() == shape, + "{name}: expected {shape:?}, got {:?}", + view.shape() + ); + ensure!( + view.dtype() == Dtype::F32, + "{name} is {:?}, the pinned checkpoint stores f32", + view.dtype() + ); + } + Ok(expected.len()) + } + + /// Read one f32 tensor with its expected shape. Upload-time storage + /// conversion (f16/bf16 narrowing per [`storage_dtype`]) happens in + /// the executor slice. + pub fn f32(&self, name: &str, shape: &[usize]) -> Result> { + let tensors = SafeTensors::deserialize(&self.data)?; + let t = tensors + .tensor(name) + .with_context(|| format!("{name} is missing"))?; + ensure!( + t.shape() == shape, + "{name}: expected {shape:?}, got {:?}", + t.shape() + ); + ensure!(t.dtype() == Dtype::F32, "{name} is not f32"); + Ok(t.data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect()) + } +} diff --git a/src/models/vela/postprocess.py b/src/models/vela/postprocess.py index 907acb9b..2f793dbd 100644 --- a/src/models/vela/postprocess.py +++ b/src/models/vela/postprocess.py @@ -24,6 +24,17 @@ # Emitted by the hybrid members only; a 0.3B response carrying it means # the wrong member's runtime answered. HYBRID_ONLY_FIELDS = {"span_heads"} +# The exact answer field sets of the pinned runtime's +# ``system_one_response``: equality, not subset — a missing field is as +# broken as a foreign one. Set and span questions answer through their +# top-level ``sets``/``spans``/``thresholds`` payloads plus Noul views. +ANSWER_FIELDS = { + "noul": {"type", "noul"}, + "choice": {"type", "choice", "confidence", "probabilities"}, + "score": {"type", "score", "confidence", "legend", "probabilities"}, +} +SET_FIELDS = {"selected", "probabilities"} +SPAN_FIELDS = {"label", "start", "end", "text", "probability"} def choice_confidence(probabilities: Sequence[float]) -> float: @@ -102,7 +113,8 @@ def finish_response( The pinned runtime owns every value in this response; these checks only refuse a response that could not be reconstructed for the request that produced it (wrong key set or order, a member-invented - field, broken usage accounting, an answer of unexpected shape). + field, broken usage accounting, an answer view that is not an + object or carries fields its kind does not allow). ``hybrid`` selects the member: the hybrid executor's responses carry ``span_heads`` alongside ``spans`` (upstream emits both together, keyed by the span question ids); the 0.3B member's never do. @@ -135,6 +147,20 @@ def finish_response( raise RuntimeError(f"{field} do not match the request's {field} questions") elif field in response: raise RuntimeError(f"the pinned runtime returned {field} without {field} questions") + for qid, payload in (response.get("sets") or {}).items(): + if not isinstance(payload, dict) or set(payload) != SET_FIELDS: + raise RuntimeError(f"set {qid!r} does not carry exactly {sorted(SET_FIELDS)}") + for qid, spans in (response.get("spans") or {}).items(): + if not isinstance(spans, list): + raise RuntimeError(f"spans {qid!r} is not a list") + for span in spans: + if not isinstance(span, dict) or set(span) != SPAN_FIELDS: + raise RuntimeError( + f"a span of {qid!r} does not carry exactly {sorted(SPAN_FIELDS)}" + ) + for qid, threshold in (response.get("thresholds") or {}).items(): + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): + raise RuntimeError(f"threshold {qid!r} is not a number") if hybrid: span_ids = _question_ids_of(questions, "span") if span_ids: @@ -155,7 +181,13 @@ def finish_response( for name, answer in answers.items(): if not isinstance(answer, dict): raise RuntimeError(f"answer {name!r} is not an object") - shape = {"type", "noul"} if answer.get("type") == "noul" else None - if shape is not None and not set(answer) <= shape: - raise RuntimeError(f"noul view {name!r} carries unexpected fields") + kind = answer.get("type") + shape = ANSWER_FIELDS.get(kind) if isinstance(kind, str) else None + if shape is None: + raise RuntimeError(f"answer {name!r} has unknown type {kind!r}") + if set(answer) != shape: + raise RuntimeError( + f"answer {name!r} does not carry exactly {sorted(shape)} " + f"(got {sorted(answer)})" + ) return response diff --git a/src/models/vela/protocol.py b/src/models/vela/protocol.py index d9341206..df026b88 100644 --- a/src/models/vela/protocol.py +++ b/src/models/vela/protocol.py @@ -83,7 +83,7 @@ def parse_body(raw: bytes) -> dict[str, Any]: """Decode a bounded JSON body and reject ambiguous JSON values.""" if len(raw) > MAX_BODY: - raise RequestError("request body exceeds 8 MiB", 413) + raise RequestError("request exceeds the 8 MiB body limit", 413) try: value = json.loads( raw.decode("utf-8"), diff --git a/tests/laya/artifacts.rs b/tests/laya/artifacts.rs index 16e261d9..78a2cbbe 100644 --- a/tests/laya/artifacts.rs +++ b/tests/laya/artifacts.rs @@ -49,7 +49,7 @@ impl Fixture { ); } let table_manifest = json!({ - "abi": 1, + "abi": 2, "laya": "0.3.20", "hidden_size": 1024, "head_dim": 64, @@ -64,7 +64,7 @@ impl Fixture { .unwrap(); fs::write(bundle.join("liblaya_cuda.so"), b"not loaded in CPU test").unwrap(); let build_manifest = json!({ - "abi": 1, + "abi": 2, "arch": "sm_90a", "library_sha256": format!("{:x}", Sha256::digest(b"not loaded in CPU test")), }); diff --git a/tests/laya/batch.rs b/tests/laya/batch.rs index 92faf922..821649e8 100644 --- a/tests/laya/batch.rs +++ b/tests/laya/batch.rs @@ -23,7 +23,7 @@ fn prepared(n: usize) -> Prepared { #[test] fn padding_preserves_order_lengths_and_markers() { let prepared = prepared(3); - let batch = pack(&prepared).unwrap(); + let batch = pack(&prepared, 50368).unwrap(); assert_eq!((batch.b, batch.l), (4, 16)); assert_eq!(batch.lens, [3, 3, 3, 0]); assert_eq!(batch.qtypes, [2, 2, 2, 0]); @@ -37,17 +37,17 @@ fn padding_preserves_order_lengths_and_markers() { #[test] fn rejects_executor_capacity_and_invalid_markers() { - assert!(pack(&prepared(17)).is_err()); + assert!(pack(&prepared(17), 50368).is_err()); let mut input = prepared(1); input.questions[0].markers.push(3); - assert!(pack(&input).is_err()); + assert!(pack(&input, 50368).is_err()); input.questions[0].markers = vec![1; 2049]; - assert!(pack(&input).is_err()); + assert!(pack(&input, 50368).is_err()); } #[test] fn empty_request_needs_no_device_rows() { - let batch = pack(&prepared(0)).unwrap(); + let batch = pack(&prepared(0), 50368).unwrap(); assert_eq!((batch.b, batch.l), (0, 0)); assert!(batch.markers.is_empty()); } @@ -56,7 +56,7 @@ fn empty_request_needs_no_device_rows() { fn padding_matches_native_sequence_alignment() { let mut input = prepared(3); input.questions[1].ids = vec![7; 257]; - let batch = pack(&input).unwrap(); + let batch = pack(&input, 50368).unwrap(); assert_eq!((batch.b, batch.l), (4, 320)); assert_eq!(batch.lens, [3, 257, 3, 0]); assert!(batch.input_ids[3..257].iter().all(|id| *id == 50283)); diff --git a/tests/laya/data/encoder-plan-original-rope.txt b/tests/laya/data/encoder-plan-original-rope.txt new file mode 100644 index 00000000..eec8a07e --- /dev/null +++ b/tests/laya/data/encoder-plan-original-rope.txt @@ -0,0 +1,251 @@ +launch embed ids encoder.embeddings.tok_embeddings.weight encoder.embeddings.norm.weight x y +dump embedding x 0 +launch qkv y encoder.layers.0.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.0.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.0.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.0.mlp.Wi.weight g +launch down g encoder.layers.0.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.1.attn_norm.weight zeros.1024 y +dump encoder0_residual x 0 +dump encoder0_normalized y 1 +launch qkv y encoder.layers.1.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.1.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.1.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.1.mlp.Wi.weight g +launch down g encoder.layers.1.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.2.attn_norm.weight zeros.1024 y +dump encoder1_residual x 0 +dump encoder1_normalized y 1 +launch qkv y encoder.layers.2.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.2.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.2.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.2.mlp.Wi.weight g +launch down g encoder.layers.2.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.3.attn_norm.weight zeros.1024 y +dump encoder2_residual x 0 +dump encoder2_normalized y 1 +launch qkv y encoder.layers.3.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.3.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.3.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.3.mlp.Wi.weight g +launch down g encoder.layers.3.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.4.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.4.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.4.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.4.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.4.mlp.Wi.weight g +launch down g encoder.layers.4.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.5.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.5.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.5.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.5.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.5.mlp.Wi.weight g +launch down g encoder.layers.5.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.6.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.6.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.6.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.6.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.6.mlp.Wi.weight g +launch down g encoder.layers.6.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.7.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.7.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.7.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.7.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.7.mlp.Wi.weight g +launch down g encoder.layers.7.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.8.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.8.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.8.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.8.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.8.mlp.Wi.weight g +launch down g encoder.layers.8.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.9.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.9.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.9.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.9.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.9.mlp.Wi.weight g +launch down g encoder.layers.9.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.10.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.10.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.10.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.10.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.10.mlp.Wi.weight g +launch down g encoder.layers.10.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.11.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.11.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.11.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.11.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.11.mlp.Wi.weight g +launch down g encoder.layers.11.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.12.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.12.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.12.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.12.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.12.mlp.Wi.weight g +launch down g encoder.layers.12.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.13.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.13.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.13.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.13.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.13.mlp.Wi.weight g +launch down g encoder.layers.13.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.14.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.14.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.14.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.14.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.14.mlp.Wi.weight g +launch down g encoder.layers.14.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.15.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.15.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.15.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.15.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.15.mlp.Wi.weight g +launch down g encoder.layers.15.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.16.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.16.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.16.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.16.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.16.mlp.Wi.weight g +launch down g encoder.layers.16.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.17.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.17.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.17.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.17.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.17.mlp.Wi.weight g +launch down g encoder.layers.17.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.18.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.18.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.18.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.18.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.18.mlp.Wi.weight g +launch down g encoder.layers.18.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.19.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.19.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.19.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.19.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.19.mlp.Wi.weight g +launch down g encoder.layers.19.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.20.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.20.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.20.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.20.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.20.mlp.Wi.weight g +launch down g encoder.layers.20.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.21.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.21.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.21.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.21.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.21.mlp.Wi.weight g +launch down g encoder.layers.21.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.22.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.22.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.22.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.22.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.22.mlp.Wi.weight g +launch down g encoder.layers.22.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.23.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.23.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.23.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.23.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.23.mlp.Wi.weight g +launch down g encoder.layers.23.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.24.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.24.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.24.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.24.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.24.mlp.Wi.weight g +launch down g encoder.layers.24.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.25.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.25.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.25.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.25.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.25.mlp.Wi.weight g +launch down g encoder.layers.25.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.26.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.26.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.26.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.26.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.26.mlp.Wi.weight g +launch down g encoder.layers.26.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.27.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.27.attn.Wqkv.weight zeros.3072 qkv +launch rope_original qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.27.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.27.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.27.mlp.Wi.weight g +launch down g encoder.layers.27.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.final_norm.weight zeros.1024 y +dump encoder27_residual x 0 +dump encoder27_normalized y 1 +launch type y type_emb.weight types x +launch ln_bias x y head.layers.0.norm1.weight head.layers.0.norm1.bias y +launch head_in y head.layers.0.self_attn.in_proj_weight head.layers.0.self_attn.in_proj_bias qkv +launch attn_full qkv lens o +launch head_out o head.layers.0.self_attn.out_proj.weight head.layers.0.self_attn.out_proj.bias y +launch addln_bias x y head.layers.0.norm2.weight head.layers.0.norm2.bias y +launch ffn1 y head.layers.0.linear1.weight head.layers.0.linear1.bias ff +launch ffn2 ff head.layers.0.linear2.weight head.layers.0.linear2.bias y +launch residual x y +launch ln_bias x y head.layers.1.norm1.weight head.layers.1.norm1.bias y +launch head_in y head.layers.1.self_attn.in_proj_weight head.layers.1.self_attn.in_proj_bias qkv +launch attn_full qkv lens o +launch head_out o head.layers.1.self_attn.out_proj.weight head.layers.1.self_attn.out_proj.bias y +launch addln_bias x y head.layers.1.norm2.weight head.layers.1.norm2.bias y +launch ffn1 y head.layers.1.linear1.weight head.layers.1.linear1.bias ff +launch ffn2 ff head.layers.1.linear2.weight head.layers.1.linear2.bias y +launch residual x y diff --git a/tests/laya/data/encoder-plan.txt b/tests/laya/data/encoder-plan.txt new file mode 100644 index 00000000..4f1df250 --- /dev/null +++ b/tests/laya/data/encoder-plan.txt @@ -0,0 +1,251 @@ +launch embed ids encoder.embeddings.tok_embeddings.weight encoder.embeddings.norm.weight x y +dump embedding x 0 +launch qkv y encoder.layers.0.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.0.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.0.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.0.mlp.Wi.weight g +launch down g encoder.layers.0.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.1.attn_norm.weight zeros.1024 y +dump encoder0_residual x 0 +dump encoder0_normalized y 1 +launch qkv y encoder.layers.1.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.1.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.1.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.1.mlp.Wi.weight g +launch down g encoder.layers.1.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.2.attn_norm.weight zeros.1024 y +dump encoder1_residual x 0 +dump encoder1_normalized y 1 +launch qkv y encoder.layers.2.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.2.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.2.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.2.mlp.Wi.weight g +launch down g encoder.layers.2.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.3.attn_norm.weight zeros.1024 y +dump encoder2_residual x 0 +dump encoder2_normalized y 1 +launch qkv y encoder.layers.3.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.3.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.3.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.3.mlp.Wi.weight g +launch down g encoder.layers.3.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.4.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.4.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.4.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.4.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.4.mlp.Wi.weight g +launch down g encoder.layers.4.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.5.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.5.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.5.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.5.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.5.mlp.Wi.weight g +launch down g encoder.layers.5.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.6.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.6.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.6.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.6.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.6.mlp.Wi.weight g +launch down g encoder.layers.6.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.7.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.7.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.7.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.7.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.7.mlp.Wi.weight g +launch down g encoder.layers.7.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.8.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.8.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.8.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.8.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.8.mlp.Wi.weight g +launch down g encoder.layers.8.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.9.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.9.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.9.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.9.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.9.mlp.Wi.weight g +launch down g encoder.layers.9.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.10.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.10.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.10.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.10.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.10.mlp.Wi.weight g +launch down g encoder.layers.10.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.11.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.11.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.11.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.11.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.11.mlp.Wi.weight g +launch down g encoder.layers.11.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.12.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.12.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.12.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.12.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.12.mlp.Wi.weight g +launch down g encoder.layers.12.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.13.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.13.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.13.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.13.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.13.mlp.Wi.weight g +launch down g encoder.layers.13.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.14.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.14.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.14.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.14.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.14.mlp.Wi.weight g +launch down g encoder.layers.14.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.15.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.15.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.15.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.15.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.15.mlp.Wi.weight g +launch down g encoder.layers.15.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.16.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.16.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.16.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.16.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.16.mlp.Wi.weight g +launch down g encoder.layers.16.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.17.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.17.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.17.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.17.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.17.mlp.Wi.weight g +launch down g encoder.layers.17.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.18.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.18.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.18.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.18.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.18.mlp.Wi.weight g +launch down g encoder.layers.18.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.19.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.19.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.19.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.19.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.19.mlp.Wi.weight g +launch down g encoder.layers.19.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.20.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.20.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.20.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.20.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.20.mlp.Wi.weight g +launch down g encoder.layers.20.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.21.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.21.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.21.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.21.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.21.mlp.Wi.weight g +launch down g encoder.layers.21.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.22.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.22.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.22.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.22.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.22.mlp.Wi.weight g +launch down g encoder.layers.22.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.23.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.23.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.23.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.23.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.23.mlp.Wi.weight g +launch down g encoder.layers.23.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.24.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.24.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.24.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.24.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.24.mlp.Wi.weight g +launch down g encoder.layers.24.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.25.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.25.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.25.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.25.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.25.mlp.Wi.weight g +launch down g encoder.layers.25.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.26.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.26.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_local_cos rope_local_sin +launch attn_local qkv lens o +launch out o encoder.layers.26.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.26.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.26.mlp.Wi.weight g +launch down g encoder.layers.26.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.layers.27.attn_norm.weight zeros.1024 y +launch qkv y encoder.layers.27.attn.Wqkv.weight zeros.3072 qkv +launch rope qkv rope_full_cos rope_full_sin +launch attn_full qkv lens o +launch out o encoder.layers.27.attn.Wo.weight zeros.1024 y +launch addln x y encoder.layers.27.mlp_norm.weight zeros.1024 y +launch geglu y encoder.layers.27.mlp.Wi.weight g +launch down g encoder.layers.27.mlp.Wo.weight zeros.1024 y +launch addln x y encoder.final_norm.weight zeros.1024 y +dump encoder27_residual x 0 +dump encoder27_normalized y 1 +launch type y type_emb.weight types x +launch ln_bias x y head.layers.0.norm1.weight head.layers.0.norm1.bias y +launch head_in y head.layers.0.self_attn.in_proj_weight head.layers.0.self_attn.in_proj_bias qkv +launch attn_full qkv lens o +launch head_out o head.layers.0.self_attn.out_proj.weight head.layers.0.self_attn.out_proj.bias y +launch addln_bias x y head.layers.0.norm2.weight head.layers.0.norm2.bias y +launch ffn1 y head.layers.0.linear1.weight head.layers.0.linear1.bias ff +launch ffn2 ff head.layers.0.linear2.weight head.layers.0.linear2.bias y +launch residual x y +launch ln_bias x y head.layers.1.norm1.weight head.layers.1.norm1.bias y +launch head_in y head.layers.1.self_attn.in_proj_weight head.layers.1.self_attn.in_proj_bias qkv +launch attn_full qkv lens o +launch head_out o head.layers.1.self_attn.out_proj.weight head.layers.1.self_attn.out_proj.bias y +launch addln_bias x y head.layers.1.norm2.weight head.layers.1.norm2.bias y +launch ffn1 y head.layers.1.linear1.weight head.layers.1.linear1.bias ff +launch ffn2 ff head.layers.1.linear2.weight head.layers.1.linear2.bias y +launch residual x y diff --git a/tests/laya/model.rs b/tests/laya/model.rs index 922bbaee..95328fda 100644 --- a/tests/laya/model.rs +++ b/tests/laya/model.rs @@ -51,3 +51,42 @@ mod grouped; #[cfg(unix)] #[path = "original_encode.rs"] mod original_encode; + +#[test] +fn workspace_sizes_match_the_recorded_profile_layout() { + let structure = gate_structure(); + for (b, l) in [(1usize, 16usize), (1, 512), (16, 512), (4, 128)] { + let m = b * l; + let sizes = Workspace::sizes(b, l, &structure, &limits()).unwrap(); + assert_eq!( + sizes, + [ + m * 8, + b * 4, + b * 8, + m * 1024 * 4, + m * 1024 * 2, + m * 1024 * 6, + m * 1024 * 2, + m * 2624 * 2, + m * 4096 * 2, + 2048 * 4, + 17 * 4, + 2048 * 1024 * 2, + 2048 * 1024 * 2, + 2048 * 2, + 16 * 1028 * 2, + 16 * 256 * 2, + 16 * 2 * 2, + ], + "b={b} l={l}" + ); + } +} + +fn gate_structure() -> omni_modernbert::structure::ModernBertStructure { + crate::profile::laya_0_3_20().structure +} +fn limits() -> omni_modernbert::gate::BackendLimits { + crate::profile::laya_0_3_20().limits +} diff --git a/tests/laya/model_cache.rs b/tests/laya/model_cache.rs index 776656c5..ecf8268d 100644 --- a/tests/laya/model_cache.rs +++ b/tests/laya/model_cache.rs @@ -1,5 +1,12 @@ use super::*; +fn structure() -> omni_modernbert::structure::ModernBertStructure { + crate::profile::laya_0_3_20().structure +} +fn limits() -> omni_modernbert::gate::BackendLimits { + crate::profile::laya_0_3_20().limits +} + fn batch_for(b: usize, l: usize, value: i64) -> Batch { Batch { b, @@ -16,8 +23,8 @@ fn cache_limits_evict_before_allocation_and_allow_exact_fits() { let fixture = fixture::Fixture::new(&[]); let mut model = fixture_model_at(false, fixture.path()); model.graphs = true; - let a = Workspace::required_bytes(1, 16).unwrap(); - let b = Workspace::required_bytes(1, 32).unwrap(); + let a = Workspace::required_bytes(1, 16, &structure(), &limits()).unwrap(); + let b = Workspace::required_bytes(1, 32, &structure(), &limits()).unwrap(); model.cache_config = CacheConfig { max_shapes: 1, max_bytes: a.max(b), @@ -73,8 +80,8 @@ fn byte_budget_lru_and_disabled_or_oversized_fallback_reuse() { let fixture = fixture::Fixture::new(&[]); let mut model = fixture_model_at(false, fixture.path()); model.graphs = true; - let a = Workspace::required_bytes(1, 16).unwrap(); - let b = Workspace::required_bytes(1, 32).unwrap(); + let a = Workspace::required_bytes(1, 16, &structure(), &limits()).unwrap(); + let b = Workspace::required_bytes(1, 32, &structure(), &limits()).unwrap(); model.cache_config = CacheConfig { max_shapes: 8, max_bytes: a + b, @@ -141,10 +148,10 @@ fn workspace_estimate_matches_allocation_for_all_supported_shapes() { let weights = unsafe { live() }; for b in [1, 2, 4, 8, 16] { for l in (16..=512).step_by(16) { - let workspace = Workspace::new(&model.cuda, b, l).unwrap(); + let workspace = Workspace::new(&model.cuda, b, l, &structure(), &limits()).unwrap(); let allocated = unsafe { live() } - weights; assert_eq!( - Workspace::required_bytes(b, l).unwrap(), + Workspace::required_bytes(b, l, &structure(), &limits()).unwrap(), allocated, "B={b}, L={l}" ); @@ -154,7 +161,7 @@ fn workspace_estimate_matches_allocation_for_all_supported_shapes() { } } for (b, l) in [(0, 16), (3, 16), (32, 16), (1, 0), (1, 17), (1, 528)] { - assert!(Workspace::required_bytes(b, l).is_err()); + assert!(Workspace::required_bytes(b, l, &structure(), &limits()).is_err()); } } @@ -165,7 +172,7 @@ fn clear_releases_graph_and_eager_workspaces_and_allows_reuse() { model.graphs = true; model.cache_config = CacheConfig { max_shapes: 1, - max_bytes: Workspace::required_bytes(1, 16).unwrap(), + max_bytes: Workspace::required_bytes(1, 16, &structure(), &limits()).unwrap(), }; let live = unsafe { model @@ -222,6 +229,6 @@ fn failed_capture_admission_releases_partial_workspace_and_can_retry() { assert_eq!(model.cache.len(), 1); assert_eq!( model.cached_bytes, - Workspace::required_bytes(1, 16).unwrap() + Workspace::required_bytes(1, 16, &structure(), &limits()).unwrap() ); } diff --git a/tests/laya/model_dispatch.rs b/tests/laya/model_dispatch.rs index d2feb506..eb721958 100644 --- a/tests/laya/model_dispatch.rs +++ b/tests/laya/model_dispatch.rs @@ -1,5 +1,12 @@ use super::*; +fn structure() -> omni_modernbert::structure::ModernBertStructure { + crate::profile::laya_0_3_20().structure +} +fn limits() -> omni_modernbert::gate::BackendLimits { + crate::profile::laya_0_3_20().limits +} + #[test] fn prepared_plan_preserves_original_launch_trace_for_all_shape_selectors() { let fixture = fixture::Fixture::new(&[]); @@ -16,7 +23,7 @@ fn prepared_plan_preserves_original_launch_trace_for_all_shape_selectors() { let model = fixture_model_at(original_rope, fixture.path()); for b in [1, 2, 4, 8, 16] { for l in [16, 512] { - let s = Workspace::new(&model.cuda, b, l).unwrap(); + let s = Workspace::new(&model.cuda, b, l, &structure(), &limits()).unwrap(); unsafe { reset() }; model.encode_original(&s, original_rope).unwrap(); let expected = unsafe { std::ffi::CStr::from_ptr(trace()) } @@ -41,12 +48,12 @@ fn prepared_plan_preserves_original_launch_trace_for_all_shape_selectors() { fn missing_specialized_attention_only_rejects_its_shape() { let fixture = fixture::Fixture::new(&["OMIT_SPECIALIZED_ATTN"]); let model = fixture_model_at(false, fixture.path()); - let short = Workspace::new(&model.cuda, 1, 16).unwrap(); + let short = Workspace::new(&model.cuda, 1, 16, &structure(), &limits()).unwrap(); model.encode(&short).unwrap(); - let generic_long = Workspace::new(&model.cuda, 2, 512).unwrap(); + let generic_long = Workspace::new(&model.cuda, 2, 512, &structure(), &limits()).unwrap(); model.encode(&generic_long).unwrap(); for b in [1, 4] { - let specialized = Workspace::new(&model.cuda, b, 512).unwrap(); + let specialized = Workspace::new(&model.cuda, b, 512, &structure(), &limits()).unwrap(); assert!( model .encode(&specialized) diff --git a/tests/laya/model_fixture.rs b/tests/laya/model_fixture.rs index 09c9e62c..477ec27b 100644 --- a/tests/laya/model_fixture.rs +++ b/tests/laya/model_fixture.rs @@ -57,9 +57,10 @@ pub(super) fn fixture_model_at(original_rope: bool, path: &Path) -> Model { } } let (encoder, agent) = configs(); - let plan = prepare_encoder(&cuda, &weights, original_rope).unwrap(); + let config = load_config(&encoder, &agent).unwrap(); + let plan = prepare_encoder(&cuda, &weights, &config.structure, original_rope).unwrap(); Model { - config: load_config(&encoder, &agent).unwrap(), + config, blas: Blas::new(&cuda).unwrap(), cuda, weights, diff --git a/tests/laya/plan.rs b/tests/laya/plan.rs new file mode 100644 index 00000000..b59ff74f --- /dev/null +++ b/tests/laya/plan.rs @@ -0,0 +1,47 @@ +use omni_laya::plan::{plan_encoder, render_plan}; + +// The golden fixtures are the rendered encoder plans of the registered +// Laya 0.3.20 profile: `encoder-plan.txt` drives the rope-table kernel, +// `encoder-plan-original-rope.txt` its `rope_original` reference +// variant. The dispatch trace test in model_dispatch.rs cross-checks +// the same plans against the frozen CUDA-ABI launch sequence. + +fn laya_structure() -> omni_modernbert::structure::ModernBertStructure { + omni_laya::profile::laya_0_3_20().structure +} + +#[test] +fn encoder_plan_matches_the_recorded_sequence() { + assert_eq!( + render_plan(&plan_encoder(&laya_structure(), false)), + include_str!("data/encoder-plan.txt") + ); + assert_eq!( + render_plan(&plan_encoder(&laya_structure(), true)), + include_str!("data/encoder-plan-original-rope.txt") + ); +} + +#[test] +fn encoder_plan_shape_is_the_supported_profile() { + let plan = plan_encoder(&laya_structure(), false); + let launches = plan + .iter() + .filter(|s| matches!(s, omni_laya::plan::PlannedStep::Launch { .. })) + .count(); + // embed + 28 layers x 8 kernels + type add + 2 head layers x 8 kernels. + assert_eq!(launches, 1 + 28 * 8 + 1 + 2 * 8); + let rendered = render_plan(&plan); + assert_eq!(rendered.matches("attn_full").count(), 12); + assert_eq!(rendered.matches("attn_local").count(), 18); + assert!(rendered.contains("addln x y encoder.final_norm.weight zeros.1024 y")); + assert!(rendered.ends_with("residual x y\n")); + // The two rope variants differ in exactly the 28 rope kernel names. + let original = render_plan(&plan_encoder(&laya_structure(), true)); + let differences = rendered + .lines() + .zip(original.lines()) + .filter(|(a, b)| a != b) + .count(); + assert_eq!(differences, 28); +} diff --git a/tests/laya/profile.rs b/tests/laya/profile.rs new file mode 100644 index 00000000..739f60f3 --- /dev/null +++ b/tests/laya/profile.rs @@ -0,0 +1,180 @@ +use omni_laya::profile::{LayaBundleAdapter, laya_0_3_20}; +use omni_modernbert::structure::{AttentionKind, FieldOrigin}; +use serde_json::{Value, json}; +use std::fs; +use std::path::PathBuf; +use tempfile::{TempDir, tempdir}; + +/// The Laya 0.3.20 bundle's `encoder/config.json` as shipped: the keys +/// the native worker consumes. The HF-side members the dialect ignores +/// (`max_position_embeddings`, `head_dim`) are added by name where a +/// test exercises that behavior. +fn laya_encoder_config() -> Value { + json!({ + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 1024, "intermediate_size": 2624, + "num_attention_heads": 16, "num_hidden_layers": 28, + "vocab_size": 50368, "norm_eps": 0.00001, "local_attention": 128, + "layer_types": (["full_attention", "sliding_attention", "sliding_attention"] + .repeat(10))[..28], + "rope_parameters": { + "full_attention": {"rope_type": "default", "rope_theta": 160000.0}, + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0} + } + }) +} + +fn laya_bundle(encoder: &Value) -> (TempDir, PathBuf) { + let dir = tempdir().unwrap(); + fs::create_dir(dir.path().join("encoder")).unwrap(); + fs::write(dir.path().join("encoder/config.json"), encoder.to_string()).unwrap(); + let path = dir.path().to_path_buf(); + (dir, path) +} + +fn parse_laya(encoder: &Value) -> omni_modernbert::structure::ModernBertStructure { + let (_dir, path) = laya_bundle(encoder); + LayaBundleAdapter::read(&path).unwrap().0 +} + +#[test] +fn bundle_adapter_parses_the_profile_with_provenance() { + let (_dir, path) = laya_bundle(&laya_encoder_config()); + let (structure, meta) = LayaBundleAdapter::read(&path).unwrap(); + assert_eq!(structure.hidden_size, 1024); + assert_eq!(structure.intermediate_size, 2624); + assert_eq!(structure.num_hidden_layers, 28); + assert_eq!(structure.num_attention_heads, 16); + assert_eq!(structure.head_dim, 64); + assert_eq!(structure.local_window, 128); + assert_eq!(structure.vocab_size, 50368); + assert_eq!(structure.type_embedding_rows, Some(3)); + assert_eq!(structure.rope_full.theta, 160000.0); + assert_eq!(structure.rope_sliding.theta, 10000.0); + assert_eq!(structure.max_positions, 512); + assert_eq!(meta.model_type, "modernbert"); + assert_eq!(meta.config_files, ["encoder/config.json"]); + assert_eq!(meta.unrecognized_keys, Vec::::new()); + assert_eq!( + structure.provenance.get("schedule"), + Some(&FieldOrigin::Source( + "encoder/config.json#layer_types".into() + )) + ); + assert_eq!( + structure.provenance.get("head_dim"), + Some(&FieldOrigin::Derived( + "hidden_size / num_attention_heads".into() + )) + ); + assert!(matches!( + structure.provenance.get("max_positions"), + Some(FieldOrigin::Default(_)) + )); +} + +/// The dialect ignores the HF position and head-dim members the native +/// worker never consumed: a bundle declaring them (as stock HF +/// ModernBERT configs do) still admits with the bundle's own 512 +/// exercised positions and derived head dim, and the keys are recorded +/// as unrecognized instead of gated. +#[test] +fn bundle_adapter_ignores_hf_position_and_head_dim_members() { + let mut config = laya_encoder_config(); + config["max_position_embeddings"] = json!(8192); + config["head_dim"] = json!(32); + let structure = parse_laya(&config); + assert_eq!(structure.max_positions, 512); + assert_eq!(structure.head_dim, 64); + let (_dir, path) = laya_bundle(&config); + let (_, meta) = LayaBundleAdapter::read(&path).unwrap(); + let mut keys = meta.unrecognized_keys.clone(); + keys.sort(); + assert_eq!(keys, ["head_dim", "max_position_embeddings"]); + laya_0_3_20().admit(&structure).unwrap(); +} + +#[test] +fn bundle_adapter_rejects_marker_and_family_violations() { + for (pointer, value) in [ + ("/model_type", json!("bert")), + ("/hidden_activation", json!("relu")), + ("/attention_bias", json!(true)), + ("/mlp_bias", json!(true)), + ("/norm_bias", json!(true)), + ("/layer_types/0", json!("sparse_attention")), + ("/rope_parameters/full_attention/rope_type", json!("linear")), + ("/rope_parameters/sliding_attention/rope_theta", json!(null)), + ] { + let mut bad = laya_encoder_config(); + *bad.pointer_mut(pointer).unwrap() = value; + let (_dir, path) = laya_bundle(&bad); + assert!( + LayaBundleAdapter::read(&path).is_err(), + "accepted {pointer}" + ); + } + let mut missing = laya_encoder_config(); + missing.as_object_mut().unwrap().remove("layer_types"); + let (_dir, path) = laya_bundle(&missing); + assert!(LayaBundleAdapter::read(&path).is_err()); + let mut short = laya_encoder_config(); + short["layer_types"] = json!(["full_attention"]); + let (_dir, path) = laya_bundle(&short); + assert!(LayaBundleAdapter::read(&path).is_err()); +} + +#[test] +fn the_registered_profile_admits_its_structure_and_rejects_field_drift() { + let profile = laya_0_3_20(); + let structure = parse_laya(&laya_encoder_config()); + let identity = profile.admit(&structure).unwrap(); + assert_eq!(identity.schema_version, 1); + assert_eq!(identity.family, "modernbert"); + assert_eq!(identity.backend, "laya_cuda"); + assert_eq!(identity.abi_version, 1); + + for mutate in [ + ("hidden_size", { + let mut s = parse_laya(&laya_encoder_config()); + s.hidden_size = 768; + s + }), + ("rope", { + let mut s = parse_laya(&laya_encoder_config()); + s.rope_sliding.theta = 160000.0; + s + }), + ("schedule", { + let mut s = parse_laya(&laya_encoder_config()); + s.schedule[1] = AttentionKind::Full; + s + }), + ("type_embedding_rows", { + let mut s = parse_laya(&laya_encoder_config()); + s.type_embedding_rows = None; + s + }), + ("max_positions", { + let mut s = parse_laya(&laya_encoder_config()); + s.max_positions = 8192; + s + }), + ] { + assert!(profile.admit(&mutate.1).is_err(), "accepted {}", mutate.0); + } +} + +#[test] +fn the_registered_limits_check_the_request_envelope() { + let limits = &laya_0_3_20().limits; + limits.check_request(16, 512).unwrap(); + limits.check_request(1, 16).unwrap(); + for (batch, len) in [(32, 512), (3, 512), (1, 8), (1, 513), (1, 520)] { + assert!( + limits.check_request(batch, len).is_err(), + "accepted batch {batch} len {len}" + ); + } +} diff --git a/tests/modernbert/plan.rs b/tests/modernbert/plan.rs new file mode 100644 index 00000000..03b9d2b9 --- /dev/null +++ b/tests/modernbert/plan.rs @@ -0,0 +1,175 @@ +use omni_modernbert::adapter::ModernBertConfigAdapter; +use omni_modernbert::plan::{self, EncoderPlanOptions, PlanAnchor, PlanArg, PlanStep, ScratchSlot}; +use omni_modernbert::structure::ModernBertStructure; +use serde_json::json; + +/// A 1024-hidden member with a type embedding and plain rotary tables, +/// parsed as a stock HF config (the same geometry the Laya bundle +/// dialect produces), so the family plan tests stay model-agnostic. +fn laya_structure() -> ModernBertStructure { + let config = json!({ + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 1024, "intermediate_size": 2624, + "num_attention_heads": 16, "num_hidden_layers": 28, + "vocab_size": 50368, "norm_eps": 0.00001, "local_attention": 128, + "type_vocab_size": 3, "max_position_embeddings": 512, + "layer_types": (["full_attention", "sliding_attention", "sliding_attention"] + .repeat(10))[..28], + "global_rope_theta": 160000.0, "local_rope_theta": 10000.0 + }); + ModernBertConfigAdapter::from_value(&config, "config.json") + .unwrap() + .0 +} + +fn vela_structure() -> ModernBertStructure { + let config = json!({ + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 768, "intermediate_size": 1152, + "num_attention_heads": 12, "num_hidden_layers": 22, + "vocab_size": 256008, "norm_eps": 1e-5, + "local_attention": 128, "global_attn_every_n_layers": 3, + "global_rope_theta": 160000, "local_rope_theta": 160000, + "max_position_embeddings": 32768, + "rope_scaling": { + "rope_type": "yarn", "factor": 4.0, + "original_max_position_embeddings": 8192, "truncate": true + } + }); + ModernBertConfigAdapter::from_value(&config, "config.json#encoder_config") + .unwrap() + .0 +} + +fn launches(steps: &[PlanStep]) -> Vec<&PlanStep> { + steps + .iter() + .filter(|s| matches!(s, PlanStep::Launch { .. })) + .collect() +} + +#[test] +fn laya_plan_reproduces_the_encoder_step_sequence() { + let steps = plan::encoder_plan(&laya_structure(), &EncoderPlanOptions::default()); + // embed + 28 layers x 8 kernels + the type add. + assert_eq!(launches(&steps).len(), 1 + 28 * 8 + 1); + assert_eq!( + steps[0].render(), + "embed ids encoder.embeddings.tok_embeddings.weight encoder.embeddings.norm.weight x y" + ); + assert_eq!(steps[1], PlanStep::Anchor(PlanAnchor::AfterEmbed)); + + // Layer 0 runs full attention on the full-kind tables. + assert_eq!( + steps[2].render(), + "qkv y encoder.layers.0.attn.Wqkv.weight zeros.3072 qkv" + ); + assert_eq!(steps[3].render(), "rope qkv rope_full_cos rope_full_sin"); + assert_eq!(steps[4].render(), "attn_full qkv lens o"); + assert_eq!( + steps[5].render(), + "out o encoder.layers.0.attn.Wo.weight zeros.1024 y" + ); + assert_eq!( + steps[6].render(), + "addln x y encoder.layers.0.mlp_norm.weight zeros.1024 y" + ); + assert_eq!( + steps[7].render(), + "geglu y encoder.layers.0.mlp.Wi.weight g" + ); + assert_eq!( + steps[8].render(), + "down g encoder.layers.0.mlp.Wo.weight zeros.1024 y" + ); + // Layer 0's post-MLP norm is layer 1's attention norm. + assert_eq!( + steps[9].render(), + "addln x y encoder.layers.1.attn_norm.weight zeros.1024 y" + ); + + // Layer 1 is sliding (layer 0 ends with its anchor step at index 10). + assert!(steps[12].render().starts_with("rope qkv rope_local")); + assert_eq!(steps[13].render(), "attn_local qkv lens o"); + + // The last layer's post-MLP norm is the final norm. + let last_layer_addln = steps + .iter() + .position(|s| s.render().contains("encoder.final_norm.weight")) + .unwrap(); + assert_eq!( + steps[last_layer_addln].render(), + "addln x y encoder.final_norm.weight zeros.1024 y" + ); + + // The type add closes the encoder plan. + let last = launches(&steps).last().unwrap().render(); + assert_eq!(last, "type y type_emb.weight types x"); + assert_eq!( + steps.last(), + Some(&PlanStep::Launch { + kernel: "type".into(), + args: vec![ + PlanArg::Scratch(ScratchSlot::Y), + PlanArg::Weight("type_emb.weight".into()), + PlanArg::Scratch(ScratchSlot::Types), + PlanArg::Scratch(ScratchSlot::X), + ] + }) + ); +} + +#[test] +fn vela_plan_runs_its_own_schedule_without_a_type_add() { + let steps = plan::encoder_plan(&vela_structure(), &EncoderPlanOptions::default()); + assert_eq!(launches(&steps).len(), 1 + 22 * 8); + assert!(!steps.iter().any(|s| s.render().contains("type_emb"))); + // Layer 21 is the last full-attention layer; zeros follow its width. + assert!( + steps + .iter() + .any(|s| s.render() == "qkv y encoder.layers.21.attn.Wqkv.weight zeros.2304 qkv") + ); + assert!(steps.iter().any(|s| s.render() == "attn_full qkv lens o")); + assert_eq!( + steps + .iter() + .filter(|s| s.render() == "attn_local qkv lens o") + .count(), + 14 + ); + assert!(steps.contains(&PlanStep::Anchor(PlanAnchor::LayerEnd(21)))); + assert_eq!( + steps + .iter() + .filter(|s| matches!(s, PlanStep::Anchor(PlanAnchor::LayerEnd(_)))) + .count(), + 22 + ); +} + +#[test] +fn original_rope_option_selects_the_reference_variant() { + let steps = plan::encoder_plan( + &laya_structure(), + &EncoderPlanOptions { + original_rope: true, + }, + ); + assert!( + steps + .iter() + .any(|s| s.render() == "rope_original qkv rope_full_cos rope_full_sin") + ); +} + +#[test] +fn plan_rendering_is_deterministic() { + let once = plan::encoder_plan(&laya_structure(), &EncoderPlanOptions::default()); + let twice = plan::encoder_plan(&laya_structure(), &EncoderPlanOptions::default()); + let render = |steps: &[PlanStep]| steps.iter().map(PlanStep::render).collect::>(); + assert_eq!(render(&once), render(&twice)); + assert_eq!(render(&once).len(), once.len()); +} diff --git a/tests/modernbert/structure.rs b/tests/modernbert/structure.rs new file mode 100644 index 00000000..f075b661 --- /dev/null +++ b/tests/modernbert/structure.rs @@ -0,0 +1,439 @@ +use omni_modernbert::adapter::ModernBertConfigAdapter; +use omni_modernbert::gate::{self, structure_identity}; +use omni_modernbert::geometry::{self, EncoderScratch}; +use omni_modernbert::structure::{ + Activation, AttentionKind, FieldOrigin, ModernBertStructure, RopeSpec, +}; +use serde_json::{Value, json}; + +/// A 1024-hidden family member with an explicit schedule, a type +/// embedding and plain rotary tables — the same geometry the Laya +/// bundle dialect produces, expressed as a stock HF config so the +/// family tests stay model-agnostic. +fn family_config() -> Value { + json!({ + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 1024, "intermediate_size": 2624, + "num_attention_heads": 16, "num_hidden_layers": 28, + "vocab_size": 50368, "norm_eps": 0.00001, "local_attention": 128, + "type_vocab_size": 3, "max_position_embeddings": 512, + "layer_types": (["full_attention", "sliding_attention", "sliding_attention"] + .repeat(10))[..28], + "global_rope_theta": 160000.0, "local_rope_theta": 10000.0 + }) +} + +/// The pinned vllm-sr/Vela-2.0-0.3B@d6f03aa9 `encoder_config` member. +fn vela_encoder_config() -> Value { + json!({ + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 768, "intermediate_size": 1152, + "num_attention_heads": 12, "num_hidden_layers": 22, + "vocab_size": 256008, "norm_eps": 1e-5, "layer_norm_eps": 1e-5, + "local_attention": 128, "global_attn_every_n_layers": 3, + "global_rope_theta": 160000, "local_rope_theta": 160000, + "max_position_embeddings": 32768, "position_embedding_type": "sans_pos", + "rope_scaling": { + "rope_type": "yarn", "factor": 4.0, + "original_max_position_embeddings": 8192, + "beta_fast": 32.0, "beta_slow": 1.0, "truncate": true + }, + "cls_token_id": 1, "pad_token_id": 0 + }) +} + +fn parse(config: &Value) -> ModernBertStructure { + ModernBertConfigAdapter::from_value(config, "config.json") + .unwrap() + .0 +} + +#[test] +fn hf_adapter_parses_the_vela_profile_with_derived_schedule() { + let (structure, meta) = + ModernBertConfigAdapter::from_value(&vela_encoder_config(), "config.json#encoder_config") + .unwrap(); + assert_eq!(structure.hidden_size, 768); + assert_eq!(structure.intermediate_size, 1152); + assert_eq!(structure.num_hidden_layers, 22); + assert_eq!(structure.num_attention_heads, 12); + assert_eq!(structure.head_dim, 64); + assert_eq!(structure.vocab_size, 256008); + assert_eq!(structure.type_embedding_rows, None); + assert_eq!(structure.max_positions, 32768); + let full_layers: Vec<_> = structure + .schedule + .iter() + .enumerate() + .filter(|(_, k)| **k == AttentionKind::Full) + .map(|(i, _)| i) + .collect(); + assert_eq!(full_layers, [0, 3, 6, 9, 12, 15, 18, 21]); + assert_eq!(structure.rope_full.theta, 160000.0); + assert_eq!(structure.rope_sliding.theta, 160000.0); + let yarn = structure.rope_full.scaling.unwrap(); + assert_eq!(yarn.factor, 4.0); + assert_eq!(yarn.original_max_positions, 8192); + assert_eq!(structure.rope_full.scaling, structure.rope_sliding.scaling); + assert_eq!(meta.unrecognized_keys, ["cls_token_id", "pad_token_id"]); + match structure.provenance.get("schedule") { + Some(FieldOrigin::Derived(how)) => { + assert!(how.contains("global_attn_every_n_layers=3"), "{how}") + } + other => panic!("schedule provenance {other:?}"), + } +} + +#[test] +fn hf_adapter_accepts_explicit_layer_types_over_the_interval() { + let mut config = vela_encoder_config(); + config["global_attn_every_n_layers"] = json!(2); + config["layer_types"] = json!((vec!["sliding_attention"; 22])); + let structure = parse(&config); + assert!( + structure + .schedule + .iter() + .all(|k| *k == AttentionKind::Sliding) + ); + assert!(matches!( + structure.provenance.get("schedule"), + Some(FieldOrigin::Source(_)) + )); +} + +#[test] +fn hf_adapter_falls_back_to_the_generic_rope_theta() { + let mut config = family_config(); + let object = config.as_object_mut().unwrap(); + object.remove("global_rope_theta"); + object.remove("local_rope_theta"); + object.insert("rope_theta".into(), json!(500000.0)); + let structure = parse(&config); + assert_eq!(structure.rope_full.theta, 500000.0); + assert_eq!(structure.rope_sliding.theta, 500000.0); + for field in ["rope_full", "rope_sliding"] { + match structure.provenance.get(field) { + Some(FieldOrigin::Source(key)) => assert_eq!(key, "config.json#rope_theta"), + other => panic!("{field} provenance {other:?}"), + } + } +} + +#[test] +fn hf_adapter_prefers_the_per_kind_theta_over_the_generic_key() { + let mut config = family_config(); + let object = config.as_object_mut().unwrap(); + object.remove("global_rope_theta"); + object.insert("rope_theta".into(), json!(500000.0)); + let structure = parse(&config); + assert_eq!(structure.rope_full.theta, 500000.0); + assert_eq!(structure.rope_sliding.theta, 10000.0); + assert!(matches!( + structure.provenance.get("rope_sliding"), + Some(FieldOrigin::Source(_)) + )); +} + +#[test] +fn hf_adapter_reads_type_vocab_and_position_members() { + let (structure, meta) = + ModernBertConfigAdapter::from_value(&family_config(), "encoder/config.json").unwrap(); + assert_eq!(structure.hidden_size, 1024); + assert_eq!(structure.head_dim, 64); + assert_eq!(structure.type_embedding_rows, Some(3)); + assert_eq!(structure.max_positions, 512); + assert_eq!(structure.rope_full.theta, 160000.0); + assert_eq!(structure.rope_sliding.theta, 10000.0); + assert_eq!(meta.unrecognized_keys, Vec::::new()); + assert_eq!( + structure.provenance.get("type_embedding_rows"), + Some(&FieldOrigin::Source( + "encoder/config.json#type_vocab_size".into() + )) + ); + assert_eq!( + structure.provenance.get("head_dim"), + Some(&FieldOrigin::Derived( + "hidden_size / num_attention_heads".into() + )) + ); +} + +#[test] +fn hf_adapter_rejects_unknown_variants() { + for (name, config) in [ + ("model_type", { + let mut c = vela_encoder_config(); + c["model_type"] = json!("vela2-unified"); + c + }), + ("no schedule source", { + let mut c = vela_encoder_config(); + c.as_object_mut() + .unwrap() + .remove("global_attn_every_n_layers"); + c + }), + ("zero interval", { + // The derived schedule tests i % every; zero must be an + // error, not a division-by-zero panic. + let mut c = vela_encoder_config(); + c["global_attn_every_n_layers"] = json!(0); + c + }), + ("scaling type", { + let mut c = vela_encoder_config(); + c["rope_scaling"]["rope_type"] = json!("linear"); + c + }), + ("scaling truncate", { + let mut c = vela_encoder_config(); + c["rope_scaling"]["truncate"] = json!(false); + c + }), + ("position embedding", { + let mut c = vela_encoder_config(); + c["position_embedding_type"] = json!("absolute"); + c + }), + ("norm eps", { + let mut c = vela_encoder_config(); + c.as_object_mut().unwrap().remove("norm_eps"); + c.as_object_mut().unwrap().remove("layer_norm_eps"); + c + }), + ("yarn factor", { + let mut c = vela_encoder_config(); + c["rope_scaling"]["factor"] = json!(1.0); + c + }), + ] { + assert!( + ModernBertConfigAdapter::from_value(&config, "config.json#encoder_config").is_err(), + "accepted {name}" + ); + } +} + +#[test] +fn structure_invariants_name_the_violated_field() { + let mut structure = parse(&family_config()); + structure.num_attention_heads = 15; + assert!(structure.validate().unwrap_err().to_string().contains("15")); + + let mut structure = parse(&family_config()); + structure.schedule.pop(); + let message = structure.validate().unwrap_err().to_string(); + assert!(message.contains("schedule"), "{message}"); + + for mutate in [ + ("local_window", 0), + ("intermediate_size", 0), + ("vocab_size", 0), + ("num_hidden_layers", 0), + ] { + let mut structure = parse(&family_config()); + match mutate.0 { + "local_window" => structure.local_window = mutate.1, + "intermediate_size" => structure.intermediate_size = mutate.1, + "vocab_size" => structure.vocab_size = mutate.1, + _ => structure.num_hidden_layers = mutate.1, + } + assert!(structure.validate().is_err(), "accepted {} = 0", mutate.0); + } + let mut structure = parse(&family_config()); + structure.rope_full.theta = 0.0; + assert!(structure.validate().is_err()); +} + +/// A test-local registered profile: the gate is family machinery, so +/// its accept/reject set is pinned here against a synthetic profile, +/// independent of any model crate's registered constants. +fn test_profile() -> gate::RegisteredProfile { + let structure = parse(&family_config()); + gate::RegisteredProfile { + name: "test-profile", + structure, + limits: gate::BackendLimits { + hidden_sizes: vec![1024], + batch_max: 16, + len_min: 16, + len_max: 512, + len_step: 16, + yarn_scaling: false, + }, + backend: "test_backend", + abi_version: 1, + } +} + +#[test] +fn admission_accepts_the_registered_structure_and_rejects_drift() { + let profile = test_profile(); + let structure = parse(&family_config()); + let identity = profile.admit(&structure).unwrap(); + assert_eq!(identity.schema_version, 1); + assert_eq!(identity.family, "modernbert"); + assert_eq!(identity.backend, "test_backend"); + assert_eq!(identity.abi_version, 1); + + let vela = parse(&vela_encoder_config()); + let message = profile.admit(&vela).unwrap_err().to_string(); + assert!(message.contains("hidden_size 768"), "{message}"); + assert!(message.contains("1024"), "{message}"); + + for mutate in [ + ("rope", { + let mut s = parse(&family_config()); + s.rope_sliding.theta = 160000.0; + s + }), + ("schedule", { + let mut s = parse(&family_config()); + s.schedule[1] = AttentionKind::Full; + s + }), + ("type_embedding_rows", { + let mut s = parse(&family_config()); + s.type_embedding_rows = None; + s + }), + ("norm_eps", { + let mut s = parse(&family_config()); + s.norm_eps = 1e-6; + s + }), + ("max_positions", { + let mut s = parse(&family_config()); + s.max_positions = 8192; + s + }), + ] { + assert!(profile.admit(&mutate.1).is_err(), "accepted {}", mutate.0); + } +} + +#[test] +fn backend_limits_check_the_request_envelope() { + let limits = &test_profile().limits; + limits.check_request(16, 512).unwrap(); + limits.check_request(1, 16).unwrap(); + for (batch, len) in [(32, 512), (3, 512), (1, 8), (1, 513), (1, 520)] { + assert!( + limits.check_request(batch, len).is_err(), + "accepted batch {batch} len {len}" + ); + } +} + +#[test] +fn structure_identity_is_stable_and_sensitive() { + let member = parse(&family_config()); + let again = parse(&family_config()); + assert_eq!( + structure_identity(&member, "laya_cuda", 1), + structure_identity(&again, "laya_cuda", 1) + ); + let mut drifted = member.clone(); + drifted.rope_sliding.theta = 160000.0; + assert_ne!( + structure_identity(&member, "laya_cuda", 1), + structure_identity(&drifted, "laya_cuda", 1) + ); + let vela = parse(&vela_encoder_config()); + assert_ne!( + structure_identity(&member, "laya_cuda", 1), + structure_identity(&vela, "laya_cuda", 1) + ); + assert_ne!( + structure_identity(&member, "laya_cuda", 1), + structure_identity(&member, "laya_cuda", 2) + ); + // Provenance is not part of identity: same semantics, same digest. + let mut reworded = member.clone(); + reworded.provenance.clear(); + assert_eq!( + structure_identity(&member, "laya_cuda", 1), + structure_identity(&reworded, "laya_cuda", 1) + ); +} + +#[test] +fn derived_geometry_matches_the_1024_member_inventory() { + let member = parse(&family_config()); + let tensors = geometry::encoder_tensors(&member); + assert_eq!(tensors.len(), 171); + assert_eq!(tensors.len(), 3 + 28 * 5 + 27 + 1); + let find = |name: &str| tensors.iter().find(|(n, _)| n == name).unwrap(); + assert_eq!( + find("encoder.embeddings.tok_embeddings.weight").1, + vec![50368, 1024] + ); + assert_eq!(find("encoder.layers.0.mlp_norm.weight").1, vec![1024]); + assert_eq!(find("encoder.layers.0.mlp.Wi.weight").1, vec![5248, 1024]); + let attn_norm_count = tensors + .iter() + .filter(|(name, _)| name.ends_with("attn_norm.weight")) + .count(); + assert_eq!(attn_norm_count, 27); + assert_eq!(tensors.last().unwrap().0, "type_emb.weight"); + assert_eq!(tensors.last().unwrap().1, vec![3, 1024]); + + let vela = parse(&vela_encoder_config()); + let vela_tensors = geometry::encoder_tensors(&vela); + assert_eq!(vela_tensors.len(), 3 + 22 * 5 + 21); + let find_vela = |name: &str| vela_tensors.iter().find(|(n, _)| n == name).unwrap(); + assert_eq!( + find_vela("encoder.layers.0.mlp.Wi.weight").1, + vec![2304, 768] + ); + assert!( + vela_tensors + .iter() + .all(|(name, _)| name != "type_emb.weight") + ); + + assert_eq!(geometry::rotary_half(&member), 32); + assert_eq!(geometry::rope_table_bytes(&member), 512 * 32 * 4); +} + +#[test] +fn encoder_scratch_sizes_follow_the_structure() { + let member = parse(&family_config()); + let scratch = EncoderScratch::new(1, 512, &member); + assert_eq!( + scratch.sizes(), + vec![ + 512 * 8, + 4, + 512 * 1024 * 4, + 512 * 1024 * 2, + 512 * 3072 * 2, + 512 * 1024 * 2, + 512 * 2624 * 2 + ] + ); + let total: usize = scratch.sizes().iter().sum(); + assert_eq!(scratch.total(), total); + let vela = parse(&vela_encoder_config()); + assert_eq!(EncoderScratch::new(1, 512, &vela).qkv, 512 * 3 * 768 * 2); +} + +#[test] +fn admitted_structures_carry_the_family_activation() { + // The frozen-equality check includes the activation; today the + // family admits gelu only, and both adapters reject anything else + // at the source, so this pins the field's presence in the IR. + let member = parse(&family_config()); + assert_eq!(member.activation, Activation::Gelu); + assert_eq!( + member.rope_full, + RopeSpec { + theta: 160000.0, + scaling: None + } + ); +} diff --git a/tests/qwen3_5/structure.rs b/tests/qwen3_5/structure.rs new file mode 100644 index 00000000..a27fe9db --- /dev/null +++ b/tests/qwen3_5/structure.rs @@ -0,0 +1,266 @@ +use omni_qwen3_5_native::model::Config; +use omni_qwen3_5_native::structure::{self, LayerKind, PrefixLayout, QwenHybridStructure}; +use serde_json::json; +use std::fs; +use std::path::Path; +use tempfile::tempdir; + +/// A minimal gate-passing structure; every test mutates from here. +fn hybrid( + hidden: usize, + intermediate: usize, + layers: usize, + heads: usize, + kv_heads: usize, + lin_k_heads: usize, + lin_v_heads: usize, +) -> QwenHybridStructure { + let (_keep, path) = fixture_dir(default_config()); + let mut s = structure::read_structure(&path).unwrap(); + s.hidden = hidden; + s.intermediate = intermediate; + s.heads = heads; + s.kv_heads = kv_heads; + s.lin_k_heads = lin_k_heads; + s.lin_v_heads = lin_v_heads; + s.layer_kinds = (0..layers) + .map(|i| { + if i % 4 == 3 { + LayerKind::FullAttention + } else { + LayerKind::GatedDeltaNet + } + }) + .collect(); + s +} + +fn vela_08b() -> QwenHybridStructure { + hybrid(1024, 3584, 24, 8, 2, 16, 16) +} + +fn vela_4b() -> QwenHybridStructure { + hybrid(2560, 9216, 32, 16, 4, 16, 32) +} + +fn vela_9b() -> QwenHybridStructure { + hybrid(4096, 12288, 32, 16, 4, 16, 32) +} + +fn default_config() -> serde_json::Value { + json!({ + "model_type": "qwen3_5_text", + "hidden_act": "silu", "attn_output_gate": true, "attention_bias": false, + "hidden_size": 5120, "intermediate_size": 17408, + "num_attention_heads": 24, "num_key_value_heads": 4, "head_dim": 256, + "num_hidden_layers": 4, + "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"], + "rms_norm_eps": 1e-6, + "linear_num_key_heads": 16, "linear_num_value_heads": 48, + "linear_key_head_dim": 128, "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "max_position_embeddings": 262144, + "rope_parameters": { + "rope_type": "default", "rope_theta": 10000000.0, + "partial_rotary_factor": 0.25, "mrope_interleaved": true, + "mrope_section": [11, 11, 10] + } + }) +} + +fn fixture_dir(config: serde_json::Value) -> (tempfile::TempDir, std::path::PathBuf) { + let dir = tempdir().unwrap(); + fs::write(dir.path().join("config.json"), config.to_string()).unwrap(); + let path = dir.path().to_path_buf(); + (dir, path) +} + +#[test] +fn adapter_reads_the_checkpoint_and_records_provenance() { + let (_keep, path) = fixture_dir(default_config()); + let s = structure::read_structure(&path).unwrap(); + assert_eq!(s.hidden, 5120); + assert_eq!(s.head_dim, 256); + assert_eq!(s.rotary_half, 32); + assert_eq!(s.image_token_id, None); + assert_eq!( + s.layer_kinds, + [ + LayerKind::GatedDeltaNet, + LayerKind::GatedDeltaNet, + LayerKind::GatedDeltaNet, + LayerKind::FullAttention + ] + ); + assert_eq!( + s.provenance.get("hidden"), + Some(&"config.json#hidden_size".to_string()) + ); + assert_eq!(s.hidden_act.as_deref(), Some("silu")); + assert_eq!(s.conv_kernel_dim, 4); + structure::check_abi(&s).unwrap(); +} + +#[test] +fn adapter_rejects_broken_sources() { + for (name, config) in [ + ("missing layer_types", { + let mut c = default_config(); + c.as_object_mut().unwrap().remove("layer_types"); + c + }), + ("unknown kind", { + let mut c = default_config(); + c["layer_types"][0] = json!("sparse_attention"); + c + }), + ("length mismatch", { + let mut c = default_config(); + c["layer_types"] = json!(["linear_attention"]); + c + }), + ("missing conv", { + let mut c = default_config(); + c.as_object_mut().unwrap().remove("linear_conv_kernel_dim"); + c + }), + ] { + assert!( + { + let (_keep, path) = fixture_dir(config); + structure::read_structure(&path).is_err() + }, + "accepted {name}" + ); + } +} + +#[test] +fn the_frozen_gate_accepts_the_three_vela_hybrid_backbones() { + for (name, s) in [("0.8B", vela_08b()), ("4B", vela_4b()), ("9B", vela_9b())] { + structure::check_abi(&s).unwrap_or_else(|e| panic!("{name} rejected: {e}")); + } +} + +#[test] +fn the_frozen_gate_rejects_every_capability_violation() { + type Mutation = Box; + let rejects: Vec<(&str, Mutation)> = vec![ + ( + "head_dim", + Box::new(|s: &mut QwenHybridStructure| s.head_dim = 128), + ), + ("lin_k_dim", Box::new(|s| s.lin_k_dim = 64)), + ("lin_v_dim", Box::new(|s| s.lin_v_dim = 64)), + ("conv", Box::new(|s| s.conv_kernel_dim = 5)), + ( + "hidden_act", + Box::new(|s| s.hidden_act = Some("relu".into())), + ), + ("output gate", Box::new(|s| s.attn_output_gate = false)), + ("attention bias", Box::new(|s| s.attention_bias = true)), + ("rope type", Box::new(|s| s.rope_type = "linear".into())), + ("mrope", Box::new(|s| s.mrope_interleaved = false)), + ("mrope sum", Box::new(|s| s.mrope_section = [11, 11, 11])), + ("rotary", Box::new(|s| s.rotary_half = 16)), + ( + "max positions upper", + Box::new(|s| s.max_positions = 1 << 25), + ), + ("max positions zero", Box::new(|s| s.max_positions = 0)), + ("rope theta zero", Box::new(|s| s.rope_theta = 0.0)), + ("rope theta negative", Box::new(|s| s.rope_theta = -1.0)), + ("rope theta nan", Box::new(|s| s.rope_theta = f64::NAN)), + ( + "heads divisibility", + Box::new(|s: &mut QwenHybridStructure| { + s.heads = 9; + s.kv_heads = 2; + }), + ), + ( + "linear heads divisibility", + Box::new(|s: &mut QwenHybridStructure| { + s.lin_v_heads = 17; + }), + ), + ("hidden alignment", Box::new(|s| s.hidden = 1025)), + ]; + assert_eq!(rejects.len(), 19); + for (name, mutate) in rejects { + let (_keep, path) = fixture_dir(default_config()); + let mut s = structure::read_structure(&path).unwrap(); + mutate(&mut s); + assert!(structure::check_abi(&s).is_err(), "gate accepted {name}"); + } + let (_keep, path) = fixture_dir(default_config()); + let mut s = structure::read_structure(&path).unwrap(); + s.rope_type = "linear".into(); + let message = structure::check_abi(&s).unwrap_err().to_string(); + assert!(message.contains("\"linear\""), "{message}"); +} + +#[test] +fn config_load_rejects_a_defective_checkpoint_end_to_end() { + let mut config = default_config(); + config["hidden_act"] = json!("relu"); + let (_keep, path) = fixture_dir(config); + assert!(Config::load(&path).is_err()); +} + +#[test] +fn the_open_jev_27b_checkpoint_passes_the_gate_unchanged() { + let dir = Path::new(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../../tests/qwen3_5/data" + )); + let s = structure::read_structure(dir).unwrap(); + structure::check_abi(&s).unwrap(); + assert_eq!(s.hidden, 5120); + assert_eq!(s.layer_kinds.len(), 64); + assert_eq!(s.full_attention_layers(), 16); + assert_eq!(s.gdn_layers(), 48); + // The facade projects the same values it always did. + let cfg = Config::load(dir).unwrap(); + assert_eq!(cfg.hidden, s.hidden); + assert_eq!(cfg.full_attention.len(), s.layer_kinds.len()); + assert_eq!( + cfg.full_attention.iter().filter(|&&f| f).count(), + s.full_attention_layers() + ); +} + +#[test] +fn prefix_layout_sizes_derive_from_the_structure() { + let layout = PrefixLayout::of(&vela_08b(), 64); + assert_eq!(layout.full_layers, 6); + assert_eq!(layout.gdn_layers, 18); + assert_eq!(layout.full_kv_bytes, 2 * 64 * 2 * 256 * 2); + assert_eq!(layout.gdn_state_bytes, 16 * 128 * 128 * 4); + assert_eq!(layout.conv_tail_bytes, 3 * 8224 * 2); + assert_eq!( + layout.total_bytes(), + 6 * layout.full_kv_bytes + 18 * (layout.gdn_state_bytes + layout.conv_tail_bytes) + ); + // Grows linearly in tokens through the full-attention rows only. + let longer = PrefixLayout::of(&vela_08b(), 128); + assert_eq!(longer.full_kv_bytes, 2 * layout.full_kv_bytes); + assert_eq!(longer.gdn_state_bytes, layout.gdn_state_bytes); +} + +#[test] +fn identity_is_stable_and_sensitive() { + let a = structure::structure_identity(&vela_08b()); + let b = structure::structure_identity(&vela_08b()); + assert_eq!(a, b); + assert_eq!(a.abi_version, 9); + assert_eq!(a.family, "qwen3_5_hybrid"); + assert_ne!(a, structure::structure_identity(&vela_4b())); + let mut drifted = vela_08b(); + drifted.rope_theta = 5e6; + assert_ne!(a, structure::structure_identity(&drifted)); + // Provenance is not part of identity. + let mut reworded = vela_08b(); + reworded.provenance.clear(); + assert_eq!(a, structure::structure_identity(&reworded)); +} diff --git a/tests/vela/native/common/mod.rs b/tests/vela/native/common/mod.rs new file mode 100644 index 00000000..16404f1f --- /dev/null +++ b/tests/vela/native/common/mod.rs @@ -0,0 +1,22 @@ +//! Shared fixtures for the vela-native test targets. The root config is +//! the pinned vllm-sr/Vela-2.0-0.3B@d6f03aa9 config, structure-relevant +//! members only; recipe fields stay worker-side. Regenerate it together +//! with the census and yarn references when the pinned revision moves. + +use serde_json::Value; +use std::fs; +use std::path::PathBuf; +use tempfile::{TempDir, tempdir}; + +pub fn pinned_config() -> Value { + serde_json::from_str(include_str!("../data/vela-config.json")).unwrap() +} + +/// Writes the pinned root config into a temp bundle dir. Keep the +/// returned `TempDir` alive for as long as the path is read. +pub fn bundle_dir() -> (TempDir, PathBuf) { + let dir = tempdir().unwrap(); + fs::write(dir.path().join("config.json"), pinned_config().to_string()).unwrap(); + let path = dir.path().to_path_buf(); + (dir, path) +} diff --git a/tests/vela/native/data/pinned-checkpoint-census.json b/tests/vela/native/data/pinned-checkpoint-census.json new file mode 100644 index 00000000..620b7351 --- /dev/null +++ b/tests/vela/native/data/pinned-checkpoint-census.json @@ -0,0 +1,1006 @@ +{ +"cls_mlp.0.bias": { +"dtype": "F32", +"shape": [ +1536 +] +}, +"cls_mlp.0.weight": { +"dtype": "F32", +"shape": [ +1536, +768 +] +}, +"cls_mlp.3.bias": { +"dtype": "F32", +"shape": [ +1 +] +}, +"cls_mlp.3.weight": { +"dtype": "F32", +"shape": [ +1, +1536 +] +}, +"encoder.embeddings.norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.embeddings.tok_embeddings.weight": { +"dtype": "F32", +"shape": [ +256008, +768 +] +}, +"encoder.final_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.0.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.0.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.0.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.0.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.0.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.1.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.1.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.1.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.1.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.1.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.1.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.10.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.10.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.10.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.10.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.10.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.10.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.11.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.11.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.11.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.11.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.11.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.11.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.12.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.12.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.12.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.12.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.12.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.12.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.13.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.13.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.13.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.13.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.13.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.13.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.14.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.14.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.14.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.14.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.14.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.14.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.15.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.15.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.15.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.15.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.15.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.15.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.16.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.16.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.16.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.16.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.16.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.16.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.17.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.17.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.17.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.17.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.17.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.17.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.18.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.18.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.18.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.18.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.18.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.18.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.19.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.19.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.19.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.19.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.19.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.19.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.2.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.2.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.2.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.2.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.2.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.2.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.20.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.20.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.20.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.20.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.20.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.20.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.21.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.21.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.21.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.21.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.21.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.21.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.3.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.3.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.3.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.3.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.3.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.3.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.4.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.4.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.4.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.4.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.4.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.4.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.5.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.5.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.5.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.5.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.5.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.5.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.6.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.6.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.6.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.6.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.6.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.6.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.7.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.7.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.7.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.7.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.7.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.7.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.8.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.8.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.8.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.8.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.8.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.8.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.9.attn.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +768 +] +}, +"encoder.layers.9.attn.Wqkv.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.9.attn_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"encoder.layers.9.mlp.Wi.weight": { +"dtype": "F32", +"shape": [ +2304, +768 +] +}, +"encoder.layers.9.mlp.Wo.weight": { +"dtype": "F32", +"shape": [ +768, +1152 +] +}, +"encoder.layers.9.mlp_norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"log_tau": { +"dtype": "F32", +"shape": [] +}, +"log_tau_span": { +"dtype": "F32", +"shape": [] +}, +"norm.bias": { +"dtype": "F32", +"shape": [ +768 +] +}, +"norm.weight": { +"dtype": "F32", +"shape": [ +768 +] +}, +"w_e.bias": { +"dtype": "F32", +"shape": [ +256 +] +}, +"w_e.weight": { +"dtype": "F32", +"shape": [ +256, +768 +] +}, +"w_o.bias": { +"dtype": "F32", +"shape": [ +256 +] +}, +"w_o.weight": { +"dtype": "F32", +"shape": [ +256, +768 +] +}, +"w_p.bias": { +"dtype": "F32", +"shape": [ +256 +] +}, +"w_p.weight": { +"dtype": "F32", +"shape": [ +256, +768 +] +}, +"w_q.bias": { +"dtype": "F32", +"shape": [ +256 +] +}, +"w_q.weight": { +"dtype": "F32", +"shape": [ +256, +768 +] +}, +"w_t.bias": { +"dtype": "F32", +"shape": [ +256 +] +}, +"w_t.weight": { +"dtype": "F32", +"shape": [ +256, +768 +] +} +} \ No newline at end of file diff --git a/tests/vela/native/data/vela-config.json b/tests/vela/native/data/vela-config.json new file mode 100644 index 00000000..c747b765 --- /dev/null +++ b/tests/vela/native/data/vela-config.json @@ -0,0 +1,33 @@ +{ + "model_type": "vela2-unified", + "format_version": 1, + "encoder_config": { + "model_type": "modernbert", + "hidden_activation": "gelu", + "attention_bias": false, + "mlp_bias": false, + "norm_bias": false, + "hidden_size": 768, + "intermediate_size": 1152, + "num_attention_heads": 12, + "num_hidden_layers": 22, + "vocab_size": 256008, + "norm_eps": 1e-5, + "local_attention": 128, + "global_attn_every_n_layers": 3, + "global_rope_theta": 160000, + "local_rope_theta": 160000, + "max_position_embeddings": 32768, + "position_embedding_type": "sans_pos", + "rope_scaling": { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 8192, + "beta_fast": 32.0, + "beta_slow": 1.0, + "truncate": true + } + }, + "max_length": 8192, + "proj_dim": 256 +} diff --git a/tests/vela/native/data/yarn-reference.json b/tests/vela/native/data/yarn-reference.json new file mode 100644 index 00000000..e8a26adf --- /dev/null +++ b/tests/vela/native/data/yarn-reference.json @@ -0,0 +1,1010 @@ +{ +"source": "transformers 5.19.0 ModernBertRotaryEmbedding, ModernBertConfig(Vela-2.0-0.3B@d6f03aa9 encoder_config), CPU float32", +"positions": [ +0, +1, +2, +63, +64, +1024, +4095, +8191, +8192, +8193, +12288, +16384, +24576, +32767 +], +"attention_scaling": 1.138629436111989, +"inv_freq": [ +1.0, +0.687656044960022, +0.472870796918869, +0.32517245411872864, +0.22360679507255554, +0.15376456081867218, +0.1057371199131012, +0.07271077483892441, +0.05000000074505806, +0.03438280150294304, +0.022031478583812714, +0.0140415383502841, +0.0088934525847435, +0.005591438151896, +0.003484518965706229, +0.002148272702470422, +0.0013068180996924639, +0.0007814273121766746, +0.0004567502182908356, +0.0002586598857305944, +0.00013975425099488348, +9.610284905647859e-05, +6.608570402022451e-05, +4.5444234274327755e-05, +3.125000148429535e-05, +2.148924977518618e-05, +1.4777213436900638e-05, +1.0161638783756644e-05, +6.987712822592584e-06, +4.8051424528239295e-06, +3.3042852010112256e-06, +2.272211759191123e-06 +], +"cos_half": [ +[ +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292 +], +[ +0.6152041554450989, +0.8798598647117615, +1.0136812925338745, +1.0789602994918823, +1.110282063484192, +1.1251952648162842, +1.1322702169418335, +1.1356208324432373, +1.1372064352035522, +1.1379565000534058, +1.1383531093597412, +1.138517141342163, +1.1385844945907593, +1.1386116743087769, +1.138622522354126, +1.1386268138885498, +1.1386284828186035, +1.1386290788650513, +1.1386293172836304, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292 +], +[ +-0.4738370478153229, +0.22116917371749878, +0.6662595272064209, +0.9062067270278931, +1.0266516208648682, +1.0852099657058716, +1.1132636070251465, +1.1266111135482788, +1.1329410076141357, +1.1359384059906006, +1.1375242471694946, +1.1381804943084717, +1.1384493112564087, +1.138558268547058, +1.1386017799377441, +1.1386189460754395, +1.1386256217956543, +1.1386281251907349, +1.1386289596557617, +1.1386293172836304, +1.1386293172836304, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292, +1.13862943649292 +], +[ +1.1225708723068237, +0.8995316624641418, +-0.061762042343616486, +-0.07454098761081696, +0.056838519871234894, +-1.0996575355529785, +1.0581409931182861, +-0.14942322671413422, +-1.1385891437530518, +-0.6385136246681213, +0.20699888467788696, +0.7214195728302002, +0.964535653591156, +1.0687121152877808, +1.1113035678863525, +1.128217101097107, +1.134772777557373, +1.1372499465942383, +1.1381580829620361, +1.1384782791137695, +1.1385853290557861, +1.1386085748672485, +1.1386195421218872, +1.1386247873306274, +1.138627290725708, +1.138628363609314, +1.1386289596557617, +1.1386291980743408, +1.1386293172836304, +1.1386293172836304, +1.13862943649292, +1.13862943649292 +], +[ +0.44618019461631775, +1.1381930112838745, +0.46283289790153503, +-0.4336148500442505, +-0.1967509239912033, +-1.0414478778839111, +1.00785231590271, +-0.0670258104801178, +-1.1366877555847168, +-0.6705442667007446, +0.18228302896022797, +0.7089793086051941, +0.9591161608695984, +1.0664987564086914, +1.1104328632354736, +1.1278843879699707, +1.134649395942688, +1.137205719947815, +1.1381429433822632, +1.138473391532898, +1.1385838985443115, +1.1386079788208008, +1.138619303703308, +1.138624668121338, +1.1386271715164185, +1.138628363609314, +1.1386289596557617, +1.1386291980743408, +1.1386293172836304, +1.1386293172836304, +1.13862943649292, +1.13862943649292 +], +[ +1.124229907989502, +1.0286998748779297, +1.042241096496582, +1.138038158416748, +-1.064469814300537, +1.0593875646591187, +0.12515044212341309, +0.6693502068519592, +0.6765803694725037, +-0.906120777130127, +-0.9591747522354126, +-0.2721683084964752, +-1.0815833806991577, +0.9661869406700134, +-1.0366042852401733, +-0.6699293255805969, +0.262479692697525, +0.7931424975395203, +1.0163429975509644, +1.0989222526550293, +1.126989722251892, +1.1331204175949097, +1.1360232830047607, +1.1373968124389648, +1.1380465030670166, +1.138353705406189, +1.1384990215301514, +1.1385678052902222, +1.1386003494262695, +1.138615608215332, +1.1386228799819946, +1.1386263370513916 +], +[ +-0.07512220740318298, +0.5322462320327759, +0.4286642074584961, +1.0232322216033936, +-0.11872357875108719, +0.25229644775390625, +0.9729129076004028, +-0.8703641891479492, +-0.9727998971939087, +-0.9560906291007996, +-0.7190672159194946, +0.6609950661659241, +0.3260018825531006, +-0.7026075720787048, +-0.14979302883148193, +-0.9216501116752625, +0.6790989637374878, +-1.1366914510726929, +-0.3360483646392822, +0.5574265122413635, +0.9572014212608337, +1.0515891313552856, +1.09718918800354, +1.118970274925232, +1.1293190717697144, +1.1342236995697021, +1.1365453004837036, +1.137643814086914, +1.1381633281707764, +1.1384090185165405, +1.1385252475738525, +1.1385802030563354 +], +[ +-0.7359992265701294, +-1.0927605628967285, +-1.088085412979126, +0.9505313038825989, +-1.1385061740875244, +-1.0900436639785767, +0.6277983784675598, +0.273050993680954, +0.4724070727825165, +0.5024028420448303, +-0.2057957947254181, +-0.38627222180366516, +-0.946363627910614, +-0.2777009606361389, +-1.0981756448745728, +0.35573408007621765, +-0.32715022563934326, +1.1307802200317383, +-0.9399774074554443, +-0.5930944681167603, +0.47058945894241333, +0.803697943687439, +0.9758458137512207, +1.0606529712677002, +1.1015311479568481, +1.121036171913147, +1.1302987337112427, +1.1346875429153442, +1.1367648839950562, +1.1377476453781128, +1.1382124423980713, +1.138432264328003 +], +[ +0.33339276909828186, +-1.0475496053695679, +-1.1214704513549805, +1.1009937524795532, +-1.1064395904541016, +-1.1275886297225952, +0.724525511264801, +0.352614164352417, +0.42005088925361633, +0.5372417569160461, +-0.18107281625270844, +-0.40126997232437134, +-0.9406938552856445, +-0.2838718891143799, +-1.097120761871338, +0.35805627703666687, +-0.3257249891757965, +1.1306756734848022, +-0.9396838545799255, +-0.5933458209037781, +0.47044458985328674, +0.8036203980445862, +0.9758070707321167, +1.0606341361999512, +1.1015220880508423, +1.1210318803787231, +1.1302967071533203, +1.1346865892410278, +1.136764407157898, +1.1377474069595337, +1.1382123231887817, +1.138432264328003 +], +[ +1.0962649583816528, +-0.5264052152633667, +-0.9086950421333313, +1.1360465288162231, +-1.0192683935165405, +-1.138510823249817, +0.8131633996963501, +0.4303148686885834, +0.36664533615112305, +0.5714452266693115, +-0.15626192092895508, +-0.41618862748146057, +-0.9349496364593506, +-0.2900339663028717, +-1.0960525274276733, +0.3603768050670624, +-0.32429924607276917, +1.1305702924728394, +-0.9393899440765381, +-0.5935972332954407, +0.47029972076416016, +0.8035429120063782, +0.9757682085037231, +1.0606153011322021, +1.101513147354126, +1.1210275888442993, +1.130294680595398, +1.1346856355667114, +1.1367639303207397, +1.1377471685409546, +1.1382122039794922, +1.1384321451187134 +], +[ +-0.37935808300971985, +0.6471661329269409, +0.29339906573295593, +1.0545871257781982, +-0.39833804965019226, +-0.23624123632907867, +0.2807825803756714, +0.3507510721683502, +0.24699416756629944, +0.05504527688026428, +0.9731765389442444, +-1.1045781373977661, +-0.8903589248657227, +1.0454181432724, +0.44996723532676697, +0.34253257513046265, +-1.069512128829956, +-1.1207594871520996, +0.8920308947563171, +-1.1378576755523682, +-0.16621750593185425, +0.43277212977409363, +0.7833817601203918, +0.965664803981781, +1.0557070970535278, +1.0991625785827637, +1.119909405708313, +1.129764437675476, +1.1344345808029175, +1.1366451978683472, +1.1376910209655762, +1.1381856203079224 +], +[ +-0.9433934092521667, +0.7888809442520142, +1.0705108642578125, +0.9905748963356018, +1.0116900205612183, +1.0946804285049438, +-0.2165784388780594, +-0.9202321767807007, +-0.8287081718444824, +-0.6316537857055664, +-1.081038475036621, +-0.8558024168014526, +0.4157041013240814, +-0.9970850944519043, +0.9756210446357727, +-0.9134389162063599, +-0.9522706866264343, +1.106925368309021, +0.4123680293560028, +-0.5202381610870361, +-0.7498846650123596, +-0.004272983875125647, +0.5339065790176392, +0.8373335599899292, +0.9926188588142395, +1.0687830448150635, +1.105420470237732, +1.1228853464126587, +1.1311755180358887, +1.13510262966156, +1.1369612216949463, +1.1378405094146729 +], +[ +-0.8858473896980286, +-0.40296605229377747, +-0.9874248504638672, +0.8148664236068726, +-0.85992032289505, +-1.0405994653701782, +-1.000149130821228, +-0.9225339889526367, +-1.0314723253250122, +-1.1333072185516357, +0.5249013900756836, +1.0044609308242798, +0.2538149356842041, +0.781045138835907, +-0.7829904556274414, +-0.932542085647583, +0.8705511689186096, +1.067710518836975, +0.2590492367744446, +1.1355434656143188, +-1.0901004076004028, +-0.8096519708633423, +-0.06068970635533333, +0.49931955337524414, +0.8190178871154785, +0.9834980964660645, +1.0643649101257324, +1.1033074855804443, +1.1218808889389038, +1.1306992769241333, +1.1348772048950195, +1.1368545293807983 +], +[ +1.118433952331543, +0.6868389248847961, +1.1105233430862427, +0.24214112758636475, +0.8486597537994385, +0.862530529499054, +-1.0054718255996704, +0.4266851544380188, +0.010710631497204304, +-0.4013756811618805, +0.8989120125770569, +0.16365931928157806, +-0.8281826376914978, +0.6130115389823914, +0.5367735028266907, +0.32928648591041565, +0.45283031463623047, +1.013983964920044, +-0.839590311050415, +-0.6629980802536011, +-0.15106093883514404, +-1.1385982036590576, +-0.6378665566444397, +0.09295102953910828, +0.5920644998550415, +0.867828905582428, +1.007738709449768, +1.0760924816131592, +1.1089129447937012, +1.1245449781417847, +1.1319620609283447, +1.1354750394821167 +] +], +"sin_half": [ +[ +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0, +0.0 +], +[ +0.9581236243247986, +0.7227196097373962, +0.5185818672180176, +0.36376047134399414, +0.2524888515472412, +0.1743917465209961, +0.12017118185758591, +0.08271769434213638, +0.05690775439143181, +0.03914155811071396, +0.02508366107940674, +0.015987584367394447, +0.010126213543117046, +0.006366542540490627, +0.003967567812651396, +0.002446084748953581, +0.0014879811787977815, +0.000889756076503545, +0.0005200692103244364, +0.0002945177548099309, +0.0001591283071320504, +0.00010942553490167484, +7.524713146267459e-05, +5.174414400244132e-05, +3.558217213139869e-05, +2.4468292394885793e-05, +1.6825770217110403e-05, +1.1570341484912205e-05, +7.95641517470358e-06, +5.471276836033212e-06, +3.762356300285319e-06, +2.587207291071536e-06 +], +[ +1.0353528261184692, +1.1169427633285522, +0.9233499765396118, +0.6893956661224365, +0.49240580201148987, +0.34466835856437683, +0.23900005221366882, +0.16499827802181244, +0.11367326974868774, +0.07823684811592102, +0.0501551479101181, +0.0319720134139061, +0.020251624286174774, +0.012732886709272861, +0.007935087196528912, +0.004892158322036266, +0.002975959563627839, +0.001779511570930481, +0.001040138304233551, +0.0005890355096198618, +0.0003182566142641008, +0.00021885106980334967, +0.00015049426292534918, +0.00010348828800488263, +7.116434426279739e-05, +4.893658478977159e-05, +3.3651540434220806e-05, +2.314068296982441e-05, +1.591283034940716e-05, +1.0942553672066424e-05, +7.524712600570638e-06, +5.174414582143072e-06 +], +[ +0.1905561238527298, +-0.6980829238891602, +-1.1369531154632568, +1.1361868381500244, +1.1372098922729492, +-0.2953476309776306, +0.4204932153224945, +-1.1287823915481567, +-0.009572848677635193, +0.9427499175071716, +1.1196554899215698, +0.8809261322021484, +0.6051015853881836, +0.39285045862197876, +0.24795448780059814, +0.15363341569900513, +0.09363697469234467, +0.05603199824690819, +0.03275984153151512, +0.018553797155618668, +0.01002495363354683, +0.0068937670439481735, +0.004740555305033922, +0.003259876277297735, +0.0022416755091398954, +0.0015415019588544965, +0.0010600233217701316, +0.000728931394405663, +0.0005012541660107672, +0.0003446904302109033, +0.00023702844919171184, +0.0001629940525162965 +], +[ +1.0475687980651855, +0.03152431175112724, +-1.040318489074707, +1.052832007408142, +1.1215016841888428, +-0.4602859914302826, +0.5298212766647339, +-1.1366549730300903, +-0.06646657735109329, +0.9202430844306946, +1.1239439249038696, +0.8909687399864197, +0.6136556267738342, +0.3988199234008789, +0.25182533264160156, +0.15605679154396057, +0.09511983394622803, +0.05692065507173538, +0.03327969089150429, +0.018848275765776634, +0.010184075683355331, +0.007003189995884895, +0.004815801978111267, +0.0033116203267127275, +0.0022772573865950108, +0.0015659701311960816, +0.001076849177479744, +0.0007405017386190593, +0.0005092105711810291, +0.00035016171750612557, +0.0002407908032182604, +0.0001655812666285783 +], +[ +-0.1805107742547989, +0.488112211227417, +0.4584871828556061, +-0.036689698696136475, +0.4042043685913086, +0.41734281182289124, +1.1317306756973267, +-0.9211119413375854, +0.9158143997192383, +-0.6895085573196411, +-0.6135640740394592, +1.10562264919281, +0.3558853268623352, +-0.6024614572525024, +-0.4710928201675415, +0.9206908941268921, +1.107962727546692, +0.8169467449188232, +0.5133458375930786, +0.2980722486972809, +0.1623917520046234, +0.11187097430229187, +0.07699426263570786, +0.0529668815433979, +0.03642992675304413, +0.025053510442376137, +0.017228931188583374, +0.011847815476357937, +0.008147300221025944, +0.005602564662694931, +0.0038526456337422132, +0.0026492977049201727 +], +[ +-1.1361485719680786, +1.0065739154815674, +1.0548573732376099, +-0.49947258830070496, +-1.132422924041748, +1.110325813293457, +-0.5915381908416748, +0.7341274619102478, +-0.5917240381240845, +0.6183589100837708, +0.8828473091125488, +0.9271259307861328, +-1.0909627676010132, +-0.8960019946098328, +1.1287333965301514, +0.6686090230941772, +-0.9139483571052551, +-0.06640370190143585, +1.0879101753234863, +0.9928507804870605, +0.6166380643844604, +0.43662029504776, +0.3043896555900574, +0.21067138016223907, +0.1453116238117218, +0.10006839036941528, +0.06885948032140732, +0.047366876155138016, +0.03257707878947258, +0.022403430193662643, +0.015406379476189613, +0.010594460181891918 +], +[ +-0.8687819838523865, +0.31992387771606445, +0.3354804813861847, +-0.6268709897994995, +-0.0167547594755888, +0.3290621042251587, +-0.9499190449714661, +-1.1054049730300903, +1.036005973815918, +-1.021796703338623, +-1.1198772192001343, +1.071107268333435, +-0.6331452131271362, +1.1042460203170776, +-0.30081096291542053, +-1.0816330909729004, +-1.0906189680099487, +0.13346514105796814, +-0.6425880193710327, +0.9719650149345398, +1.0368329286575317, +0.8065646290779114, +0.5866872668266296, +0.4141162037849426, +0.28828129172325134, +0.19938646256923676, +0.1374835968017578, +0.09466326981782913, +0.0651354193687439, +0.04480365663766861, +0.03081369958817959, +0.021190589293837547 +], +[ +-1.0887268781661987, +-0.44622525572776794, +-0.1969287246465683, +-0.29032671451568604, +-0.2688279449939728, +0.1581796556711197, +-0.878373384475708, +-1.0826542377471924, +1.058316707611084, +-1.0039165019989014, +-1.1241395473480225, +1.0655794143676758, +-0.6415387988090515, +1.1026756763458252, +-0.3046359717845917, +-1.0808666944503784, +-1.0910454988479614, +0.13434885442256927, +-0.6430173516273499, +0.9718115329742432, +1.0368987321853638, +0.8066418766975403, +0.5867518186569214, +0.4141644239425659, +0.2883157432079315, +0.1994105577468872, +0.13750030100345612, +0.09467480331659317, +0.06514336168766022, +0.04480912163853645, +0.030817458406090736, +0.02119317837059498 +], +[ +-0.30770131945610046, +-1.0096408128738403, +-0.6861124634742737, +0.07665055245161057, +-0.5075125694274902, +-0.01643751561641693, +-0.7970208525657654, +-1.0541850328445435, +1.0779833793640137, +-0.9848489165306091, +-1.1278560161590576, +1.059841513633728, +-0.6498815417289734, +1.1010710000991821, +-0.30845728516578674, +-1.0800951719284058, +-1.0914701223373413, +0.13523246347904205, +-0.6434465050697327, +0.9716580510139465, +1.0369644165039062, +0.8067191243171692, +0.5868162512779236, +0.4142126142978668, +0.2883501648902893, +0.19943465292453766, +0.13751700520515442, +0.09468633681535721, +0.06515131145715714, +0.044814590364694595, +0.030821220949292183, +0.02119576558470726 +], +[ +-1.0735756158828735, +-0.9368313550949097, +-1.1001790761947632, +-0.4293286204338074, +1.0666788816452026, +-1.1138523817062378, +-1.1034663915634155, +1.0832593441009521, +-1.111517310142517, +1.1372981071472168, +0.5911043882369995, +0.27637675404548645, +0.7097449898719788, +-0.45119622349739075, +-1.0459476709365845, +1.085886001586914, +-0.39066722989082336, +-0.2009357362985611, +-0.7076425552368164, +-0.041914887726306915, +1.1264318227767944, +1.0531786680221558, +0.8263110518455505, +0.6032980680465698, +0.4265671968460083, +0.29718443751335144, +0.20562073588371277, +0.1418071836233139, +0.09764833748340607, +0.06719198822975159, +0.04621913284063339, +0.03178746998310089 +], +[ +-0.6375624537467957, +0.8210626840591431, +0.387922078371048, +-0.5614608526229858, +0.5224559903144836, +-0.3132917284965515, +-1.1178419589996338, +-0.670559287071228, +0.7808455228805542, +-0.9473597407341003, +0.35753706097602844, +-0.7510520815849304, +1.0600316524505615, +-0.5498165488243103, +0.5870609283447266, +-0.6797841191291809, +0.6242254972457886, +0.2668207287788391, +1.0613338947296143, +-1.0128322839736938, +0.8568255305290222, +1.13862144947052, +1.0056941509246826, +0.7715888619422913, +0.5578393340110779, +0.3926573097705841, +0.27298811078071594, +0.18869392573833466, +0.13007332384586334, +0.08954882621765137, +0.061612337827682495, +0.04237901419401169 +], +[ +0.7153680324554443, +-1.0649391412734985, +-0.5669825673103333, +-0.7952796220779419, +-0.7463337182998657, +0.4622010588645935, +-0.5442230105400085, +0.6673889756202698, +-0.4822259247303009, +0.10996182262897491, +1.0104234218597412, +-0.5362231135368347, +-1.1099797487258911, +-0.8285199999809265, +-0.826681911945343, +0.6533316969871521, +0.7339057326316833, +0.39556437730789185, +-1.1087697744369507, +0.08377296477556229, +-0.32887381315231323, +0.8005877137184143, +1.1370108127593994, +1.0233068466186523, +0.7910036444664001, +0.5737669467926025, +0.40448033809661865, +0.2814062237739563, +0.1945771872997284, +0.13414978981018066, +0.09236208349466324, +0.06355016678571701 +], +[ +0.2135004848241806, +0.9081460237503052, +0.251426100730896, +-1.1125847101211548, +0.7591137290000916, +-0.7433156371116638, +0.5343251824378967, +1.0556594133377075, +-1.1385791301727295, +1.0655395984649658, +-0.6988806128501892, +1.126806378364563, +0.7814028859138489, +0.9595279097557068, +1.0041669607162476, +1.0899759531021118, +-1.0447112321853638, +0.5179899334907532, +0.769132673740387, +0.9256946444511414, +-1.1285643577575684, +-0.008436508476734161, +0.9431877732276917, +1.1348291635513306, +0.9725927114486694, +0.7371227741241455, +0.530037522315979, +0.3721587061882019, +0.2584358751773834, +0.178537517786026, +0.12304040789604187, +0.08469671756029129 +] +] +} \ No newline at end of file diff --git a/tests/vela/native/executor.rs b/tests/vela/native/executor.rs new file mode 100644 index 00000000..561cb8f5 --- /dev/null +++ b/tests/vela/native/executor.rs @@ -0,0 +1,39 @@ +//! Device-gated encoder execution: runs only when a checkpoint bundle +//! and the regenerated CUDA library are provided, on the hardware the +//! bundle was built for. +//! +//! VELA_NATIVE_CHECKPOINT= +//! VELA_NATIVE_LIBRARY= + +use omni_vela_native::executor::Encoder; +use std::path::PathBuf; + +fn gated(name: &str) -> Option { + std::env::var_os(name).map(PathBuf::from) +} + +#[test] +fn encoder_runs_and_produces_finite_hidden() { + let (Some(checkpoint), Some(library)) = ( + gated("VELA_NATIVE_CHECKPOINT"), + gated("VELA_NATIVE_LIBRARY"), + ) else { + eprintln!("skipping: set VELA_NATIVE_CHECKPOINT and VELA_NATIVE_LIBRARY"); + return; + }; + let encoder = unsafe { Encoder::load(&checkpoint, &library) }.unwrap(); + assert_eq!(encoder.structure.hidden_size, 768); + // The embedding kernel indexes the token table unchecked, so ids + // outside the vocabulary must fail on the host before any launch. + let vocab = encoder.structure.vocab_size as i64; + for bad in [-1, vocab] { + assert!(encoder.encode(&[bad], &[0]).is_err(), "accepted id {bad}"); + } + let ids: Vec = (1..=97).collect(); + let hidden = encoder.encode(&ids, &[0, 47, 96]).unwrap(); + assert_eq!(hidden.len(), 3); + for row in &hidden { + assert_eq!(row.len(), 768); + assert!(row.iter().all(|v| v.is_finite()), "non-finite hidden"); + } +} diff --git a/tests/vela/native/profile.rs b/tests/vela/native/profile.rs new file mode 100644 index 00000000..fdf3c66b --- /dev/null +++ b/tests/vela/native/profile.rs @@ -0,0 +1,102 @@ +mod common; + +use omni_modernbert::plan::{self, EncoderPlanOptions, PlanStep}; +use omni_vela_native::{profile, read_structure}; +use serde_json::json; +use std::fs; +use tempfile::tempdir; + +use tempfile::TempDir; + +fn bundle_dir() -> (TempDir, std::path::PathBuf) { + common::bundle_dir() +} + +#[test] +fn the_pinned_bundle_parses_into_the_registered_profile() { + let (_keep, path) = bundle_dir(); + let (structure, meta) = read_structure(&path).unwrap(); + let identity = profile().admit(&structure).unwrap(); + assert_eq!(identity.backend, "laya_cuda"); + assert_eq!(identity.abi_version, 2); + assert_eq!(meta.model_type, "modernbert"); + let registered = profile().structure; + assert_eq!(structure.hidden_size, registered.hidden_size); + assert_eq!(structure.schedule, registered.schedule); + assert_eq!(structure.rope_full, registered.rope_full); + assert_eq!(structure.type_embedding_rows, None); +} + +#[test] +fn the_bundle_rejects_foreign_roots_and_encoders() { + for (name, mutate) in [("root class", "vela2-decoder"), ("encoder class", "bert")] { + let dir = tempdir().unwrap(); + let mut config = json!({ + "model_type": "vela2-unified", + "encoder_config": { + "model_type": "modernbert", "hidden_activation": "gelu", + "attention_bias": false, "mlp_bias": false, "norm_bias": false, + "hidden_size": 768, "intermediate_size": 1152, + "num_attention_heads": 12, "num_hidden_layers": 22, + "vocab_size": 256008, "norm_eps": 1e-5, "local_attention": 128, + "global_attn_every_n_layers": 3, + "global_rope_theta": 160000, "local_rope_theta": 160000, + "max_position_embeddings": 32768 + } + }); + match name { + "root class" => config["model_type"] = json!(mutate), + _ => config["encoder_config"]["model_type"] = json!(mutate), + } + fs::write(dir.path().join("config.json"), config.to_string()).unwrap(); + assert!(read_structure(dir.path()).is_err(), "accepted {name}"); + } + let dir = tempdir().unwrap(); + fs::write( + dir.path().join("config.json"), + json!({"model_type": "vela2-unified"}).to_string(), + ) + .unwrap(); + assert!(read_structure(dir.path()).is_err()); +} + +#[test] +fn profile_drift_is_rejected_with_field_diffs() { + let (_keep, path) = bundle_dir(); + let (mut structure, _) = read_structure(&path).unwrap(); + structure.vocab_size = 50368; + let message = profile().admit(&structure).unwrap_err().to_string(); + assert!(message.contains("vocab_size 50368"), "{message}"); + + let (_keep, path) = bundle_dir(); + let (mut structure, _) = read_structure(&path).unwrap(); + structure.schedule[3] = omni_modernbert::AttentionKind::Sliding; + assert!(profile().admit(&structure).is_err()); + + let (_keep, path) = bundle_dir(); + let (mut structure, _) = read_structure(&path).unwrap(); + structure.rope_full.scaling = None; + assert!(profile().admit(&structure).is_err()); +} + +#[test] +fn the_encoder_plan_runs_the_profile_schedule_without_a_type_add() { + let (_keep, path) = bundle_dir(); + let (structure, _) = read_structure(&path).unwrap(); + let steps = plan::encoder_plan(&structure, &EncoderPlanOptions::default()); + let launches = steps + .iter() + .filter(|s| matches!(s, PlanStep::Launch { .. })) + .count(); + assert_eq!(launches, 1 + 22 * 8); + assert!(!steps.iter().any(|s| s.render().contains("type_emb"))); + let limits = &profile().limits; + limits.check_request(1, 8192).unwrap(); + limits.check_request(4, 512).unwrap(); + for (batch, len) in [(8, 512), (1, 8208), (1, 8190)] { + assert!( + limits.check_request(batch, len).is_err(), + "accepted batch {batch} len {len}" + ); + } +} diff --git a/tests/vela/native/rope.rs b/tests/vela/native/rope.rs new file mode 100644 index 00000000..147f9286 --- /dev/null +++ b/tests/vela/native/rope.rs @@ -0,0 +1,132 @@ +use omni_vela_native::{cos_sin_tables, profile, yarn}; +use serde_json::Value; + +/// Reference tables dumped from the pinned runtime itself: +/// transformers 5.19.0's ModernBertRotaryEmbedding over a +/// ModernBertConfig built from the pinned encoder_config, CPU float32. +fn reference_value() -> Value { + serde_json::from_str(include_str!("data/yarn-reference.json")).unwrap() +} + +fn rope() -> omni_modernbert::structure::RopeSpec { + profile().structure.rope_full +} + +#[test] +fn inverse_frequencies_match_the_reference_runtime() { + let tables = yarn(&rope(), 32).unwrap(); + let expected: Vec = reference_value()["inv_freq"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap()) + .collect(); + assert_eq!(tables.inv_freq.len(), expected.len()); + let mut worst = 0.0f64; + for (actual, wanted) in tables.inv_freq.iter().zip(&expected) { + let diff = (*actual as f64 - wanted).abs(); + worst = worst.max(diff); + } + assert!(worst < 1e-7, "max inverse-frequency delta {worst}"); +} + +#[test] +fn the_attention_scaling_is_the_reference_mscale() { + let tables = yarn(&rope(), 32).unwrap(); + let expected = reference_value()["attention_scaling"].as_f64().unwrap(); + assert!((tables.attention_scaling - expected).abs() < 1e-12); + assert_eq!(tables.attention_scaling, 1.0 + 0.1 * 4.0f64.ln()); +} + +#[test] +fn cos_sin_tables_reproduce_the_reference_rows() { + let structure = &profile().structure; + let (cos, sin, tables) = cos_sin_tables(structure, 32768).unwrap(); + let half = tables.inv_freq.len(); + assert_eq!(cos.len(), 32768 * half); + let reference = reference_value(); + let positions: Vec = reference_value()["positions"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as usize) + .collect(); + let mut worst_served = 0.0f64; + let mut worst_full = 0.0f64; + for (row, pos) in positions.iter().enumerate() { + for j in 0..half { + let cos_ref = reference["cos_half"][row][j].as_f64().unwrap(); + let sin_ref = reference["sin_half"][row][j].as_f64().unwrap(); + let cos_actual = cos[pos * half + j] as f64; + let sin_actual = sin[pos * half + j] as f64; + let delta = (cos_actual - cos_ref) + .abs() + .max((sin_actual - sin_ref).abs()); + if *pos <= 8192 { + worst_served = worst_served.max(delta); + } + worst_full = worst_full.max(delta); + } + } + // The reference computes f32 powers and f32 angle products; this + // generator runs f64 and narrows once. The residual is f32-ulp + // frequency noise amplified by the position, so the served window + // (at most 8 192 tokens) holds a tight gate and the extension + // range a documented looser one. + assert!(worst_served < 1e-4, "served-window delta {worst_served}"); + assert!(worst_full < 5e-4, "full-range delta {worst_full}"); +} + +#[test] +fn the_ramp_upper_bound_clamps_to_the_head_dimension_like_the_reference() { + // theta=1000, original_max=8192, beta_slow=1 puts the unclamped + // correction range above the ramp array (high = 34 > 31): the + // reference clamps to head_dim - 1 = 63 and lets the mask saturate + // beyond the array, so the in-band blend denominator stays + // 34 - low. Clamping to rotary_half - 1 would rescale every + // in-band blend. + let spec = omni_modernbert::structure::RopeSpec { + theta: 1000.0, + scaling: Some(omni_modernbert::structure::YarnScaling { + factor: 4.0, + original_max_positions: 8192, + beta_fast: 32.0, + beta_slow: 1.0, + }), + }; + let tables = yarn(&spec, 32).unwrap(); + let correction = |rotations: f64| { + (2.0 * 32.0 * (8192.0f64 / (rotations * std::f64::consts::TAU)).ln()) / (2.0 * 1000f64.ln()) + }; + let low = correction(32.0).floor().max(0.0); + let high = correction(1.0).ceil().min(2.0 * 32.0 - 1.0); + assert!(high > 31.0, "the case must exceed the ramp array ({high})"); + for (j, actual) in tables.inv_freq.iter().enumerate() { + let j = j as f64; + let pos_freq = 1000.0f64.powf(j / 32.0); + let ramp = ((j - low) / (high - low)).clamp(0.0, 1.0); + let expected = 1.0 / (4.0 * pos_freq) * ramp + 1.0 / pos_freq * (1.0 - ramp); + assert!( + (*actual as f64 - expected).abs() < 1e-6, + "inv_freq[{j}] = {actual}, expected {expected}" + ); + } +} + +#[test] +fn table_lengths_follow_the_structure_bounds() { + let structure = &profile().structure; + assert!(cos_sin_tables(structure, 0).is_err()); + assert!(cos_sin_tables(structure, 32769).is_err()); + let (cos, sin, _) = cos_sin_tables(structure, 8192).unwrap(); + assert_eq!(cos.len(), 8192 * 32); + assert_eq!(sin.len(), cos.len()); + // Inside the ramp band the frequency is a blend of extrapolation + // and interpolation, so it sits strictly between plain theta + // rope and the fully interpolated (factor-divided) variant. + let tables = yarn(&rope(), 32).unwrap(); + let plain = 160000.0f64.powf(-(15.0 / 32.0)); + let blended = tables.inv_freq[15] as f64; + assert!(blended < plain, "{blended} vs {plain}"); + assert!(blended > plain / 4.0, "{blended} vs {}", plain / 4.0); +} diff --git a/tests/vela/native/weights.rs b/tests/vela/native/weights.rs new file mode 100644 index 00000000..feb0dd74 --- /dev/null +++ b/tests/vela/native/weights.rs @@ -0,0 +1,148 @@ +use omni_vela_native::{Weights, head_tensors, read_structure, storage_dtype}; +use serde_json::{Value, json}; +use std::fs; +use tempfile::{TempDir, tempdir}; + +/// The recorded census of the pinned checkpoint's safetensors header +/// (vllm-sr/Vela-2.0-0.3B@d6f03aa9, model.safetensors): every tensor +/// name with dtype and shape, fetched from the first bytes of the file. +fn census() -> Value { + serde_json::from_str(include_str!("data/pinned-checkpoint-census.json")).unwrap() +} + +/// Builds a synthetic checkpoint holding exactly the given tensors. +fn checkpoint(tensors: &Value) -> (TempDir, std::path::PathBuf) { + let dir = tempdir().unwrap(); + let path = dir.path().join("model.safetensors"); + let mut entries: Vec<(String, Vec, &str, Vec)> = Vec::new(); + for (name, spec) in tensors.as_object().unwrap() { + let shape: Vec = spec["shape"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as usize) + .collect(); + let dtype = spec["dtype"].as_str().unwrap(); + let width = if dtype == "F32" { 4 } else { 2 }; + let bytes: Vec = vec![0; shape.iter().product::() * width]; + entries.push((name.clone(), shape, dtype, bytes)); + } + entries.sort_by(|a, b| a.0.cmp(&b.0)); + let mut offsets = 0usize; + let mut header = serde_json::Map::new(); + let mut blob: Vec = Vec::new(); + for (name, shape, dtype, bytes) in &entries { + let end = offsets + bytes.len(); + header.insert( + name.clone(), + json!({"dtype": dtype, "shape": shape, "data_offsets": [offsets, end]}), + ); + offsets = end; + blob.extend_from_slice(bytes); + } + let mut header_text = Value::Object(header).to_string(); + header_text.extend(std::iter::repeat_n(' ', (8 - header_text.len() % 8) % 8)); + let mut file = (header_text.len() as u64).to_le_bytes().to_vec(); + file.extend_from_slice(header_text.as_bytes()); + file.extend_from_slice(&blob); + fs::write(&path, file).unwrap(); + (dir, path) +} + +mod common; + +fn structure() -> omni_modernbert::structure::ModernBertStructure { + let (_keep, path) = common::bundle_dir(); + let (structure, _) = read_structure(&path).unwrap(); + structure +} + +#[test] +fn the_geometry_inventory_matches_the_pinned_census_exactly() { + let expected = omni_modernbert::geometry::encoder_tensors(&structure()); + let heads = head_tensors(); + assert_eq!(expected.len(), 134); + assert_eq!(heads.len(), 18); + for (name, shape) in expected.iter().chain(heads.iter()) { + let entry = &census()[name]; + assert!(entry.is_object(), "{name} missing from the census"); + let census_shape: Vec = entry["shape"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as usize) + .collect(); + assert_eq!(&census_shape, shape, "{name}"); + assert_eq!(entry["dtype"], "F32", "{name}"); + } + // And the census holds nothing beyond these two inventories. + assert_eq!(census().as_object().unwrap().len(), 134 + 18); +} + +#[test] +fn a_checkpoint_matching_the_profile_validates_completely() { + let (_dir, path) = checkpoint(&census()); + let weights = Weights::open(&path).unwrap(); + assert_eq!(weights.validate(&structure()).unwrap(), 152); + // Storage classes mirror the Laya baseline: embedding f16, norms, + // calibration scalars and head biases f32, GEMM-class bf16. + assert_eq!( + storage_dtype("encoder.embeddings.tok_embeddings.weight"), + "f16" + ); + assert_eq!(storage_dtype("encoder.layers.0.attn_norm.weight"), "f32"); + assert_eq!(storage_dtype("encoder.layers.0.attn.Wqkv.weight"), "bf16"); + assert_eq!(storage_dtype("norm.weight"), "f32"); + assert_eq!(storage_dtype("log_tau"), "f32"); + assert_eq!(storage_dtype("log_tau_span"), "f32"); + assert_eq!(storage_dtype("cls_mlp.3.bias"), "f32"); + assert_eq!(storage_dtype("w_e.bias"), "f32"); + assert_eq!(storage_dtype("cls_mlp.0.weight"), "bf16"); + assert_eq!(storage_dtype("w_o.weight"), "bf16"); + let scalar = weights.f32("log_tau", &[]).unwrap(); + assert_eq!(scalar.len(), 1); +} + +#[test] +fn inventory_drift_fails_naming_both_sides() { + for (name, mutate) in [ + ("missing", { + let mut c = census(); + c.as_object_mut() + .unwrap() + .remove("encoder.layers.5.attn.Wqkv.weight"); + c + }), + ("extra", { + let mut c = census(); + c.as_object_mut().unwrap().insert( + "encoder.layers.22.attn.Wqkv.weight".into(), + json!({"dtype": "F32", "shape": [2304, 768]}), + ); + c + }), + ("shape", { + let mut c = census(); + c["encoder.layers.0.mlp.Wi.weight"]["shape"] = json!([2303, 768]); + c + }), + ("dtype", { + let mut c = census(); + c["encoder.layers.0.attn.Wo.weight"]["dtype"] = json!("F16"); + c + }), + ] { + let (_dir, path) = checkpoint(&mutate); + let weights = Weights::open(&path).unwrap(); + let result = weights.validate(&structure()); + assert!(result.is_err(), "accepted {name}"); + let message = result.unwrap_err().to_string(); + match name { + "missing" => assert!(message.contains("missing"), "{message}"), + "extra" => assert!(message.contains("unexpected"), "{message}"), + "shape" => assert!(message.contains("expected"), "{message}"), + "dtype" => assert!(message.contains("f32"), "{message}"), + _ => {} + } + } +} diff --git a/tests/vela/test_engine_manifest.py b/tests/vela/test_engine_manifest.py index 4968bdaa..5fe1b2ee 100644 --- a/tests/vela/test_engine_manifest.py +++ b/tests/vela/test_engine_manifest.py @@ -16,6 +16,7 @@ import py_compile import re import sys +import types from pathlib import Path import pytest @@ -305,6 +306,21 @@ def test_a_runtime_module_from_another_directory_is_refused(tmp_path): sys.modules.pop("modeling_vela2", None) +def test_a_runtime_module_without_a_file_is_refused(tmp_path): + # A namespace package or builtin resolves with __file__ = None; it + # cannot be this package's flat module, so the gate fails closed + # instead of crashing on the None path. + package = tmp_path / "pkg" + package.mkdir() + build_package(package, small_files()) + sys.modules["vela2_inference"] = types.ModuleType("vela2_inference") + try: + with pytest.raises(ValueError, match="namespace"): + verify_no_foreign_runtime(package) + finally: + sys.modules.pop("vela2_inference", None) + + def test_the_package_root_is_forced_to_the_front_of_sys_path(tmp_path): # A PYTHONPATH copy of the root listed later must not leave a foreign # directory ahead of it: verify_no_foreign_runtime cannot catch this, diff --git a/tests/vela/test_oracle_fixtures.py b/tests/vela/test_oracle_fixtures.py new file mode 100644 index 00000000..54163334 --- /dev/null +++ b/tests/vela/test_oracle_fixtures.py @@ -0,0 +1,52 @@ +"""Ties the frozen oracle fixtures to the recipe examples they came from. + +The four oracle fixtures under ``tests/vela/fixtures`` are the frozen +comparison baseline for the later native encoder slices: each one's +``request`` is the recipe's example request, and its recorded +``system_one`` response passes the worker's own response gate. This +keeps the tie mechanical instead of by-convention. No weights, no +torch. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from models.vela.postprocess import finish_response +from models.vela.protocol import MODEL_NAME + +TESTS = Path(__file__).parent +RECIPE = TESTS.parent.parent / "recipe" / "vela" + +CASES = [ + ("oracle-example-request", "example-request.json"), + ("oracle-example-request-cjk", "example-request-cjk.json"), + ("oracle-example-request-set-threshold", "example-request-set-threshold.json"), + ("oracle-example-request-score", "example-request-score.json"), +] + + +def oracle(case: str) -> dict: + return json.loads((TESTS / "fixtures" / f"{case}.json").read_text()) + + +@pytest.mark.parametrize("case,example", CASES) +def test_the_oracle_request_is_the_recipe_example(case: str, example: str): + recorded = oracle(case) + served = json.loads((RECIPE / example).read_text()) + assert recorded["request"] == served + # Same insertion order, so byte-for-byte serialization holds too. + assert list(recorded["request"]) == list(served) + assert recorded["request"]["model"] == MODEL_NAME + + +@pytest.mark.parametrize("case,example", CASES) +def test_the_recorded_response_passes_the_worker_gate(case: str, example: str): + recorded = oracle(case) + response = recorded["system_one"] + assert finish_response(recorded["request"]["questions"], response) is response + # The frozen case label still names the recipe file it belongs to. + assert recorded["case"] == example.removesuffix(".json") diff --git a/tests/vela/test_postprocess.py b/tests/vela/test_postprocess.py index 9f0fdbf8..e67f1642 100644 --- a/tests/vela/test_postprocess.py +++ b/tests/vela/test_postprocess.py @@ -140,8 +140,22 @@ def test_finish_response_matches_span_and_set_key_sets(): lambda r: r.update(usage={"input_tokens": 5}), lambda r: r.update(usage={"input_tokens": -1, "output_tokens": 0}), lambda r: r.update(usage={"input_tokens": 10, "output_tokens": 3}), - # Answer-shape violations. + # Answer-shape violations: foreign fields, missing fields, and + # unknown kinds — exact key sets per answer kind. lambda r: r["answers"]["billing"].update(probabilities={"yes": 0.75}), + lambda r: r["answers"]["billing"].pop("noul"), + lambda r: r["answers"]["area"].pop("probabilities"), + lambda r: r["answers"]["area"].update(extra=1), + lambda r: r["answers"]["severity"].pop("legend"), + lambda r: r["answers"]["severity"].update(confidence2=0.5), + lambda r: r["answers"].update({"billing": {"type": "abstain"}}), + # Payload-shape violations. + lambda r: r["sets"]["issues"].update(threshold=0.3), + lambda r: r["sets"]["issues"].pop("selected"), + lambda r: r["spans"]["pii"].append({"label": "PERSON", "start": 0}), + lambda r: r["spans"]["pii"][0].update(weight=1.0), + lambda r: r["thresholds"].update({"pii": "0.75"}), + lambda r: r["thresholds"].update({"pii": True}), ], ) def test_finish_response_rejects_contract_violations(mutate): diff --git a/tests/vela/test_worker_server.py b/tests/vela/test_worker_server.py index 9221a933..242e2dbf 100644 --- a/tests/vela/test_worker_server.py +++ b/tests/vela/test_worker_server.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import socket import threading from http.client import HTTPConnection @@ -118,6 +119,23 @@ def test_models_lists_the_served_member_upstream_verbatim(server): } +def test_body_short_of_content_length_answers_400(server): + payload = json.dumps(five_kind_body()).encode("utf-8") + conn = HTTPConnection("127.0.0.1", server.server_address[1], timeout=10) + conn.putrequest("POST", "/v1/systemone", skip_accept_encoding=True) + conn.putheader("Content-Type", "application/json") + conn.putheader("Content-Length", str(len(payload) + 64)) + conn.endheaders(payload) + # Half-close: the client will never send the promised remaining bytes. + conn.sock.shutdown(socket.SHUT_WR) + response = conn.getresponse() + raw = response.read() + status = response.status + conn.close() + assert status == 400 + assert json.loads(raw) == {"detail": "incomplete body"} + + def test_valid_five_kind_decision_returns_200(server): status, ctype, raw = post(server, json.dumps(five_kind_body()).encode("utf-8")) assert status == 200 and ctype == "application/json"