Commit Graph

120 Commits

Author SHA1 Message Date
145e60b4dd workspace: Update LLVM, XLA, StableHLO, and PJRT plugins to latest versions. 2023-12-13 10:10:32 +00:00
6a4a7fb9a1 zml/module.zig: Remove unnecessary optional unwrapping. 2023-12-05 12:27:08 +00:00
37725cdaa6 Update PJRT, runtime, and ZML modules to use per‑target output folders and expose profiler.dumpDataAsJson for JSON profiling output. 2023-12-04 10:38:10 +00:00
22a846de72 Update llama example to use per‑target output folders and call profiler.dumpDataAsJson for testing the new compilation layout. 2023-12-01 16:05:59 +00:00
46fbbf43a2 Update tutorial documentation in write_first_model.md with quick fixes. 2023-11-30 12:14:33 +00:00
737f7cbdee Add example build runner scripts and config for Zig code completion. 2023-11-21 14:55:34 +00:00
ec37c8f731 Update Bazel build files and helper scripts to integrate the custom build runner for ZLS code completion. 2023-11-20 15:29:01 +00:00
6e4fef8844 zml: Introduce arena allocator in CompilationContext. Expose arena allocator to replace existing allocator, enabling safe allocation for ops without misusing std.BoundedArray. Includes breaking changes to chunkAllowTrailing and split. Upgrade axis_ types to anytype for tag handling and add TODOs for upcoming Tensor API. 2023-11-16 15:11:23 +00:00
57bf667c90 Add struct‑based client creation flags to the Zig PJRT API and update context.autoPlatform to accept a flag struct. 2023-11-13 12:45:17 +00:00
cb6fcbbb1a Update docs and Zig examples to demonstrate the new client creation flags API. 2023-11-09 12:31:11 +00:00
9f4194ad97 Fix test layer. Add tests to detect silent breakage of testLayer and regression in mapAlloc with zero-size struct fields. Add Python venv directory to .gitignore. 2023-11-06 11:25:57 +00:00
237a877a29 zml: Add support for Llama 3.2 text-only models. Implement transpose over embed_tokens as a replacement for missing lm_head and make lm_head optional for compatibility. Add repositories and executions to Bazel and update README. 2023-11-01 10:16:48 +00:00
1c9749c25e docs: move image in concepts.md 2023-10-31 10:21:14 +00:00
eb20548241 update instructions
following, `prepare` doesn't alloc anymore, `ExeWithWeights` is
`ModuleExe`
2023-10-26 13:56:56 +00:00
27c8309424 async: add intrusive queue
all code contributed by @steeve

* add intrusive queue
* change the constructor of Channel with default AsyncThread executor

---------

Co-authored-by: Steeve Morin <steeve@zml.ai>
2023-10-24 14:36:22 +00:00
98b512c495 Implement func.call emission and function caching across MLIR dialects and ZML module/ops, propagating tags and donations. 2023-10-19 17:01:55 +00:00
37de7b9613 Add Llama example showcasing the new func.call emission and function caching behavior. 2023-10-17 11:00:37 +00:00
7d36913b31 Refactor ZML API: move compile, compileFn and related types to exe.zig, update BaseExe allocation and inline caching in compileInternal, and clean up supporting modules (func.zig, meta.zig, signature.zig, cuda.zig, testing.zig, zml.zig). 2023-10-13 16:08:08 +00:00
35395c13f8 Update example programs (benchmark, llama, mnist, simple_layer) to use the new Exe API and reflect BaseExe allocation changes. 2023-10-10 11:12:34 +00:00
3bc6ad98be Update module.zig to donate all buffers except the token_index buffer for the Llama+Neuron example. 2023-10-06 10:10:56 +00:00
474f76cd75 Enable buffer donation in the Llama example, donating all buffers except the token_index buffer. 2023-10-03 16:32:40 +00:00
5122ca0203 Refactor rope implementation to compute only required offsets, eliminating full cos/sin matrix generation in module, nn, and tensor code. 2023-09-27 11:45:33 +00:00
06865f5876 Update Llama example to use the new direct rope IR implementation. 2023-09-25 10:22:05 +00:00
b5c4fb7c58 zml: fix float8 <-> float32 conversions, support for Tensor.constant(.{}, .{ .f8 = 1.0})
Mostly:
* fix float8 <-> float32 conversions
* support for `Tensor.constant(.{}, .{ .f8 = 1.0})`

