Triton monolithic head-loss kernel (#507)

fast_llm/layers/language_model/loss/config.py

@@ -92,6 +92,9 @@ class CombinableLossConfig(LanguageModelLossConfig):
 
     _abstract: typing.ClassVar[bool] = True
 
+    # Slot this loss occupies in the triton monolithic kernel; `None` when it has no triton implementation.
+    triton_kind: typing.ClassVar[str | None] = None
+
     use_triton: bool | None = Field(
         default=None,
         desc="Enable triton implementation. Default: use if available.",
@@ -102,6 +105,7 @@ class CombinableLossConfig(LanguageModelLossConfig):
 @config_class(dynamic_type={LanguageModelLossConfig: "label"})
 class LanguageModelLabelEntropyLossConfig(CombinableLossConfig):
     _abstract: typing.ClassVar[bool] = False
+    triton_kind: typing.ClassVar[str | None] = "ce"
 
     loss_type: EntropyLossType = Field(
         default=EntropyLossType.cross_entropy,
@@ -200,6 +204,7 @@ class LanguageModelZLossConfig(CombinableLossConfig):
     """Z-loss regularization to prevent overconfidence."""
 
     _abstract: typing.ClassVar[bool] = False
+    triton_kind: typing.ClassVar[str | None] = "z"
 
     @property
     def loss_class(self) -> "type[LanguageModelZLoss]":
@@ -244,6 +249,7 @@ class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig, Combina
     """Group-Relative Policy Optimization: per-token IS-ratio clipping."""
 
     _abstract: typing.ClassVar[bool] = False
+    triton_kind: typing.ClassVar[str | None] = "grpo"
 
     @property
     def loss_class(self) -> "type[LanguageModelGRPOLoss]":
@@ -257,6 +263,7 @@ class LanguageModelGSPOLossConfig(LanguageModelPolicyGradientLossConfig, Combina
     """Group Sequence Policy Optimization: sequence-level geometric-mean IS-ratio clipping."""
 
     _abstract: typing.ClassVar[bool] = False
+    triton_kind: typing.ClassVar[str | None] = "gspo"
 
     @property
     def loss_class(self) -> "type[LanguageModelGSPOLoss]":
@@ -276,6 +283,11 @@ class MonolithicLossConfig(LanguageModelLossConfig):
         desc="The combinable losses sharing a single softmax pass. They must agree on `logits_scale_factor`.",
         hint=FieldHint.core,
     )
+    use_triton: bool | None = Field(
+        default=None,
+        desc="Enable the triton implementation of the shared-softmax kernel. Default: use if available.",
+        hint=FieldHint.expert,
+    )
 
     def _validate(self) -> None:
         super()._validate()
@@ -289,6 +301,17 @@ def _validate(self) -> None:
                 raise ValueError(f"Loss `{name}` sets `use_triton`, which has no effect on a fused child loss.")
         # A single softmax serves one effective scale (stacked with the common model scale).
         Assert.eq(len({loss.logits_scale_factor for loss in self.losses.values()}), 1)
+        if self.use_triton:
+            # The triton kernel has one slot per loss kind, so it fuses at most one of each.
+            seen_kinds = set()
+            for name, loss in self.losses.items():
+                if loss.triton_kind is None:
+                    raise ValueError(f"Loss `{name}` (`{type(loss).__name__}`) has no triton fused implementation.")
+                if loss.triton_kind in seen_kinds:
+                    raise ValueError(
+                        f"The triton path fuses at most one `{loss.triton_kind}` loss; `{name}` is a duplicate."
+                    )
+                seen_kinds.add(loss.triton_kind)
 
     @property
     def loss_class(self) -> "type[LanguageModelLoss]":

fast_llm/layers/language_model/loss/entropy_loss.py

@@ -16,6 +16,9 @@
 )
 from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss
 
+if typing.TYPE_CHECKING:
+    from fast_llm.layers.language_model.loss.monolithic import _TritonContext
+
 
 class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](
     CombinableLoss, SingleLoss[ConfigType]
@@ -74,6 +77,21 @@ def fused_core(
             grad = grad * loss_mask.unsqueeze(-1)
         return loss, grad, None
 
+    def triton_add_inputs(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> None:
+        labels, grad_output, divisor = self.get_inputs(kwargs, split_index, register)
+        if context.labels is None:
+            context.labels = labels
+        if context.divisor is None:
+            context.divisor = divisor
+        context.ce = (grad_output,)
+
+    def triton_finish(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> tuple[torch.Tensor, None]:
+        return context.ce_loss, None
+
 
 class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig](
     CombinableLoss, SingleLoss[ConfigType]

fast_llm/layers/language_model/loss/loss.py

@@ -12,6 +12,9 @@
 from fast_llm.layers.language_model.loss.config import LanguageModelLossConfig, LanguageModelLossKwargs
 from fast_llm.utils import Assert
 
+if typing.TYPE_CHECKING:
+    from fast_llm.layers.language_model.loss.monolithic import _TritonContext
+
 
 class LanguageModelLoss[ConfigType: LanguageModelLossConfig](Configurable[ConfigType], torch.nn.Module):
     def __init__(
@@ -237,6 +240,36 @@ def finish(
         loss with an eager seam runs it here from the forward state its `fused_core` returned as `extra`."""
         return loss, extra, grad_logits
 
+    # The triton monolithic kernel has a fixed per-slot signature, so instead of the compiled loop each loss
+    # fills its slot in a shared `_TritonContext`: `triton_add_inputs` packs its kernel inputs (reusing the
+    # eager `get_inputs`) and `triton_finish` reads its `(loss, extra)` back from the kernel outputs. A loss
+    # whose gradient can't be produced in one kernel launch (an eager seam) sets `triton_needs_forward` and
+    # runs the seam in `triton_seam` over the pre-computed shared softmax.
+    triton_needs_forward: typing.ClassVar[bool] = False
+
+    def triton_add_inputs(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> None:
+        raise NotImplementedError()
+
+    def triton_seam(
+        self,
+        context: "_TritonContext",
+        softmax: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None],
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+    ) -> None:
+        """Run the eager pre-kernel seam over the shared softmax. No-op unless `triton_needs_forward`."""
+
+    def triton_finish(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> tuple[torch.Tensor, typing.Any]:
+        raise NotImplementedError()
+
+    def triton_metrics_enabled(self, register: bool) -> bool:
+        """Whether this loss emits diagnostic metrics from the kernel's shared softmax this step."""
+        return False
+
 
 def loss_forward_backward(
     grad_output: float | None, fn: typing.Callable, input_: torch.Tensor, *args, **kwargs

fast_llm/layers/language_model/loss/monolithic.py

@@ -1,16 +1,44 @@
+import dataclasses
 import typing
 
 import torch
 
 from fast_llm.core.distributed import ProcessGroup
 from fast_llm.engine.base_model.config import LossDef
 from fast_llm.engine.distributed.config import DistributedConfig
+from fast_llm.functional.config import TritonConfig
 from fast_llm.functional.entropy_loss import softmax_base
 from fast_llm.layers.language_model.loss.config import MonolithicLossConfig
 from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss
 from fast_llm.utils import safe_merge_dicts
 
 
+@dataclasses.dataclass
+class _TritonContext:
+    """Scratch threaded through the triton monolithic path. Each child packs its kernel-slot inputs (and the
+    shared labels / divisor) in `triton_add_inputs`, and reads its outputs back in `triton_finish`, so the
+    driver iterates children with no per-kind branching. The single `@triton.jit` kernel keeps its fixed
+    per-slot signature — that boundary is the one place the loss kinds are still named."""
+
+    # Shared inputs (first non-`None` child wins) and per-slot kernel inputs.
+    labels: torch.Tensor | None = None
+    divisor: float | None = None
+    ce: tuple | None = None
+    z: tuple | None = None
+    grpo: tuple | None = None
+    gspo_coeff: torch.Tensor | None = None
+    # Per-slot kernel outputs; GSPO's loss / new-log-probs instead come from its eager seam.
+    ce_loss: torch.Tensor | None = None
+    z_loss: torch.Tensor | None = None
+    grpo_loss: torch.Tensor | None = None
+    grpo_new_logprobs: torch.Tensor | None = None
+    gspo_loss: torch.Tensor | None = None
+    gspo_new_logprobs: torch.Tensor | None = None
+    # Shared metrics precursors from the kernel's softmax + `Σ exp·logits_norm` (set only when metrics are on).
+    new_log_probs: torch.Tensor | None = None
+    entropy_per_token: torch.Tensor | None = None
+
+
 @torch.compile
 def _monolithic_core(
     children: tuple[LanguageModelLoss, ...],
@@ -93,6 +121,9 @@ def __init__(
         )
         # The shared softmax serves one effective scale; the config validates the children agree on it.
         self._softmax_scale_factor = self._children[0]._logits_scale_factor
+        # The triton kernel fuses at most one loss per kind; anything else falls back to the compiled path.
+        triton_kinds = [child._config.triton_kind for child in self._children]
+        self._triton_valid = None not in triton_kinds and len(set(triton_kinds)) == len(triton_kinds)
 
     def forward_backward(
         self,
@@ -102,6 +133,8 @@ def forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> "tuple[torch.Tensor | None, torch.Tensor | None]":
+        if self._triton_valid and TritonConfig.enabled(logits.device, self._config.use_triton):
+            return self._triton_forward_backward(logits, kwargs, losses, split_index, grad_logits)
         register = losses is not None
         arguments = tuple(child.get_inputs(kwargs, split_index, register) for child in self._children)
         group = self._parallel_dim.group if self._vocab_parallel else None
@@ -121,6 +154,78 @@ def forward_backward(
             total_loss = weighted if total_loss is None else total_loss + weighted
         return total_loss, grad_logits
 
+    def _triton_forward_backward(
+        self,
+        logits: torch.Tensor,
+        kwargs: dict[str, typing.Any],
+        losses: dict | None,
+        split_index: int,
+        grad_logits: torch.Tensor | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        from fast_llm.functional.triton.monolithic_loss import (
+            _monolithic_forward_reduce,
+            triton_monolithic_loss_forward_backward,
+        )
+
+        register = losses is not None
+        group = self._parallel_dim.group if self._vocab_parallel else None
+
+        # Each child packs its own kernel slot (and the shared labels / divisor) — no per-kind branching.
+        context = _TritonContext()
+        for child in self._children:
+            child.triton_add_inputs(context, kwargs, split_index, register)
+        compute_metrics = any(child.triton_metrics_enabled(register) for child in self._children)
+
+        # A child with an eager segment seam (GSPO) can't produce its gradient in one kernel launch: it needs
+        # the shared softmax first, so run the reduced forward once, then let each such child add its
+        # per-token backward coefficient the kernel superposes with the in-kernel losses.
+        softmax = None
+        if any(child.triton_needs_forward for child in self._children):
+            softmax = _monolithic_forward_reduce(logits, context.labels, group, self._softmax_scale_factor)
+            for child in self._children:
+                child.triton_seam(context, softmax, kwargs, split_index)
+
+        (
+            context.ce_loss,
+            context.z_loss,
+            context.grpo_loss,
+            context.grpo_new_logprobs,
+            grad_logits,
+            softmax,
+            weighted_logits_sum,
+        ) = triton_monolithic_loss_forward_backward(
+            logits,
+            None if context.labels is None else context.labels.contiguous(),
+            grad_logits,
+            self._softmax_scale_factor,
+            group,
+            1.0 if context.divisor is None else context.divisor,
+            ce=context.ce,
+            z=context.z,
+            grpo=context.grpo,
+            gspo_coeff=context.gspo_coeff,
+            softmax=softmax,
+            compute_metrics=compute_metrics,
+        )
+
+        # Metrics reuse the kernel's shared softmax: per-token new log-probs and the entropy from the kernel's
+        # `Σ exp·logits_norm`, so no second softmax pass. Derived once here and shared by each metric loss.
+        if compute_metrics:
+            max_logits, sum_exp_logits, predicted_logits = softmax
+            log_sum_exp_logits = sum_exp_logits.log()
+            context.new_log_probs = predicted_logits - max_logits - log_sum_exp_logits
+            context.entropy_per_token = log_sum_exp_logits - weighted_logits_sum / sum_exp_logits
+
+        total_loss = None
+        for child in self._children:
+            loss, extra = child.triton_finish(context, kwargs, split_index, register)
+            if child._do_register_loss:
+                child._register_loss(child.name, loss, losses)
+            child.register_combinable_extras(extra, kwargs, losses)
+            weighted = loss if child.weight == 1 else loss * child.weight
+            total_loss = weighted if total_loss is None else total_loss + weighted
+        return total_loss, grad_logits
+
     def get_preprocessing_config(self) -> dict[str, typing.Any]:
         return safe_merge_dicts(*(child.get_preprocessing_config() for child in self._children))