ejkernel 0.0.80__tar.gz → 0.0.80.1__tar.gz

{ejkernel-0.0.80 → ejkernel-0.0.80.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ejkernel
-Version: 0.0.80
+Version: 0.0.80.1
 Summary: Accelerate, Optimize performance with streamlined training and serving options with JAX.
 Keywords: Deep Learning,Machine Learning,JAX,CUDA,XLA,Triton,Pallas
 Author: Erfan Zare Chavoshi
@@ -25,9 +25,9 @@ Requires-Dist: jax[cuda13]~=0.10.0 ; extra == 'cuda'
 Requires-Dist: jax-cuda13-plugin[with-cuda]~=0.10.0 ; extra == 'cuda'
 Requires-Dist: jax-cuda13-pjrt~=0.10.0 ; extra == 'cuda'
 Requires-Dist: triton==3.6.0 ; extra == 'cuda'
-Requires-Dist: nvidia-cutlass-dsl[cu13]==4.4.0.dev1 ; extra == 'cuda'
-Requires-Dist: nvidia-cutlass-dsl-libs-base==4.4.0.dev1 ; extra == 'cuda'
-Requires-Dist: nvidia-cutlass-dsl-libs-cu13==4.4.0.dev1 ; extra == 'cuda'
+Requires-Dist: nvidia-cutlass-dsl[cu13]==4.5.2 ; extra == 'cuda'
+Requires-Dist: nvidia-cutlass-dsl-libs-base==4.5.2 ; extra == 'cuda'
+Requires-Dist: nvidia-cutlass-dsl-libs-cu13==4.5.2 ; extra == 'cuda'
 Requires-Dist: jax-tvm-ffi==0.1.2 ; extra == 'cuda'
 Requires-Dist: apache-tvm-ffi==0.1.8.post2 ; extra == 'cuda'
 Requires-Dist: eformer ; extra == 'dev'

{ejkernel-0.0.80 → ejkernel-0.0.80.1}/ejkernel/__init__.py

@@ -65,7 +65,7 @@ import os as _os
 _os.environ.setdefault("TF_GPU_ALLOCATOR", "cuda_malloc_async")
 _os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1")
 
-__version__ = "0.0.80"
+__version__ = "0.0.80.1"
 
 from . import errors, kernels, modules, types, utils, xla_utils
 from .errors import EjkernelRuntimeError

{ejkernel-0.0.80 → ejkernel-0.0.80.1}/ejkernel/kernels/_xla/fused_kl_divergence/_xla_impl_fwd.py

@@ -155,8 +155,10 @@ def fused_kl_divergence(
         temperature: Softmax temperature ``T``; loss scaled by ``T²``.
         beta: JSD interpolation factor (only used when
             ``direction="jsd"``).
-        vocab_parallel_axis: Mesh axis name for TP (forward + ``T=1`` only;
-            other combos fall back to the non-fused path).
+        vocab_parallel_axis: Mesh axis name for TP. Forward direction only
+            (reverse/JSD fall back to the non-fused path); any temperature is
+            supported via a ``1/T`` pre-scale + ``T**2`` post-scale around the
+            vocab-parallel core.
     """
     if reduction not in ("none", "sum", "mean"):
         raise ValueError(f"Invalid reduction '{reduction}'; expected one of none/sum/mean.")
@@ -181,12 +183,24 @@ def fused_kl_divergence(
         wts = weights.astype(jnp.float32)
 
     if vocab_parallel_axis is not None:
-        if direction != "forward" or temperature != 1.0:
+        if direction != "forward":
             raise NotImplementedError(
-                "vocab-parallel mode currently only supports direction='forward' with "
-                "temperature=1.0 in the XLA fallback."
+                "vocab-parallel mode currently only supports direction='forward' in the XLA fallback."
             )
-        per_row = _fused_kl_core_tp(student_logits, teacher_logits, wts, vocab_parallel_axis)
+        # Temperature is a linear 1/T pre-scale on the (still per-shard ``V_local``) logits plus a
+        # T^2 post-scale on the per-row loss. Both compose through ``_fused_kl_core_tp``'s custom_vjp
+        # via outer autodiff, yielding the exact T^2*(1/T)=T-scaled distillation gradient WITHOUT
+        # ever materializing the full vocabulary (the whole point of the vocab-parallel path).
+        if temperature != 1.0:
+            inv_T = 1.0 / float(temperature)
+            per_row = _fused_kl_core_tp(
+                student_logits * inv_T,
+                teacher_logits * inv_T,
+                wts,
+                vocab_parallel_axis,
+            ) * (float(temperature) ** 2)
+        else:
+            per_row = _fused_kl_core_tp(student_logits, teacher_logits, wts, vocab_parallel_axis)
     elif direction == "forward" and temperature == 1.0:
         per_row = _fused_kl_core(student_logits, teacher_logits, wts)
     else:

{ejkernel-0.0.80 → ejkernel-0.0.80.1}/ejkernel/modules/operations/fused_kl_divergence.py

@@ -516,7 +516,8 @@ def fused_kl_divergence(
             temperatures.
         beta: JSD interpolation factor in ``(0, 1)``; ignored unless
             ``direction="jsd"``.
-        vocab_parallel_axis: TP mesh axis (forward + ``T=1`` only).
+        vocab_parallel_axis: TP mesh axis (forward direction only; any
+            temperature supported).
         return_teacher_entropy: When ``True``, also return the teacher
             entropy ``H(p_t)`` on the output (same reduction + ``T²`` scaling
             as ``loss``, detached). Lets distillation callers recover the

{ejkernel-0.0.80 → ejkernel-0.0.80.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ejkernel"
-version = "0.0.80"
+version = "0.0.80.1"
 authors = [{ name = "Erfan Zare Chavoshi", email = "Erfanzare810@gmail.com" }]
 description = "Accelerate, Optimize performance with streamlined training and serving options with JAX."
 readme = "README.md"
@@ -66,9 +66,9 @@ cuda = [
     "jax-cuda13-plugin[with-cuda]~=0.10.0",
     "jax-cuda13-pjrt~=0.10.0",
     "triton==3.6.0",
-    "nvidia-cutlass-dsl[cu13]==4.4.0.dev1",
-    "nvidia-cutlass-dsl-libs-base==4.4.0.dev1",
-    "nvidia-cutlass-dsl-libs-cu13==4.4.0.dev1",
+    "nvidia-cutlass-dsl[cu13]==4.5.2",
+    "nvidia-cutlass-dsl-libs-base==4.5.2",
+    "nvidia-cutlass-dsl-libs-cu13==4.5.2",
     "jax-tvm-ffi==0.1.2",
     "apache-tvm-ffi==0.1.8.post2",
 ]