Repository navigation
jax-tt-plugin 0.1.2
This release improves LLM serving performance, multi-device execution, numerical coverage, and JAX compatibility on Tenstorrent Blackhole devices.
Highlights
- Added memory-mapped checkpoint loading with borrowed strided host buffers, on-device BFP4 conversion, and streamlined MoE layout handling.
- Improved Qwen3 and Qwen3.8 Flash Next serving, including faster overlap scheduling, decode attention, rotary embeddings, QKV handling, and FP32 matmul accumulation.
- Added fused MoE execution and optimized router softmax and top-k operations.
- Improved multi-device support with submeshes, multiple meshes over shared devices, and more reliable all-gather and all-reduce synchronization.
- Added a unified indexed-rows kernel for gather, scatter, and dynamic-update operations.
Numerical and operation coverage
- Expanded Float16 support across arithmetic, conversion, and data-movement operations.
- Added vector-unit FP32 matmul support for
Precision.HIGHEST. - Improved exact integer reductions, sorting,
mulhi,arange, gather indices, and scatter edge cases. - Added TTNN-backed cumulative min, max, and log-sum-exp operations.
- Added TTNN implementations for bitwise reductions and
reduce_precision. - Improved floating-point sorting, including NaN ordering and signed-zero handling.
- Improved integer and bfloat16 sorting and top-k behavior.
Performance and reliability
- Faster tilization for short and wide row-major tensors.
- Reduced host overhead for decode-sized workloads and batched dot operations.
- Improved direct all-reduce behavior for packed decode rows.
- Fixed program-cache aliasing and several collective synchronization issues.
- Improved compilation diagnostics and reduced unnecessary MLIR build dependencies.
- Updated the tt-xla, tt-mlir, tt-metal, and UMD dependencies.