annote triton constexpr #56

foundation-model-stack · Jul 31, 2024 · 1ef30a9 · 1ef30a9
1 parent 80edec8
commit 1ef30a9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/...ins/fused-ops-and-kernels/src/fms_acceleration_foak/kernels/unsloth/cross_entropy_loss.py b/...ins/fused-ops-and-kernels/src/fms_acceleration_foak/kernels/unsloth/cross_entropy_loss.py
@@ -186,7 +186,7 @@ def _cross_entropy_backward(
 pass
 
 
-MAX_FUSED_SIZE = 65536 # 2**16
+MAX_FUSED_SIZE: tl.constexpr = 65536 # 2**16
 
 class Fast_CrossEntropyLoss(torch.autograd.Function):
     @staticmethod

diff --git a/plugins/fused-ops-and-kernels/src/fms_acceleration_foak/kernels/unsloth/rope_embedding.py b/plugins/fused-ops-and-kernels/src/fms_acceleration_foak/kernels/unsloth/rope_embedding.py
@@ -17,7 +17,7 @@
 import torch
 from .utils import calculate_settings
 
-ROPE_GROUP_SIZE = 4
+ROPE_GROUP_SIZE: tl.constexpr = 4
 
 @triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],})
 @triton.jit