Misc:
* fix small inconsistencies between different versions of sdpa
* better error message for broadcast
* bazelrc: --config=debug
2023-09-21 11:15:50 +00:00
455bb3877f runtimes/cuda: obtain NCCL from the pip package, matching XLA behavior. 2023-09-20 17:41:44 +00:00
0d5389ceda Update CUDA runtime sandboxing and dynamic symbol renaming, switch to pre‑built jax‑cuda‑pjrt plugin, and bump CUDA to 12.6.2 and cuDNN to 9.5.1. 2023-09-14 13:28:25 +00:00
4abdd32f0d Update llama example BUILD to use jax-cuda-pjrt plugin and bump CUDA (12.6.2) / CuDNN (9.5.1) versions. 2023-09-12 15:40:21 +00:00
c8c99d7d5a zml/pjrtx: prefer the built‑in stablehlo version when a plugin reports a newer version, ensuring artifact serialization uses the correct stablehlo version. 2023-09-07 17:06:19 +00:00
9505992e00 workspace: log diagnostic message before returning NotFound to aid debugging. 2023-09-04 13:34:37 +00:00
937cdec324 examples/loader: add missing stdx dependency. 2023-08-30 13:03:59 +00:00
aa7fae449e zml/pjrtx: execute bufferFromHostBuffer on the thread pool to avoid blocking and improve weight loading performance. 2023-08-29 10:28:51 +00:00
c081cb9ad6 zml/platform: increase maximum device limit to support up to 32 devices per platform. 2023-08-24 12:23:07 +00:00
af0630616c Update docs (deploy_on_server, dockerize_models, getting_started) and example Bazel files to include AWS Neuron/Trainium/Inferentia deployment guidance. 2023-08-21 09:15:48 +00:00
7d24329d0a Add Bazel build rules and runtime implementation for AWS Neuron/Trainium/Inferentia support. 2023-08-18 17:11:27 +00:00
0709b1b32f zml: reduce memory usage of sdpaMemEfficient by using zml.ops.while instead of zml.ops.for, avoiding concatenation of intermediate results. 2023-08-14 14:24:11 +00:00
022baf782b Update examples/MODULE.bazel to reference the bumped LLVM, XLA, StableHLO, and PJRT plugin versions. 2023-08-11 16:57:15 +00:00
01eff33fa0 Update workspace dependencies to newer LLVM, XLA, StableHLO, and PJRT versions and expose new pjrt plugin attribute APIs and stablehlo version APIs in build and runtime configurations. 2023-08-07 12:28:36 +00:00
726a2d0691 Update docs and examples to showcase the new async runtime with coroutines and cross‑thread signaling. 2023-08-03 11:35:24 +00:00
bcde3962ce Rework async runtime with coroutine support, rename async API (async_→asyncc, await_→awaitt), improve type inference, bump libxev (default epoll) and update related stdx and zml modules. 2023-08-01 11:35:04 +00:00
b53462b515 Fix crash in for_ by ensuring values are pushed to their block before opening a new block, adding asserts for block state, and guaranteeing first_step is used. Adjust padding syntax to improve usability. 2023-07-25 14:25:47 +00:00
0fa258cd88 Update examples to reflect recent async module changes, renaming asyncGeneric to asyncc. 2023-07-24 09:34:35 +00:00
f675a203c2 zml.ops.makeBlock now returns the inner tensor to propagate tags. The function returns both the created mlir.Block and tensors from the supplied function, allowing shape and tag propagation without exposing mlir.Values. Updated tests to run on non‑CPU platforms. 2023-07-21 09:01:01 +00:00
be8aa4fa8e Fix several compileError calls introduced by recent changes; ensure Zig compiler catches errors at comptime. 2023-07-17 09:10:27 +00:00
0f9a92f27d module-cache: raise max_pjrt_executable_size limit to 400 MB to accommodate large PJRT executables. 2023-07-14 17:58:22 +00:00
88c7a74ccf third_party/modules/zig-protobuf: revert indentation changes to maintain compatibility with older branches. 2023-07-13 11:44:53 +00:00
4681ce2f24 PJRT: add conversion of profiling protobuf output to JSON format. 2023-07-05 13:34:05 +00:00
f7bac1af10 Update example programs (llama and loader) with hotfixes for issue. 2023-07-04 13:40:05 +00:00
63aca9f9c2 Hotfixes for build rule, math utilities, module system, and NN implementation (fixes,) 2023-06-29 10:26:54 +00:00
7985716562 Add new Zig example programs (benchmark, llama, loader, mnist, simple_layer) and include a test for the llama example. 2023-06-27 14:23:22 +00:00
9b7eea8ac2 Add stdx utilities and rework async signature inference; tidy executable logging. 2023-06-21 14:45:14 +00:00