5ae36f2c08
Tentative JACCL examples
2026-04-22 01:29:40 -07:00
c284e0a231
Enable swap for all CI building CUDA ( #3437 )
2026-04-22 13:13:24 +09:00
b9b1bfb9a5
Generate qmm implementaions with cmake ( #3424 )
2026-04-22 13:11:55 +09:00
68cf2fddd8
Fix mx.prod vjp for complex types ( #3433 )
2026-04-21 17:35:20 -07:00
c594e6ec38
Fix use after move in Compiled primitive ( #3427 )
2026-04-21 15:22:45 -07:00
7d40a4fd5a
Throw meaningful error when Metal device is not found ( #3428 )
2026-04-21 15:21:08 -07:00
5f519ef6f9
Fix bytes_per_key truncation in random kernels (Metal + CUDA) ( #3432 )
2026-04-21 15:15:11 -07:00
705c828feb
Fix synchronize for ThreadLocalStream ( #3429 )
2026-04-20 11:29:49 -07:00
b4ddf9b374
Fix flaky TestVmap.test_vmap_masked_scatter ( #3421 )
2026-04-20 17:19:20 +09:00
1f5a413a27
Make Scheduler::enqueue thread safe ( #3423 )
2026-04-20 14:30:05 +09:00
a6222f53d5
Speed up NAX split-K by better tuning and routing and fix NAX addmm ( #3422 )
...
I 'll merge now and comment with more benchmarks later since this also fixes two bugs so worst case we 'll do another tuning, it isn't like we won't need the functionality of this PR.
2026-04-19 18:05:39 -07:00
fa4320d5fa
[CUDA] Handle residue k in qmm_naive ( #3379 )
2026-04-18 13:30:07 +09:00
859f22fbb0
[CUDA] GatherQMM matrix-matrix sm80/naive path ( #3417 )
...
Co-authored-by: Cheng <git@zcbenz.com >
2026-04-18 10:59:47 +09:00
d142de6a20
[CUDA] gather_mm ( #3414 )
2026-04-17 16:53:44 +09:00
940ba473fe
Segmented mm nax kernel ( #3419 )
2026-04-16 17:26:29 -07:00
8e649be4d0
Fix jaccl init bug ( #3418 )
2026-04-16 01:23:35 -07:00
dec6b4d10f
ThreadLocalStream in C++ ( #3405 )
2026-04-15 15:46:11 -07:00
fd8e849e26
Document sort stability and NaN handling ( #3400 )
2026-04-15 14:32:42 -07:00
50ae31241a
Validate safetensors data offsets against file boundaries ( #3410 )
...
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com >
2026-04-15 14:30:55 -07:00
6cef1e995e
Validate safetensors data offsets ( #3364 )
2026-04-15 00:52:42 -07:00
57bcced8cb
Fixes for CUDA CI ( #3413 )
2026-04-14 23:52:52 -07:00
4400504ad5
Jaccl refactor ( #3412 )
2026-04-14 23:52:21 -07:00
1fa764fbec
Update nanobind version to v2.12.0 ( #3396 )
2026-04-14 17:21:00 -07:00
435f0b6cdb
Add clear_streams API for cleanup before exit ( #3395 )
2026-04-14 18:41:32 +09:00
520cea2bec
Avoid joining threads on exit ( #3388 )
2026-04-11 09:22:34 +09:00
a33b791615
Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K ( #3361 )
...
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com >
2026-04-10 00:01:47 -07:00
d6d9b24801
Conjugate VJP and JVP support ( #3386 )
2026-04-09 15:04:46 -07:00
8332e228e4
Fix test "test get streams" missing initialization ( #3376 )
2026-04-09 08:29:04 +09:00
4403165843
[CUDA] Thread safety ( #3367 )
2026-04-09 08:18:00 +09:00
a8776b7bbd
Fix: Correct cross-attention query routing in Post-LN TransformerDecoderLayer ( #3382 )
2026-04-07 09:16:12 -07:00
b98831ad0e
fix: fail build when Metal compiler header resolution fails ( #3332 )
2026-04-06 12:49:25 -07:00
d025111b1d
[CUDA] Add GatherQMM for quantized gather matmul ( #3321 )
2026-04-06 12:48:18 -07:00
9239808225
Fix CMake finding wrong Python during pip install ( #3375 )
2026-04-06 12:32:16 -07:00
6a9a121d09
Add a convenience for making local streams in python ( #3355 )
2026-04-02 18:43:02 -07:00
befe42d303
Add printoptions ( #3333 )
2026-04-01 22:24:48 -07:00
80a1c206f9
Use metal as the front-end for the metal linker ( #3354 )
2026-04-01 16:52:07 -07:00
b0748ad8de
Fix regression in array creation ( #3353 )
2026-04-01 11:30:36 -07:00
2ffafe07f4
[CUDA] 3/5/6-bit quants for qmm_naive ( #3352 )
2026-04-01 20:13:01 +09:00
5e2c44259f
Make CommandEncoder thread local ( #3348 )
2026-04-01 18:42:49 +09:00
1c9ee2f655
[CUDA] Fallback QMM ( #3315 )
2026-04-01 12:41:26 +09:00
7cd73c4202
[Metal] Support sorting complex numbers ( #3314 )
2026-04-01 12:40:50 +09:00
2105df91da
Add fftfreq, rfftfreq and scalar axes for fftshift/ifftshift ( #3298 )
2026-03-31 18:29:16 -07:00
1944cf67a2
Add vmap for BroadcastAxes ( #3344 )
2026-03-31 17:08:56 -07:00
939e425c7a
Decouple CommandEncoder from Device ( #3316 )
2026-04-01 08:51:17 +09:00
8439b1f501
Fix use after move ( #3343 )
2026-03-31 10:37:40 -07:00
117b4f1806
Bump actions/deploy-pages from 4 to 5 ( #3334 )
...
Signed-off-by: dependabot[bot] <support@github.com >
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-03-31 09:04:19 +09:00
66f58032dc
Remove no longer needed const_cast ( #3325 )
2026-03-31 08:10:49 +09:00
8a6d28713c
Fix np bfloat16 misinterpreted as complex ( #3146 )
...
Co-authored-by: Cheng <git@zcbenz.com >
2026-03-31 08:04:55 +09:00
0ff1115a46
[CUDA] Implement BlockMaskedMM ( #3299 )
...
Co-authored-by: Cheng <git@zcbenz.com >
2026-03-27 06:57:26 +09:00
df7f7db943
Make each thread have its own default stream ( #3281 )
2026-03-25 15:48:49 +09:00
Angelos Katharopoulos