fix: update global variable syntax since its not support from triton 3.2.0 (#148)

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

plugins/attention-and-distributed-packing/pyproject.toml

@@ -23,7 +23,7 @@ classifiers=[
     "Programming Language :: Python :: 3.11",
 ]
 
-dependencies = ["numba", "trl"]
+dependencies = ["numba", "trl>=0.19.1,<0.20.0"]
 
 [tool.hatch.build.targets.wheel]
 only-include = ["src/fms_acceleration_aadp"]

plugins/attention-and-distributed-packing/src/fms_acceleration_aadp/framework_plugin_padding_free.py

@@ -20,7 +20,9 @@
 from fms_acceleration import AccelerationPlugin
 from peft import LoraConfig
 from transformers import DataCollatorForSeq2Seq, TrainingArguments
-from trl import DataCollatorForCompletionOnlyLM  # pylint: disable=import-error
+from trl import (  # pylint: disable=import-error, no-name-in-module
+    DataCollatorForCompletionOnlyLM,
+)
 import torch
 
 

plugins/fused-ops-and-kernels/src/fms_acceleration_foak/kernels/unsloth/cross_entropy_loss.py

@@ -187,7 +187,7 @@ def _cross_entropy_backward(
 pass
 
 
-MAX_FUSED_SIZE: tl.constexpr = 65536 # 2**16
+MAX_FUSED_SIZE = tl.constexpr(65536) # 2**16
 
 class Fast_CrossEntropyLoss(torch.autograd.Function):
     @staticmethod

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: tl.constexpr = 4
+ROPE_GROUP_SIZE = tl.constexpr(4)
 
 @triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],})
 @triton.jit