[example] fused_linear_jsd

Author: v0i0

benchmarks/run.py

@@ -109,6 +109,11 @@
             ("examples.matmul_split_k", "matmul_split_k_tritonbench"),
         ],
     ),
+    "fused_linear_jsd": (
+        "tritonbench.operators.fused_linear_jsd.operator",
+        "examples.fused_linear_jsd",
+        "fused_linear_jsd_fwd",
+    ),
 }
 
 
@@ -407,7 +412,9 @@ def helion_method(
 
                 def _inner() -> Callable[..., Any] | object:
                     # BENCHMARK HOT PATH, do not add any new logic here
+                    kfunc._self = self
                     result = kfunc(*args, **kwargs)
+                    result = kfunc(*args)
                     if callable(result):
                         return result()
                     return result

examples/fused_linear_jsd.py

@@ -0,0 +1,32 @@
+import torch
+import helion
+import helion.language as hl
+
+@helion.kernel()
+def _kernel(beta, ignore_index, temperature, student_weight, teacher_weight, student_input, teacher_input):
+    student_logits = student_input @ student_weight.T
+    teacher_logits = teacher_input @ teacher_weight.T
+    loss = student_logits.new_empty(student_input.shape[0], dtype=torch.float)
+    for batch in hl.tile(student_logits.shape[0]):
+        student_prob = torch.log_softmax(student_logits[batch,:] / temperature, dim=-1)
+        teacher_prob = torch.log_softmax(teacher_logits[batch,:] / temperature, dim=-1)
+        student_prob = student_prob.to(torch.float).view(-1, student_prob.size(-1))
+        teacher_prob = teacher_prob.to(torch.float).view(-1, teacher_prob.size(-1))
+        beta_ = beta
+        m = torch.exp(student_prob) + beta_ * (torch.exp(teacher_prob) - torch.exp(student_prob))
+        teacher_div = torch.nn.functional.kl_div(torch.log(m), teacher_prob, reduction="none", log_target=True).sum(dim=-1)
+        student_div = torch.nn.functional.kl_div(torch.log(m), student_prob, reduction="none", log_target=True).sum(dim=-1)
+        batch_loss = student_div + beta * (teacher_div - student_div)
+        loss[batch] = batch_loss
+    return (loss / student_logits.shape[0]).sum()
+
+
+def fused_linear_jsd_fwd(student_input, teacher_input, label=None):
+    assert label is None
+    baseline_op = fused_linear_jsd_fwd._self.baseline_op
+    beta = baseline_op.jsd.beta
+    ignore_index = baseline_op.jsd.ignore_index
+    temperature = baseline_op.temperature
+    student_weight = baseline_op.student_lin.weight
+    teacher_weight = baseline_op.teacher_lin.weight
+    return _kernel(beta, ignore_index, temperature, student_weight, teacher_weight, student_input, teacher_input)

helion/autotuner/base_search.py

@@ -52,6 +52,7 @@
                 "misaligned address",  # CUDA Error
                 "PassManager::run failed",  # Triton Error
                 "illegal memory access",  # CUDA Error
+                "exceeds triton maximum tensor numel", # Triton Error
             ],
         )
     )
@@ -147,7 +148,7 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
         except PTXASError:
             self.log.warning(f"PTXASError compiling config: {config}")
         except Exception as e:
-            if not _expected_errors_regexp.search(str(e)):
+            if not _expected_errors_regexp.search(str(e) + str(e.__cause__)):
                 raise exc.TritonError(f"{type(e).__qualname__}: {e}", config) from e
             self.log.debug(f"Benchmarking failed: {type(e).__name__}: {e}")
         return inf