Upload BD3LM

modeling_bd3lm.py

@@ -16,6 +16,14 @@ try:
   FLEX_ATTN_AVAILABLE = True
 except:
   FLEX_ATTN_AVAILABLE = False
+# Flags required to enable jit fusion kernels
+try:
+  torch._C._jit_set_profiling_mode(False)
+  torch._C._jit_set_profiling_executor(False)
+  torch._C._jit_override_can_fuse_on_cpu(True)
+  torch._C._jit_override_fcan_fuse_on_gpu(True)
+except:
+  pass
 
 from .configuration_bd3lm import BD3LMConfig
 
@@ -69,6 +77,7 @@ def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
 def fused_flex_attention(q, k, v, mask=None):
     return flex_attention(q, k, v, block_mask=mask)
 
+
 def bias_dropout_add_scale(
     x: torch.Tensor,
     bias: typing.Optional[torch.Tensor],
@@ -93,13 +102,13 @@ def get_bias_dropout_add_scale(training):
 
   return _bias_dropout_add
 
-
 # function overload
 def modulate(x: torch.Tensor,
              shift: torch.Tensor,
              scale: torch.Tensor) -> torch.Tensor:
   return x * (1 + scale) + shift
 
+@torch.jit.script
 def bias_dropout_add_scale_fused_train(
     x: torch.Tensor,
     bias: typing.Optional[torch.Tensor],
@@ -109,6 +118,7 @@ def bias_dropout_add_scale_fused_train(
   return bias_dropout_add_scale(
     x, bias, scale, residual, prob, True)
 
+@torch.jit.script
 def bias_dropout_add_scale_fused_inference(
     x: torch.Tensor,
     bias: typing.Optional[torch.Tensor],
@@ -118,6 +128,7 @@ def bias_dropout_add_scale_fused_inference(
   return bias_dropout_add_scale(
     x, bias, scale, residual, prob, False)
 
+@torch.jit.script
 def modulate_fused(x: torch.Tensor,
                    shift: torch.Tensor,
                    scale: torch.Tensor) -> torch.Tensor: