diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py
index 05eaae520..a10881a68 100644
--- a/fast_llm/functional/entropy_loss.py
+++ b/fast_llm/functional/entropy_loss.py
@@ -2,7 +2,6 @@
 
 from fast_llm.core.distributed import ProcessGroup, ReduceOp, all_reduce
 from fast_llm.functional.config import EntropyLossType, TargetFormat
-from fast_llm.functional.utils import reduce_losses
 from fast_llm.utils import Assert
 
 
@@ -92,8 +91,7 @@ def torch_entropy_loss_forward_backward(
     return loss.detach_(), grad
 
 
-@torch.compile
-def fused_softmax_base(
+def softmax_base(
     logits: torch.Tensor,  # (*batch, vocab)
     logits_scale_factor: float = 1.0,
     group: ProcessGroup | None = None,
@@ -106,6 +104,9 @@ def fused_softmax_base(
     in a numerically stable way and with tensor-parallel support.
     Warning: The returned values are regularized by `logits_max`.
         The regularization typically but not always cancels out in derived quantities.
+
+    Un-compiled so it can be inlined into a `@torch.compile` boundary that fuses several losses over a
+    single softmax; `fused_softmax_base` is the compiled standalone wrapper.
     """
     logits = logits.float()
     if logits_scale_factor != 1.0:
@@ -121,9 +122,13 @@ def fused_softmax_base(
     return logits_norm, exp_logits, sum_exp_logits, logits_max
 
 
-@torch.compile
-def _fused_reverse_kl_base_from_distribution(
-    logits: torch.Tensor,  # (*batch, vocab)
+fused_softmax_base = torch.compile(softmax_base)
+
+
+def reverse_kl_from_distribution_core(
+    logits_norm: torch.Tensor,  # (*batch, vocab)
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
     target: torch.Tensor,  # (*batch, vocab)
     grad_output: float | None,
     logits_scale_factor: float,
@@ -131,13 +136,14 @@ def _fused_reverse_kl_base_from_distribution(
     group: ProcessGroup | None = None,
     temperature: float = 1.0,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:  # (*batch,), (*batch, vocab)
+    """Reverse-KL from a precomputed student softmax (adding a teacher softmax when the target is logits).
+    Un-compiled core, inlined into a `@torch.compile` boundary."""
     assert target_format in (TargetFormat.logits, TargetFormat.probabilities)
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
     predicted_log_probability = logits_norm - sum_exp_logits.log().unsqueeze(-1)
     predicted_probability = exp_logits / sum_exp_logits.unsqueeze(-1)
 
     if target_format == TargetFormat.logits:
-        target_logits_norm, _, sum_exp_target_logits, _ = fused_softmax_base(
+        target_logits_norm, _, sum_exp_target_logits, _ = softmax_base(
             target, logits_scale_factor / temperature, group
         )
         target_log_probability = target_logits_norm - sum_exp_target_logits.log().unsqueeze(-1)
@@ -161,9 +167,10 @@ def _fused_reverse_kl_base_from_distribution(
     return per_sample_loss, grad
 
 
-@torch.compile
-def _fused_cross_entropy_base_from_distribution(
-    logits: torch.Tensor,  # (*batch, vocab)
+def cross_entropy_from_distribution_core(
+    logits_norm: torch.Tensor,  # (*batch, vocab)
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
     target: torch.Tensor,  # (*batch, vocab)
     grad_output: float | None,
     logits_scale_factor: float,
@@ -172,10 +179,10 @@ def _fused_cross_entropy_base_from_distribution(
     temperature: float = 1.0,
     return_kl_loss: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:  # (*batch,), (*batch, vocab)
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
-
+    """Cross-entropy / forward-KL from a precomputed student softmax (adding a teacher softmax when the
+    target is logits). Un-compiled core, inlined into a `@torch.compile` boundary."""
     if target_format == TargetFormat.logits:
-        target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = fused_softmax_base(
+        target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = softmax_base(
             target, logits_scale_factor / temperature, group
         )
         target = exp_logits_targets / sum_exp_target_logits.unsqueeze(-1)
@@ -207,8 +214,7 @@ def _fused_cross_entropy_base_from_distribution(
     return per_sample_loss, grad
 
 
-@torch.compile
-def fused_predicted_logits_from_labels(
+def predicted_logits_from_labels(
     logits: torch.Tensor,  # (*batch, vocab)
     target: torch.Tensor,  # (*batch,)
     loss_mask: torch.Tensor,  # (*batch,), == target>=0
@@ -218,9 +224,12 @@ def fused_predicted_logits_from_labels(
     Recover the value of the logits at the target index, with support for masking (target < 0) and tensor parallelism.
     In the simple case, equivalent to `logits.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)`
 
-    Normally used in combination with `fused_softmax_base`, may also recover probabilities or log probabilities:
+    May also recover probabilities or log probabilities:
     `predicted_probabilities = predicted_logits.exp() / sum_exp_logits`
-    `predicted_log_probabilities = predicted_logits / sum_exp_logits.log()`
+    `predicted_log_probabilities = predicted_logits - sum_exp_logits.log()`
+
+    Un-compiled core, inlined into a `@torch.compile` boundary; `fused_predicted_logits_from_labels` is the
+    compiled standalone wrapper.
     """
 
     if group is None:
@@ -243,19 +252,24 @@ def fused_predicted_logits_from_labels(
     return predicted_logits, target_masked, target_mask
 
 
-@torch.compile
-def _fused_cross_entropy_base_from_labels(
-    logits: torch.Tensor,  # (*batch, vocab)
+fused_predicted_logits_from_labels = torch.compile(predicted_logits_from_labels)
+
+
+def cross_entropy_from_labels_core(
+    logits_norm: torch.Tensor,  # (*batch, vocab), == logits - max
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
     target: torch.Tensor,  # (*batch,)
-    loss_mask: torch.Tensor,  # (*batch,)
-    grad_output: float | None,
-    logits_scale_factor: float,
+    loss_mask: torch.Tensor,  # (*batch,), == target>=0
+    grad_output: float | None,  # already normalized: raw_grad_output / divisor * logits_scale_factor
     group: ProcessGroup | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:  # (*batch,), (*batch, vocab)
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
-    predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels(
-        logits_norm, target, loss_mask, group
-    )
+) -> tuple[torch.Tensor, torch.Tensor | None]:  # (*batch,) unmasked, (*batch, vocab) unmasked
+    """
+    Cross-entropy from labels, taking the already-computed shared softmax tensors. Returns the unmasked
+    per-sample loss and (when `grad_output` is given) the unmasked gradient; the caller applies the loss
+    mask, reduction, and dtype cast. Un-compiled core, inlined into a `@torch.compile` boundary.
+    """
+    predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group)
 
     # CE loss = mean(log(sum_exp_logits) - sum(probabilities * logits))
     # KL loss is the same because P * log(P) == 0.
@@ -274,73 +288,22 @@ def _fused_cross_entropy_base_from_labels(
     return per_sample_loss, grad
 
 
-@torch.compile
-def fused_entropy_loss_forward_backward(
-    logits: torch.Tensor,  # (*batch, vocab)
-    target: torch.Tensor,  # (*batch,) or (*batch, vocab)
-    loss_mask: torch.Tensor | None,  # (*batch,)
-    grad_logits: torch.Tensor | None = None,
-    grad_output: float | None = None,
-    group: torch.distributed.ProcessGroup | None = None,
-    logits_scale_factor: float = 1.0,
-    temperature: float = 1.0,
-    target_format: TargetFormat = TargetFormat.labels,
-    entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy,
-    divisor: float | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
+def z_loss_core(
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
+    logits_max: torch.Tensor,  # (*batch,)
+    grad_output: float | None,  # already normalized: raw_grad_output / divisor * logits_scale_factor
+) -> tuple[torch.Tensor, torch.Tensor | None]:  # (*batch,) unmasked, (*batch, vocab) unmasked
     """
-    A fused implementation of cross-entropy with torch compile.
-    It is an improvement over the pytorch implementation because of the fused casting, both in speed and memory,
-    but still suboptimal because it needs multiple kernels.
+    Z-loss from the already-computed shared softmax tensors. Returns the unmasked per-sample loss term
+    (`log_sum_exp ** 2`) and (when `grad_output` is given) the unmasked gradient; the caller applies the
+    loss mask, reduction, and dtype cast. z-loss needs the un-regularized log-sum-exp, so it adds back
+    `logits_max` (cross-entropy cancels it). Un-compiled core, inlined into a `@torch.compile` boundary.
     """
-    if divisor is None:
-        divisor = logits.shape[:-1].numel()
-    grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
-    if target_format == TargetFormat.labels:
-        assert entropy_loss_type in (EntropyLossType.cross_entropy, EntropyLossType.forward_kl)
-        assert loss_mask is None
-        loss_mask = target >= 0
-        losses, grad = _fused_cross_entropy_base_from_labels(
-            logits,
-            target,
-            loss_mask,
-            grad_output,
-            logits_scale_factor,
-            group,
-        )
-    elif entropy_loss_type in (EntropyLossType.cross_entropy, EntropyLossType.forward_kl):
-        losses, grad = _fused_cross_entropy_base_from_distribution(
-            logits,
-            target,
-            grad_output,
-            logits_scale_factor,
-            target_format,
-            group,
-            temperature,
-            return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl,
-        )
-    elif entropy_loss_type == EntropyLossType.reverse_kl:
-        losses, grad = _fused_reverse_kl_base_from_distribution(
-            logits,
-            target,
-            grad_output,
-            logits_scale_factor,
-            target_format,
-            group,
-            temperature,
-        )
+    log_sum_exp_logits = sum_exp_logits.log() + logits_max
+    loss_term = log_sum_exp_logits**2
+    if grad_output is None:
+        grad = None
     else:
-        raise NotImplementedError(entropy_loss_type)
-
-    loss = reduce_losses(losses, divisor, loss_mask)
-
-    if grad is not None:
-        if loss_mask is not None:
-            grad = grad * loss_mask.unsqueeze(-1)
-        grad = grad.to(logits.dtype)
-        if grad_logits is None:
-            grad_logits = grad
-        else:
-            grad_logits.add_(grad)
-
-    return loss, grad_logits
+        grad = (2 * grad_output * (log_sum_exp_logits / sum_exp_logits)).unsqueeze(-1) * exp_logits
+    return loss_term, grad
diff --git a/fast_llm/functional/triton/monolithic_loss.py b/fast_llm/functional/triton/monolithic_loss.py
new file mode 100644
index 000000000..92706c580
--- /dev/null
+++ b/fast_llm/functional/triton/monolithic_loss.py
@@ -0,0 +1,353 @@
+import torch
+
+from fast_llm.core.distributed import ReduceOp, all_reduce
+from fast_llm.functional.triton import tl, tl_arange, tl_constexpr, triton, triton_jit
+from fast_llm.functional.triton.entropy_loss import (
+    parallel_sum_exp_logits,
+    triton_cross_entropy_forward_from_labels_parallel_kernel,
+    triton_fused_softmax_base,
+)
+from fast_llm.functional.utils import reduce_losses
+
+
+@triton_jit()
+def triton_monolithic_loss_forward_backward_kernel(
+    logits_ptr,
+    labels_ptr,
+    n_cols: tl_constexpr,
+    logits_stride_0: tl_constexpr,
+    block_size: tl_constexpr,
+    ce_losses_ptr=None,
+    z_losses_ptr=None,
+    grpo_losses_ptr=None,
+    new_logprobs_mean_parts_ptr=None,
+    z_loss_mask_ptr=None,
+    advantages_ptr=None,
+    old_log_probs_ptr=None,
+    num_labels_in_seq_ptr=None,
+    gspo_coeff_ptr=None,
+    max_logits_ptr=None,
+    sum_exp_logits_ptr=None,
+    predicted_logits_ptr=None,
+    weighted_logits_sum_ptr=None,
+    grad_logits_ptr=None,
+    grad_logits_stride_0: tl_constexpr = None,
+    grad_losses_ce=0.0,
+    grad_losses_z=0.0,
+    grad_losses_grpo=0.0,
+    col_min: tl_constexpr = 0,
+    logits_scale_factor: tl_constexpr = 1.0,
+    epsilon_low: tl_constexpr = 0.2,
+    epsilon_high: tl_constexpr = 0.2,
+    accumulate: tl_constexpr = False,
+):
+    """One shared softmax feeding several label-based losses over the same logits row. Each enabled loss
+    (selected by the presence of its output/input pointers) stores its own forward scalar, but their
+    gradients superpose into two per-row coefficients: `grad_j = prob_coeff * softmax_j - label_coeff *
+    delta_{j, label}`. The softmax is computed in-kernel when `max_logits_ptr`/`sum_exp_logits_ptr` are
+    absent (single-pass, no tensor parallelism), or loaded from a reduced forward pass otherwise."""
+    block_idx = tl.program_id(0).to(tl.int64)
+    logits_ptr = logits_ptr + block_idx * logits_stride_0
+
+    # The shared label feeds cross-entropy, GRPO, and GSPO; `labels_ptr` is set whenever any of them is
+    # present. The defaults keep both variables defined on the label-free path — triton compiles every branch,
+    # so a variable used later must be defined on all of them.
+    label_valid = False
+    label_idx = 0
+    if labels_ptr is not None:
+        label_idx = tl.load(labels_ptr + block_idx)
+        label_valid = label_idx >= 0
+        label_idx -= col_min
+
+    # A masked row of a single label-based loss (nothing else needs the softmax on a label-less row) contributes
+    # only zeros. The active-loss *set* is compile-time; only the per-token mask is runtime. Skip such a row's
+    # softmax and backward `exp` — but via a branch, not an early `return` (a mid-kernel return compiles into
+    # heavy register spills in this large kernel), so the masked-row work drops without hurting the valid path.
+    row_inactive = (
+        (not label_valid)
+        and z_losses_ptr is None
+        and gspo_coeff_ptr is None
+        and weighted_logits_sum_ptr is None
+        and (ce_losses_ptr is not None or grpo_losses_ptr is not None)
+    )
+    # Defaults keep the softmax outputs defined on the inactive path (whose backward branch never reads them).
+    exp_logits = tl.zeros([block_size], tl.float32)
+    max_logits = 0.0
+    sum_exp_logits = 1.0
+    if row_inactive:
+        pass
+    elif max_logits_ptr is None or sum_exp_logits_ptr is None:
+        exp_logits, sum_exp_logits, max_logits, col_offsets, mask = triton_fused_softmax_base(
+            logits_ptr, n_cols=n_cols, block_size=block_size, logits_scale_factor=logits_scale_factor
+        )
+    else:
+        max_logits = tl.load(max_logits_ptr + block_idx)
+        sum_exp_logits = tl.load(sum_exp_logits_ptr + block_idx)
+
+    log_sum_exp_logits = tl.log(sum_exp_logits) + max_logits
+
+    # Target-index logit shared by cross-entropy and GRPO; the defaults keep it defined when neither is present.
+    predicted_logit = 0.0
+    new_log_prob = 0.0
+    if ce_losses_ptr is not None or grpo_losses_ptr is not None:
+        if predicted_logits_ptr is not None:
+            predicted_logit = tl.load(predicted_logits_ptr + block_idx)
+        elif label_valid and label_idx >= 0 and label_idx < n_cols:
+            predicted_logit = tl.load(logits_ptr + label_idx).to(tl.float32)
+            if logits_scale_factor != 1.0:
+                predicted_logit *= logits_scale_factor
+        new_log_prob = predicted_logit - log_sum_exp_logits
+
+    prob_coeff = 0.0
+    label_coeff = 0.0
+
+    if ce_losses_ptr is not None:
+        if label_valid:
+            tl.store(ce_losses_ptr + block_idx, log_sum_exp_logits - predicted_logit)
+            grad_losses = grad_losses_ce * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_ce
+            prob_coeff += grad_losses
+            label_coeff += grad_losses
+        else:
+            tl.store(ce_losses_ptr + block_idx, 0.0)
+
+    if z_losses_ptr is not None:
+        if z_loss_mask_ptr is None or tl.load(z_loss_mask_ptr + block_idx) != 0:
+            tl.store(z_losses_ptr + block_idx, log_sum_exp_logits * log_sum_exp_logits)
+            grad_losses = grad_losses_z * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_z
+            prob_coeff += 2.0 * grad_losses * log_sum_exp_logits
+        else:
+            tl.store(z_losses_ptr + block_idx, 0.0)
+
+    if grpo_losses_ptr is not None:
+        if label_valid:
+            old_log_prob = tl.load(old_log_probs_ptr + block_idx).to(tl.float32)
+            advantage = tl.load(advantages_ptr + block_idx).to(tl.float32)
+            ratio = tl.exp(new_log_prob - old_log_prob)
+            clipped_ratio = tl.minimum(tl.maximum(ratio, 1.0 - epsilon_low), 1.0 + epsilon_high)
+            tl.store(grpo_losses_ptr + block_idx, -tl.minimum(ratio * advantage, clipped_ratio * advantage))
+            grad_losses = grad_losses_grpo * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_grpo
+            # clip_factor = clamp_min(A, 0) * (ratio <= 1 + eps_h) + clamp_max(A, 0) * (ratio >= 1 - eps_l)
+            coeff = (
+                tl.maximum(advantage, 0.0) * (ratio <= 1.0 + epsilon_high)
+                + tl.minimum(advantage, 0.0) * (ratio >= 1.0 - epsilon_low)
+            ) * (ratio * grad_losses)
+            prob_coeff += coeff
+            label_coeff += coeff
+            if new_logprobs_mean_parts_ptr is not None:
+                num_labels = tl.load(num_labels_in_seq_ptr + block_idx).to(tl.float32)
+                tl.store(new_logprobs_mean_parts_ptr + block_idx, new_log_prob / tl.maximum(num_labels, 1.0))
+        else:
+            tl.store(grpo_losses_ptr + block_idx, 0.0)
+            if new_logprobs_mean_parts_ptr is not None:
+                tl.store(new_logprobs_mean_parts_ptr + block_idx, 0.0)
+
+    if gspo_coeff_ptr is not None:
+        # The GSPO per-token coefficient is fully scaled by its eager segment seam.
+        coeff = tl.load(gspo_coeff_ptr + block_idx).to(tl.float32)
+        prob_coeff += coeff
+        label_coeff += coeff
+
+    # The backward re-stream both writes the combined gradient and accumulates the entropy's Σ exp·logits_norm,
+    # so it runs when either is requested — metrics without a gradient (e.g. a zero-weight policy loss) still
+    # emit `weighted_logits_sum`; the gradient store itself is gated on `grad_logits_ptr`.
+    if grad_logits_ptr is not None or weighted_logits_sum_ptr is not None:
+        weighted_logits_sum = 0.0
+        col_offset_start: tl.constexpr = (n_cols - 1) // block_size * block_size
+        for col_offset in tl.static_range(col_offset_start, -1, -block_size):
+            col_offsets = tl_arange(col_offset, col_offset + block_size)
+            mask = col_offsets < n_cols
+            if row_inactive:
+                # No softmax was computed for this row; its gradient is exactly zero.
+                grad_logits = col_offsets * 0.0
+            else:
+                if max_logits_ptr is not None or sum_exp_logits_ptr is not None or col_offset != col_offset_start:
+                    logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32)
+                    if logits_scale_factor != 1.0:
+                        logits *= logits_scale_factor
+                    exp_logits = tl.exp(logits - max_logits)
+                    if weighted_logits_sum_ptr is not None:
+                        # Local (per-rank) Σ exp·logits_norm feeding the policy entropy; the caller all-reduces
+                        # over the vocab group. Zero `logits_norm` on masked columns first (they load as -inf),
+                        # so the product is 0 there rather than 0·-inf = nan. `max_logits` is final here (this is
+                        # the backward re-stream), so the accumulation needs no online rescaling.
+                        weighted_logits_sum += tl.sum(exp_logits * tl.where(mask, logits - max_logits, 0.0), 0)
+                grad_logits = prob_coeff * (exp_logits / sum_exp_logits)
+                if label_valid:
+                    grad_logits = tl.where(col_offsets == label_idx, grad_logits - label_coeff, grad_logits)
+            if grad_logits_ptr is not None:
+                grad_logits_col_ptr = grad_logits_ptr + block_idx * grad_logits_stride_0 + col_offsets
+                if accumulate:
+                    grad_logits += tl.load(grad_logits_col_ptr, mask=mask)
+                tl.store(grad_logits_col_ptr, grad_logits, mask=mask)
+        if weighted_logits_sum_ptr is not None:
+            tl.store(weighted_logits_sum_ptr + block_idx, weighted_logits_sum)
+
+
+def _monolithic_forward_reduce(
+    logits: torch.Tensor,
+    labels: torch.Tensor | None,
+    group: torch.distributed.ProcessGroup | None,
+    logits_scale_factor: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Explicit forward pass: per-row softmax and target-index logit, reduced over the vocab group to the
+    global values the monolithic kernel then loads. Reuses the shared cross-entropy forward kernel. Used for
+    tensor parallelism and — regardless of parallelism — to feed the GSPO segment seam before the backward."""
+    n_rows = logits.shape[:-1].numel()
+    n_cols = logits.size(-1)
+    block_size = min(triton.next_power_of_2(n_cols), 32768)
+    num_warps = 4 if block_size < 2048 else (8 if block_size < 8192 else 16)
+    local_max_logits = torch.empty(n_rows, dtype=torch.float, device=logits.device)
+    sum_exp_logits = torch.empty_like(local_max_logits)
+    predicted_logits = torch.empty_like(local_max_logits) if labels is not None else None
+    triton_cross_entropy_forward_from_labels_parallel_kernel[(n_rows,)](
+        logits,
+        labels,
+        max_logits_ptr=local_max_logits,
+        sum_exp_logits_ptr=sum_exp_logits,
+        predicted_logits_ptr=predicted_logits,
+        col_min=n_cols * group.rank() if group is not None else 0,
+        n_cols=n_cols,
+        logits_stride_0=logits.stride(-2),
+        block_size=block_size,
+        num_warps=num_warps,
+        logits_scale_factor=logits_scale_factor,
+    )
+    if group is None:
+        return local_max_logits, sum_exp_logits, predicted_logits
+    max_logits, sum_exp_logits = parallel_sum_exp_logits(sum_exp_logits, local_max_logits, group)
+    if predicted_logits is not None:
+        all_reduce(predicted_logits, op=ReduceOp.SUM, group=group)
+    return max_logits, sum_exp_logits, predicted_logits
+
+
+def triton_monolithic_loss_forward_backward(
+    logits: torch.Tensor,  # (*batch, vocab)
+    labels: torch.Tensor | None,  # (*batch,) — shared by cross-entropy / GRPO / GSPO
+    grad_logits: torch.Tensor | None,
+    logits_scale_factor: float,
+    group: torch.distributed.ProcessGroup | None,
+    divisor: float,
+    *,
+    ce: tuple[float | None] | None = None,  # cross-entropy: (weighted grad_output,); `None` => absent
+    z: tuple[torch.Tensor | None, float | None] | None = None,  # (loss_mask, weighted grad_output)
+    # GRPO: (advantages, old_log_probabilities, weighted grad_output, epsilon_low, epsilon_high, num_labels_in_seq)
+    grpo: tuple | None = None,
+    gspo_coeff: torch.Tensor | None = None,  # per-token backward coefficient from the eager segment seam
+    softmax: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | None = None,  # precomputed (max, sum, predicted)
+    compute_metrics: bool = False,  # also emit the reduced softmax and Σ exp·logits_norm for policy metrics
+    block_size: int | None = None,
+    num_warps: int | None = None,
+) -> tuple[
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | None,
+    torch.Tensor | None,
+]:
+    """Dispatch the monolithic label-loss kernel over a single shared softmax. Returns the reduced per-loss
+    scalars `(cross_entropy, z, grpo, grpo_new_logprobs_mean)` (each `None` when the loss is absent), the
+    accumulated `grad_logits`, and — for policy metrics — the reduced `softmax` (max, sum, predicted) and the
+    all-reduced per-row `weighted_logits_sum` (`Σ exp·logits_norm`, both `None` unless `compute_metrics`).
+    `softmax` is provided by the caller when the forward was already run (the GSPO seam); otherwise it is
+    computed in-kernel (`group is None`, no metrics) or by a reduced forward pass (tensor parallel or metrics).
+    A present loss whose weighted grad_output is `None` still emits its forward scalar but no gradient term."""
+    assert logits.is_contiguous()
+    if labels is not None:
+        assert labels.is_contiguous()
+    n_rows = logits.shape[:-1].numel()
+    n_cols = logits.size(-1)
+    if block_size is None:
+        block_size = min(triton.next_power_of_2(n_cols), 32768)
+    if num_warps is None:
+        num_warps = 4 if block_size < 2048 else (8 if block_size < 8192 else 16)
+
+    # Forward-scalar buffers (needed for the total loss even when nothing is registered this step).
+    ce_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if ce is not None else None
+    z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if z is not None else None
+    grpo_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if grpo is not None else None
+
+    ce_grad_output = ce[0] if ce is not None else None
+    z_loss_mask, z_grad_output = z if z is not None else (None, None)
+    if grpo is not None:
+        advantages, old_log_probabilities, grpo_grad_output, epsilon_low, epsilon_high, num_labels_in_seq = grpo
+    else:
+        advantages = old_log_probabilities = num_labels_in_seq = grpo_grad_output = None
+        epsilon_low = epsilon_high = 0.2
+    new_logprobs_mean_parts = (
+        torch.empty(n_rows, dtype=torch.float, device=logits.device)
+        if grpo is not None and num_labels_in_seq is not None
+        else None
+    )
+    for tensor in (z_loss_mask, advantages, old_log_probabilities, num_labels_in_seq, gspo_coeff):
+        if tensor is not None:
+            assert tensor.is_contiguous()
+
+    # Metrics reuse the reduced softmax, so run the explicit forward pass even without tensor parallelism.
+    if softmax is None and (group is not None or compute_metrics):
+        softmax = _monolithic_forward_reduce(logits, labels, group, logits_scale_factor)
+    max_logits, sum_exp_logits, predicted_logits = softmax if softmax is not None else (None, None, None)
+
+    has_grad = (
+        ce_grad_output is not None
+        or z_grad_output is not None
+        or grpo_grad_output is not None
+        or gspo_coeff is not None
+    )
+    # The entropy's `Σ exp·logits_norm` is accumulated in the backward re-stream, which the kernel runs for the
+    # metrics alone when no gradient is requested (e.g. a zero-weight policy loss logging its diagnostics).
+    weighted_logits_sum = torch.empty(n_rows, dtype=torch.float, device=logits.device) if compute_metrics else None
+    backward_kwargs = {}
+    if compute_metrics:
+        backward_kwargs["weighted_logits_sum_ptr"] = weighted_logits_sum
+    if has_grad:
+        accumulate = grad_logits is not None
+        grad_logits = torch.empty_like(logits) if grad_logits is None else grad_logits
+        backward_kwargs.update(
+            grad_logits_ptr=grad_logits,
+            grad_logits_stride_0=grad_logits.stride(-2),
+            accumulate=accumulate,
+            grad_losses_ce=0.0 if ce_grad_output is None else ce_grad_output / divisor,
+            grad_losses_z=0.0 if z_grad_output is None else z_grad_output / divisor,
+            grad_losses_grpo=0.0 if grpo_grad_output is None else grpo_grad_output / divisor,
+        )
+
+    triton_monolithic_loss_forward_backward_kernel[(n_rows,)](
+        logits,
+        labels,
+        n_cols=n_cols,
+        logits_stride_0=logits.stride(-2),
+        block_size=block_size,
+        num_warps=num_warps,
+        logits_scale_factor=logits_scale_factor,
+        col_min=n_cols * group.rank() if group is not None else 0,
+        epsilon_low=epsilon_low,
+        epsilon_high=epsilon_high,
+        ce_losses_ptr=ce_losses,
+        z_losses_ptr=z_losses,
+        grpo_losses_ptr=grpo_losses,
+        new_logprobs_mean_parts_ptr=new_logprobs_mean_parts,
+        z_loss_mask_ptr=z_loss_mask,
+        advantages_ptr=advantages,
+        old_log_probs_ptr=old_log_probabilities,
+        num_labels_in_seq_ptr=num_labels_in_seq,
+        gspo_coeff_ptr=gspo_coeff,
+        max_logits_ptr=max_logits,
+        sum_exp_logits_ptr=sum_exp_logits,
+        predicted_logits_ptr=predicted_logits,
+        **backward_kwargs,
+    )
+
+    if weighted_logits_sum is not None and group is not None:
+        all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group)
+
+    return (
+        None if ce_losses is None else reduce_losses(ce_losses, divisor),
+        None if z_losses is None else reduce_losses(z_losses, divisor),
+        None if grpo_losses is None else reduce_losses(grpo_losses, divisor),
+        None if new_logprobs_mean_parts is None else new_logprobs_mean_parts.sum(),
+        grad_logits,
+        softmax,
+        weighted_logits_sum,
+    )
diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py
index 4f3cfae8d..525d03256 100644
--- a/fast_llm/layers/language_model/loss/config.py
+++ b/fast_llm/layers/language_model/loss/config.py
@@ -86,20 +86,38 @@ def get_reference_models(self) -> set[str]:
         return set()
 
 
+@config_class()
+class CombinableLossConfig(LanguageModelLossConfig):
+    """Base for losses that share the vocabulary softmax via `fused_core`, so several can be fused together."""
+
+    _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.",
+        hint=FieldHint.expert,
+    )
+
+
 @config_class(dynamic_type={LanguageModelLossConfig: "label"})
-class LanguageModelLabelEntropyLossConfig(LanguageModelLossConfig):
+class LanguageModelLabelEntropyLossConfig(CombinableLossConfig):
     _abstract: typing.ClassVar[bool] = False
+    triton_kind: typing.ClassVar[str | None] = "ce"
 
     loss_type: EntropyLossType = Field(
         default=EntropyLossType.cross_entropy,
         desc="Type of loss to use.",
         hint=FieldHint.core,
     )
-    use_triton: bool | None = Field(
-        default=None,
-        desc="Enable triton implementation. Default: use if available.",
-        hint=FieldHint.expert,
-    )
+
+    def _validate(self) -> None:
+        super()._validate()
+        # Labels are one-hot, so their forward-KL reduces to cross-entropy; reverse-KL is not defined.
+        if self.loss_type == EntropyLossType.reverse_kl:
+            raise ValueError("`reverse_kl` is not supported for a label loss.")
 
     @classmethod
     def _from_dict(cls, default: dict[str, typing.Any], strict: bool = True) -> typing.Self:
@@ -116,7 +134,7 @@ def loss_class(self) -> "type[LanguageModelLabelEntropyLoss]":
 
 
 @config_class(dynamic_type={LanguageModelLossConfig: "distillation"})
-class LanguageModelDistillationLossConfig(LanguageModelLossConfig):
+class LanguageModelDistillationLossConfig(CombinableLossConfig):
     _abstract: typing.ClassVar[bool] = False
 
     loss_type: EntropyLossType = Field(
@@ -135,11 +153,6 @@ class LanguageModelDistillationLossConfig(LanguageModelLossConfig):
         desc="Temperature for teacher softmax.",
         valid=check_field(Assert.gt, 0.0),
     )
-    use_triton: bool | None = Field(
-        default=None,
-        desc="Enable triton implementation. Default: use if available.",
-        hint=FieldHint.expert,
-    )
 
     @classmethod
     def _from_dict(cls, default: dict[str, typing.Any], strict: bool = True) -> typing.Self:
@@ -187,16 +200,11 @@ def get_reference_models(self) -> set[str]:
 
 
 @config_class(dynamic_type={LanguageModelLossConfig: "z_loss"})
-class LanguageModelZLossConfig(LanguageModelLossConfig):
+class LanguageModelZLossConfig(CombinableLossConfig):
     """Z-loss regularization to prevent overconfidence."""
 
     _abstract: typing.ClassVar[bool] = False
-
-    use_triton: bool | None = Field(
-        default=None,
-        desc="Enable triton implementation. Default: use if available.",
-        hint=FieldHint.expert,
-    )
+    triton_kind: typing.ClassVar[str | None] = "z"
 
     @property
     def loss_class(self) -> "type[LanguageModelZLoss]":
@@ -208,7 +216,7 @@ def loss_class(self) -> "type[LanguageModelZLoss]":
 class PolicyMetricsLevel(enum.StrEnum):
     none = "none"
     basic = "basic"
-    with_entropy = "with_entropy"
+    auto = "auto"
 
 
 @config_class()
@@ -220,12 +228,13 @@ class LanguageModelPolicyGradientLossConfig(LanguageModelLossConfig):
     epsilon_low: float = Field(default=0.2, desc="Lower clip parameter for ratio of log probs")
     epsilon_high: float = Field(default=0.2, desc="Upper clip parameter for ratio of log probs")
     metrics: PolicyMetricsLevel = Field(
-        default=PolicyMetricsLevel.none,
+        default=PolicyMetricsLevel.auto,
         desc=(
-            "Additional diagnostic metrics to log. "
-            "`basic`: importance-ratio, KL and advantage statistics. "
-            "`with_entropy`: also log the policy entropy. "
-            "Not supported with pipeline_parallel > 1."
+            "Diagnostic metrics to log. "
+            "`basic`: importance-ratio, KL, advantage statistics and the policy entropy. "
+            "`auto`: `basic` when `pipeline_parallel == 1`, else `none`. "
+            "`none`: disable, adding no cost. "
+            "`basic` is not supported with `pipeline_parallel > 1`."
         ),
         hint=FieldHint.feature,
     )
@@ -236,16 +245,11 @@ def loss_class(self) -> "type[LanguageModelPolicyGradientLoss]":
 
 
 @config_class(dynamic_type={LanguageModelLossConfig: "grpo"})
-class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig):
+class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig, CombinableLossConfig):
     """Group-Relative Policy Optimization: per-token IS-ratio clipping."""
 
     _abstract: typing.ClassVar[bool] = False
-
-    use_triton: bool | None = Field(
-        default=None,
-        desc="Enable triton implementation. Default: use if available.",
-        hint=FieldHint.expert,
-    )
+    triton_kind: typing.ClassVar[str | None] = "grpo"
 
     @property
     def loss_class(self) -> "type[LanguageModelGRPOLoss]":
@@ -255,19 +259,65 @@ def loss_class(self) -> "type[LanguageModelGRPOLoss]":
 
 
 @config_class(dynamic_type={LanguageModelLossConfig: "gspo"})
-class LanguageModelGSPOLossConfig(LanguageModelPolicyGradientLossConfig):
+class LanguageModelGSPOLossConfig(LanguageModelPolicyGradientLossConfig, CombinableLossConfig):
     """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]":
+        from fast_llm.layers.language_model.loss.policy_gradient import LanguageModelGSPOLoss
+
+        return LanguageModelGSPOLoss
+
 
+@config_class(dynamic_type={LanguageModelLossConfig: "monolithic"})
+class MonolithicLossConfig(LanguageModelLossConfig):
+    """A composite loss that runs one vocabulary softmax and shares it across its combinable child losses."""
+
+    _abstract: typing.ClassVar[bool] = False
+
+    losses: dict[str, LanguageModelLossConfig] = Field(
+        default_factory=dict,
+        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 triton implementation. Default: use if available.",
+        desc="Enable the triton implementation of the shared-softmax kernel. Default: use if available.",
         hint=FieldHint.expert,
     )
 
+    def _validate(self) -> None:
+        super()._validate()
+        Assert.gt(len(self.losses), 0)
+        for name, loss in self.losses.items():
+            if not isinstance(loss, CombinableLossConfig):
+                raise ValueError(
+                    f"Loss `{name}` (`{type(loss).__name__}`) is not combinable and cannot share the softmax."
+                )
+            if loss.use_triton is not 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[LanguageModelGSPOLoss]":
-        from fast_llm.layers.language_model.loss.policy_gradient import LanguageModelGSPOLoss
+    def loss_class(self) -> "type[LanguageModelLoss]":
+        from fast_llm.layers.language_model.loss.monolithic import MonolithicLoss
 
-        return LanguageModelGSPOLoss
+        return MonolithicLoss
+
+    def get_reference_models(self) -> set[str]:
+        return {reference_model for loss in self.losses.values() for reference_model in loss.get_reference_models()}
diff --git a/fast_llm/layers/language_model/loss/dpo.py b/fast_llm/layers/language_model/loss/dpo.py
index 059f808e5..644745b64 100644
--- a/fast_llm/layers/language_model/loss/dpo.py
+++ b/fast_llm/layers/language_model/loss/dpo.py
@@ -3,10 +3,10 @@
 import torch
 
 from fast_llm.layers.language_model.loss.config import LanguageModelDPOLossConfig, LanguageModelLossKwargs
-from fast_llm.layers.language_model.loss.loss import LanguageModelLoss, loss_forward_backward
+from fast_llm.layers.language_model.loss.loss import SingleLoss, loss_forward_backward
 
 
-class LanguageModelDPOLoss[ConfigType: LanguageModelDPOLossConfig](LanguageModelLoss[ConfigType]):
+class LanguageModelDPOLoss[ConfigType: LanguageModelDPOLossConfig](SingleLoss[ConfigType]):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if self._prediction_distance > 1:
diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py
index 48c1556f3..12cfa7bcf 100644
--- a/fast_llm/layers/language_model/loss/entropy_loss.py
+++ b/fast_llm/layers/language_model/loss/entropy_loss.py
@@ -2,17 +2,27 @@
 
 import torch
 
-from fast_llm.functional.config import TargetFormat, TritonConfig
-from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward
+from fast_llm.functional.config import EntropyLossType, TargetFormat, TritonConfig
+from fast_llm.functional.entropy_loss import (
+    cross_entropy_from_distribution_core,
+    cross_entropy_from_labels_core,
+    reverse_kl_from_distribution_core,
+)
 from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward
+from fast_llm.functional.utils import reduce_losses
 from fast_llm.layers.language_model.loss.config import (
     LanguageModelDistillationLossConfig,
     LanguageModelLabelEntropyLossConfig,
 )
-from fast_llm.layers.language_model.loss.loss import LanguageModelLoss
+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](LanguageModelLoss[ConfigType]):
+class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](
+    CombinableLoss, SingleLoss[ConfigType]
+):
     def _forward_backward(
         self,
         logits: "torch.Tensor",
@@ -21,25 +31,71 @@ def _forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> "tuple[torch.Tensor, torch.Tensor | None]":
-        return (
-            triton_entropy_loss_forward_backward
-            if TritonConfig.enabled(logits.device, self._config.use_triton)
-            else fused_entropy_loss_forward_backward
-        )(
-            logits,
-            self._get_labels(kwargs, split_index),
-            None,  # Labels are already masked
-            grad_logits=grad_logits,
-            grad_output=self._get_grad_output(kwargs),
-            group=self._parallel_dim.group if self._vocab_parallel else None,
-            logits_scale_factor=self._logits_scale_factor,
-            target_format=TargetFormat.labels,
-            entropy_loss_type=self._config.loss_type,
-            divisor=self._get_label_count(kwargs),
+        arguments = self.get_inputs(kwargs, split_index, losses is not None)
+        group = self._parallel_dim.group if self._vocab_parallel else None
+        if TritonConfig.enabled(logits.device, self._config.use_triton):
+            target, grad_output, divisor = arguments
+            return triton_entropy_loss_forward_backward(
+                logits,
+                target,
+                None,  # Labels are already masked
+                grad_logits=grad_logits,
+                grad_output=grad_output,
+                group=group,
+                logits_scale_factor=self._logits_scale_factor,
+                target_format=TargetFormat.labels,
+                entropy_loss_type=self._config.loss_type,
+                divisor=divisor,
+            )
+        loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments)
+        return loss, grad_logits
+
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
+        return self._get_labels(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs)
+
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, None]:
+        """Post-softmax cross-entropy-from-labels over the shared softmax. Returns the loss scalar and the
+        uncast, masked gradient (the caller casts); no extra outputs. For labels, forward-KL is identical to
+        cross-entropy (one-hot target entropy is zero)."""
+        target, grad_output, divisor = arguments
+        loss_mask = target >= 0
+        grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
+        per_sample_loss, grad = cross_entropy_from_labels_core(
+            logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group
         )
+        loss = reduce_losses(per_sample_loss, divisor, loss_mask)
+        if grad is not None:
+            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,)
 
-class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig](LanguageModelLoss[ConfigType]):
+    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]
+):
     def _forward_backward(
         self,
         logits: "torch.Tensor",
@@ -48,22 +104,80 @@ def _forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> "tuple[torch.Tensor, torch.Tensor | None]":
+        arguments = self.get_inputs(kwargs, split_index, losses is not None)
+        group = self._parallel_dim.group if self._vocab_parallel else None
+        if TritonConfig.enabled(logits.device, self._config.use_triton):
+            target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments
+            return triton_entropy_loss_forward_backward(
+                logits,
+                target,
+                loss_mask,
+                grad_output=grad_output,
+                grad_logits=grad_logits,
+                group=group,
+                logits_scale_factor=self._logits_scale_factor,
+                temperature=temperature,
+                target_format=TargetFormat.logits,
+                entropy_loss_type=entropy_loss_type,
+                divisor=divisor,
+            )
+        loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments)
+        return loss, grad_logits
+
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
         return (
-            triton_entropy_loss_forward_backward
-            if TritonConfig.enabled(logits.device, self._config.use_triton)
-            else fused_entropy_loss_forward_backward
-        )(
-            logits,
             self._get_reference_model_logits(self._config.reference_model, kwargs, split_index),
             self._get_loss_mask(kwargs, split_index),
-            grad_output=self._get_grad_output(kwargs),
-            grad_logits=grad_logits,
-            group=self._parallel_dim.group if self._vocab_parallel else None,
-            logits_scale_factor=self._logits_scale_factor,
-            target_format=TargetFormat.logits,
-            entropy_loss_type=self._config.loss_type,
-            divisor=self._get_label_count(kwargs),
+            self._get_grad_output(kwargs),
+            self._get_label_count(kwargs),
+            self._config.loss_type,
+            self._config.temperature,
         )
 
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, None]:
+        """Post-softmax distillation over the shared student softmax (cross-entropy / forward-KL / reverse-KL
+        from a teacher distribution, adding a teacher softmax at scale `logits_scale_factor / temperature`).
+        Returns the loss scalar and the uncast, masked gradient (the caller casts); no extra outputs."""
+        target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments
+        grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
+        if entropy_loss_type == EntropyLossType.reverse_kl:
+            per_sample_loss, grad = reverse_kl_from_distribution_core(
+                logits_norm,
+                exp_logits,
+                sum_exp_logits,
+                target,
+                grad_output,
+                logits_scale_factor,
+                TargetFormat.logits,
+                group,
+                temperature,
+            )
+        else:
+            per_sample_loss, grad = cross_entropy_from_distribution_core(
+                logits_norm,
+                exp_logits,
+                sum_exp_logits,
+                target,
+                grad_output,
+                logits_scale_factor,
+                TargetFormat.logits,
+                group,
+                temperature,
+                return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl,
+            )
+        loss = reduce_losses(per_sample_loss, divisor, loss_mask)
+        if grad is not None and loss_mask is not None:
+            grad = grad * loss_mask.unsqueeze(-1)
+        return loss, grad, None
+
     def get_preprocessing_config(self) -> dict[str, typing.Any]:
         return {"return_prediction_mask": True}
diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py
index 397aec966..88a99ce89 100644
--- a/fast_llm/layers/language_model/loss/loss.py
+++ b/fast_llm/layers/language_model/loss/loss.py
@@ -7,10 +7,14 @@
 from fast_llm.core.ops import split_op
 from fast_llm.engine.base_model.config import LossDef
 from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames
+from fast_llm.functional.entropy_loss import softmax_base
 from fast_llm.layers.language_model.config import LanguageModelKwargs
 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__(
@@ -42,6 +46,7 @@ def __init__(
         self._sequence_data_dim = distributed_config.get_distributed_dim(DistributedDimNames.sequence_data)
         self._sequence_data_active = self._sequence_data_dim.size > 1
 
+    @abc.abstractmethod
     def forward_backward(
         self,
         logits: "torch.Tensor",
@@ -50,23 +55,6 @@ def forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> "tuple[torch.Tensor | None, torch.Tensor | None]":
-        if losses is None and self._weight is None:
-            # Loss computation is not needed, skip.
-            return None, None
-        loss, grad = self._forward_backward(logits, kwargs, losses, split_index, grad_logits)
-        if self._do_register_loss:
-            self._register_loss(self.name, loss, losses)
-        return loss if self._weight == 1 else loss * self._weight, grad
-
-    @abc.abstractmethod
-    def _forward_backward(
-        self,
-        logits: "torch.Tensor",
-        kwargs: dict[str, typing.Any],
-        losses: dict | None = None,
-        split_index: int = 0,
-        grad_logits: torch.Tensor | None = None,
-    ) -> "tuple[torch.Tensor, torch.Tensor | None]":
         pass
 
     def get_loss_definitions(self) -> list[LossDef]:
@@ -117,7 +105,8 @@ def _prepare_target(
 
     def _get_grad_output(self, kwargs: dict[str, typing.Any]) -> float | None:
         grad_output = kwargs.get(LanguageModelKwargs.grad_output)
-        return None if grad_output is None else grad_output * self._weight
+        # A zero-weight loss contributes an all-zero gradient; return `None` so the backward term drops out.
+        return None if grad_output is None or self._weight == 0 else grad_output * self._weight
 
     def _get_labels(self, kwargs: dict[str, typing.Any], split_index: int = 0):
         return self._prepare_target(kwargs[LanguageModelLossKwargs.labels], split_index)
@@ -130,8 +119,10 @@ def _get_loss_mask(self, kwargs: dict[str, typing.Any], split_index: int = 0):
         return None if loss_mask is None else self._prepare_target(loss_mask, split_index)
 
     def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, typing.Any], split_index: int = 0):
+        # The head owns the `.logits`; split on `.losses.` so the name resolves whether this loss sits
+        # directly under the head or nested inside a composite loss (e.g. `MonolithicLoss`).
         Assert.incl(
-            logits_name := self.module_name.rsplit(".", 2)[0] + f".logits",
+            logits_name := self.module_name.split(".losses.")[0] + ".logits",
             reference_hidden_states := kwargs[f"reference_{reference_model}_hidden_states"],
         )
         # The logits are already sequence-parallel if needed, we don't want to split again.
@@ -140,6 +131,146 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty
         )
 
 
+class SingleLoss[ConfigType: LanguageModelLossConfig](LanguageModelLoss[ConfigType]):
+    """A loss emitting a single registered, weighted scalar. Subclasses implement `_forward_backward` (the
+    loss math and its gradient); this template skips the disabled case, registers the scalar, and applies the
+    per-loss weight. Composite losses satisfy `forward_backward` directly rather than through this template."""
+
+    def forward_backward(
+        self,
+        logits: "torch.Tensor",
+        kwargs: dict[str, typing.Any],
+        losses: dict | None = None,
+        split_index: int = 0,
+        grad_logits: torch.Tensor | None = None,
+    ) -> "tuple[torch.Tensor | None, torch.Tensor | None]":
+        if losses is None and self._weight is None:
+            # Loss computation is not needed, skip.
+            return None, None
+        loss, grad = self._forward_backward(logits, kwargs, losses, split_index, grad_logits)
+        if self._do_register_loss:
+            self._register_loss(self.name, loss, losses)
+        return loss if self._weight == 1 else loss * self._weight, grad
+
+    @abc.abstractmethod
+    def _forward_backward(
+        self,
+        logits: "torch.Tensor",
+        kwargs: dict[str, typing.Any],
+        losses: dict | None = None,
+        split_index: int = 0,
+        grad_logits: torch.Tensor | None = None,
+    ) -> "tuple[torch.Tensor, torch.Tensor | None]":
+        pass
+
+
+class CombinableLoss:
+    """Mixin for losses that consume the vocabulary softmax, either standalone through
+    `combinable_forward_backward` or several sharing one softmax when fused together. Subclasses implement
+    `get_inputs` (eager kwargs -> argument tuple, built outside the compiled boundary) and the `fused_core`
+    static method (post-softmax math returning `(loss, uncast_grad, extra)`), and override
+    `register_combinable_extras` when they emit outputs beyond the loss scalar. A loss whose gradient can't
+    be produced inside the compiled boundary (an eager seam between forward and backward) instead has its
+    `fused_core` return `(None, None, forward_state)` and completes its loss/gradient in `finish`."""
+
+    _logits_scale_factor: float
+
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
+        raise NotImplementedError()
+
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, typing.Any]:
+        raise NotImplementedError()
+
+    @torch.compile
+    def combinable_forward_backward(
+        self,
+        logits: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        grad_logits: torch.Tensor | None,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, typing.Any]:
+        """Standalone realization of a single combinable loss: softmax once, this loss's `fused_core`, then
+        cast-and-accumulate the gradient. Shares `fused_core` with the fused path, so it is not a second copy
+        of the math."""
+        logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, self._logits_scale_factor, group)
+        loss, grad, extra = self.fused_core(
+            logits_norm, exp_logits, sum_exp_logits, logits_max, group, self._logits_scale_factor, arguments
+        )
+        return loss, self._accumulate_grad(grad, logits.dtype, grad_logits), extra
+
+    @staticmethod
+    def _accumulate_grad(
+        grad: torch.Tensor | None, logits_dtype: torch.dtype, grad_logits: torch.Tensor | None
+    ) -> torch.Tensor | None:
+        """Cast a computed logits gradient to the logits dtype and accumulate it into `grad_logits` (in place
+        when a buffer exists, otherwise adopting it)."""
+        if grad is None:
+            return grad_logits
+        grad = grad.to(logits_dtype)
+        if grad_logits is None:
+            return grad
+        grad_logits.add_(grad)
+        return grad_logits
+
+    def register_combinable_extras(
+        self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None
+    ) -> None:
+        """Register per-loss outputs beyond the loss scalar (produced by `fused_core`). No-op by default."""
+
+    def finish(
+        self,
+        loss: torch.Tensor | None,
+        extra: typing.Any,
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+        grad_logits: torch.Tensor | None,
+        logits_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, typing.Any, torch.Tensor | None]:
+        """Complete a loss that couldn't finish inside the compiled shared-softmax boundary, accumulating its
+        gradient into `grad_logits` and returning `(loss, extra, grad_logits)`. A passthrough by default; a
+        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
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py
new file mode 100644
index 000000000..b774bee99
--- /dev/null
+++ b/fast_llm/layers/language_model/loss/monolithic.py
@@ -0,0 +1,233 @@
+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, ...],
+    logits: torch.Tensor,  # (*batch, vocab)
+    group: ProcessGroup | None,
+    logits_scale_factor: float,
+    grad_logits: torch.Tensor | None,
+    arguments: tuple[tuple, ...],
+) -> tuple[list[tuple[torch.Tensor, typing.Any]], torch.Tensor | None]:
+    """
+    One shared softmax over the logits, then each child loss's `fused_core` consuming it. The child
+    list is fixed per config, so the loop unrolls inside this single `@torch.compile` boundary and each
+    `fused_core` dispatches (and inlines) to its loss type's math — every enabled loss is fused over
+    one softmax. Gradient contributions accumulate in fp32 and cast to `logits.dtype` once at the end.
+    """
+    logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group)
+    grad = None
+    results = []
+    for child, child_arguments in zip(children, arguments, strict=True):
+        loss, child_grad, extra = child.fused_core(
+            logits_norm, exp_logits, sum_exp_logits, logits_max, group, logits_scale_factor, child_arguments
+        )
+        results.append((loss, extra))
+        if child_grad is not None:
+            grad = child_grad if grad is None else grad + child_grad
+    return results, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits)
+
+
+class MonolithicLoss[ConfigType: MonolithicLossConfig](LanguageModelLoss[ConfigType]):
+    """
+    A composite loss that runs the vocabulary softmax once and shares it across its combinable child losses,
+    emitting each child's scalar / metrics and the combined logits gradient in a single `@torch.compile`
+    boundary. It is an ordinary head loss: the head loops over it like any other and threads the same gradient
+    buffer, so non-combinable losses remain plain siblings in the head's loss list.
+    """
+
+    def __init__(
+        self,
+        config: ConfigType,
+        distributed_config: DistributedConfig,
+        *,
+        name: str,
+        prediction_distance: int = 1,
+        prediction_heads: int = 1,
+        vocab_parallel: bool = False,
+        num_splits: int = 1,
+        logits_scale_factor: float = 1.0,
+        weight: float = 1.0,
+        register_loss: bool = False,
+    ):
+        super().__init__(
+            config,
+            distributed_config,
+            name=name,
+            prediction_distance=prediction_distance,
+            prediction_heads=prediction_heads,
+            vocab_parallel=vocab_parallel,
+            num_splits=num_splits,
+            logits_scale_factor=logits_scale_factor,
+            weight=weight,
+            register_loss=register_loss,
+        )
+        # Register children as distinct losses, unless a single child equals the head total (logged anyway).
+        children_register = register_loss or len(config.losses) > 1
+        self._children: typing.Sequence[LanguageModelLoss] = torch.nn.ModuleList(
+            [
+                child_config.get_layer(
+                    distributed_config,
+                    name=child_name if prediction_distance == 1 else f"{child_name}_{prediction_distance}",
+                    prediction_distance=prediction_distance,
+                    prediction_heads=prediction_heads,
+                    vocab_parallel=vocab_parallel,
+                    num_splits=num_splits,
+                    logits_scale_factor=self._logits_scale_factor,
+                    weight=self._weight,
+                    register_loss=children_register,
+                )
+                for child_name, child_config in config.losses.items()
+            ]
+        )
+        # 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,
+        logits: "torch.Tensor",
+        kwargs: dict[str, typing.Any],
+        losses: dict | None = None,
+        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
+        results, grad_logits = _monolithic_core(
+            tuple(self._children), logits, group, self._softmax_scale_factor, grad_logits, arguments
+        )
+
+        total_loss = None
+        for child, (loss, extra) in zip(self._children, results, strict=True):
+            # A child whose gradient can't be produced inside the compiled boundary (an eager seam) had
+            # `fused_core` return `(None, None, forward_state)`; `finish` completes its loss/gradient here.
+            loss, extra, grad_logits = child.finish(loss, extra, kwargs, split_index, grad_logits, logits.dtype)
+            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 _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))
+
+    def get_loss_definitions(self) -> list[LossDef]:
+        return [loss_def for child in self._children for loss_def in child.get_loss_definitions()]
diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py
index ce532f00d..56f8f048b 100644
--- a/fast_llm/layers/language_model/loss/policy_gradient.py
+++ b/fast_llm/layers/language_model/loss/policy_gradient.py
@@ -7,7 +7,12 @@
 from fast_llm.engine.base_model.config import LossDef, ReductionType
 from fast_llm.engine.distributed.config import DistributedConfig
 from fast_llm.functional.config import TritonConfig
-from fast_llm.functional.entropy_loss import fused_predicted_logits_from_labels, fused_softmax_base
+from fast_llm.functional.entropy_loss import (
+    fused_predicted_logits_from_labels,
+    fused_softmax_base,
+    predicted_logits_from_labels,
+    softmax_base,
+)
 from fast_llm.functional.utils import reduce_losses
 from fast_llm.layers.block.config import BlockKwargs
 from fast_llm.layers.language_model.config import LanguageModelKwargs
@@ -18,9 +23,12 @@
     LanguageModelPolicyGradientLossConfig,
     PolicyMetricsLevel,
 )
-from fast_llm.layers.language_model.loss.loss import LanguageModelLoss
+from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss
 from fast_llm.utils import Assert
 
+if typing.TYPE_CHECKING:
+    from fast_llm.layers.language_model.loss.monolithic import _TritonContext
+
 
 class PolicyMetrics(typing.NamedTuple):
     # Weighted sums for the policy-gradient diagnostics, shared by the per-token (GRPO) and per-segment
@@ -40,12 +48,10 @@ class PolicyMetrics(typing.NamedTuple):
     min_advantage: torch.Tensor
     num_tokens: torch.Tensor
     num_segments: torch.Tensor | None
-    entropy: torch.Tensor | None
+    entropy: torch.Tensor
 
 
-class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig](
-    LanguageModelLoss[ConfigType]
-):
+class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig](SingleLoss[ConfigType]):
     """Shared scaffolding for policy-gradient losses (GRPO, GSPO)."""
 
     def __init__(
@@ -74,12 +80,19 @@ def __init__(
             weight=weight,
             register_loss=register_loss,
         )
-        # The extra metrics need a second softmax over the full logits, which pipeline parallelism splits.
-        Assert.custom(
-            lambda metrics, pipeline_parallel: metrics == PolicyMetricsLevel.none or pipeline_parallel == 1,
-            config.metrics,
-            distributed_config.pipeline_parallel,
-        )
+        # The metrics need the full-logits pass, which pipeline parallelism splits, so they require
+        # `pipeline_parallel == 1`. `auto` enables them only when that holds; an explicit level must satisfy it.
+        if config.metrics == PolicyMetricsLevel.auto:
+            self._metrics_level = (
+                PolicyMetricsLevel.basic if distributed_config.pipeline_parallel == 1 else PolicyMetricsLevel.none
+            )
+        else:
+            Assert.custom(
+                lambda metrics, pipeline_parallel: metrics == PolicyMetricsLevel.none or pipeline_parallel == 1,
+                config.metrics,
+                distributed_config.pipeline_parallel,
+            )
+            self._metrics_level = config.metrics
 
     def _register_new_logprobs(
         self,
@@ -94,9 +107,9 @@ def _register_new_logprobs(
         )
 
     def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]:
-        if self._config.metrics == PolicyMetricsLevel.none:
+        if self._metrics_level == PolicyMetricsLevel.none:
             return []
-        defs = [
+        return [
             LossDef(f"{self._name}_old_logprobs"),
             LossDef(f"{self._name}_ratio_new_old"),
             LossDef(f"{self._name}_ratio_new_old_sum"),
@@ -108,12 +121,12 @@ def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]:
             LossDef(f"{self._name}_min_advantage", reduction=ReductionType.minimum),
             *extra,
             LossDef(f"{self._name}_num_tokens"),
+            LossDef(f"{self._name}_entropy"),
         ]
-        if self._config.metrics == PolicyMetricsLevel.with_entropy:
-            defs.append(LossDef(f"{self._name}_entropy"))
-        return defs
 
-    def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typing.Any], losses: dict) -> None:
+    def _register_policy_metrics(
+        self, metrics: "PolicyMetrics", kwargs: dict[str, typing.Any], losses: dict | None
+    ) -> None:
         num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch]
         for name in ("old_logprobs", "ratio_new_old", "kl_new_old", "clipped_ratio_fraction", "advantage"):
             self._register_loss(f"{self._name}_{name}", getattr(metrics, name) / num_documents, losses)
@@ -127,8 +140,7 @@ def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typ
         )
         if metrics.num_segments is not None:
             self._register_loss(f"{self._name}_num_segments", metrics.num_segments, losses)
-        if metrics.entropy is not None:
-            self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses)
+        self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses)
 
     def get_loss_definitions(self) -> list[LossDef]:
         defs = super().get_loss_definitions()
@@ -143,7 +155,9 @@ def _logprob_metric_name(self) -> str:
         return f"{self._name}_new_logprobs"
 
 
-class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]):
+class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig](
+    CombinableLoss, LanguageModelPolicyGradientLoss[ConfigType]
+):
     """GRPO: per-token IS-ratio clipping."""
 
     def _forward_backward(
@@ -154,44 +168,161 @@ def _forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        arguments = self.get_inputs(kwargs, split_index, losses is not None)
+        group = self._parallel_dim.group if self._vocab_parallel else None
         if TritonConfig.enabled(logits.device, self._config.use_triton):
             from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward
 
-            fn = triton_grpo_loss_forward_backward
-        else:
-            fn = fused_grpo_loss_forward_backward
-        loss, grad, new_logprobs_mean = fn(
-            logits,
+            (
+                target,
+                advantages,
+                old_log_probabilities,
+                grad_output,
+                divisor,
+                epsilon_low,
+                epsilon_high,
+                num_labels_in_seq,
+                compute_metrics,
+            ) = arguments
+            loss, grad, new_logprobs_mean = triton_grpo_loss_forward_backward(
+                logits,
+                target,
+                advantages,
+                old_log_probabilities,
+                grad_logits=grad_logits,
+                grad_output=grad_output,
+                group=group,
+                epsilon_low=epsilon_low,
+                epsilon_high=epsilon_high,
+                logits_scale_factor=self._logits_scale_factor,
+                num_labels_in_seq=num_labels_in_seq,
+                divisor=divisor,
+            )
+            self._register_new_logprobs(new_logprobs_mean, kwargs, losses)
+            # Triton produces only loss/grad/new_logprobs; the metric family needs its own pass here.
+            if compute_metrics:
+                self._register_extra_metrics(logits, kwargs, losses, split_index)
+            return loss, grad
+
+        loss, grad_logits, extra = self.combinable_forward_backward(logits, group, grad_logits, arguments)
+        self.register_combinable_extras(extra, kwargs, losses)
+        return loss, grad_logits
+
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
+        # When nothing is logged this step, drop the metric-only outputs (`new_logprobs_mean` and the
+        # policy metric family) by passing `num_labels_in_seq=None` / `compute_metrics=False`.
+        return (
             self._get_labels(kwargs, split_index),
             self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index),
             self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index),
-            grad_logits=grad_logits,
-            grad_output=self._get_grad_output(kwargs),
-            group=self._parallel_dim.group if self._vocab_parallel else None,
-            epsilon_low=self._config.epsilon_low,
-            epsilon_high=self._config.epsilon_high,
-            logits_scale_factor=self._logits_scale_factor,
-            num_labels_in_seq=(
-                None
-                if losses is None
-                else self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index)
-            ),
-            divisor=self._get_label_count(kwargs),
+            self._get_grad_output(kwargs),
+            self._get_label_count(kwargs),
+            self._config.epsilon_low,
+            self._config.epsilon_high,
+            (self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index) if register else None),
+            register and self._metrics_level != PolicyMetricsLevel.none,
         )
 
-        self._register_new_logprobs(new_logprobs_mean, kwargs, losses)
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple]:
+        """Post-softmax GRPO over the shared softmax. Returns the loss scalar, the uncast masked gradient (the
+        caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not requested) — all
+        from one softmax and one predicted-logit lookup."""
+        (
+            target,
+            advantages,
+            old_log_probabilities,
+            grad_output,
+            divisor,
+            epsilon_low,
+            epsilon_high,
+            num_labels_in_seq,
+            compute_metrics,
+        ) = arguments
+        loss_mask = target >= 0
+        predicted_logits, target_masked, target_mask = predicted_logits_from_labels(
+            logits_norm, target, loss_mask, group
+        )
+        new_log_probs = predicted_logits - sum_exp_logits.log()
+        probability_ratio = (new_log_probs - old_log_probabilities).exp()
 
-        # Skip the extra softmax pass when there is nothing to register.
-        if losses is not None and self._config.metrics != PolicyMetricsLevel.none:
-            self._register_extra_metrics(logits, kwargs, losses, split_index)
+        losses = -torch.min(
+            probability_ratio * advantages,
+            torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantages,
+        )
+        loss = reduce_losses(losses, divisor, loss_mask)
+
+        # Sum of per-sequence mean log-probs, matching pipelinerl's new_logprobs metric:
+        #   sum_sum(new_logprobs / num_labels_in_seq, masks_shifted, segments)
+        # Dividing by num_labels_in_seq (span length broadcast per token) and summing over masked
+        # tokens gives mean logprob per sequence; summing those across sequences matches the deepspeed
+        # convention exactly (segments are redundant once num_labels_in_seq is correct).
+        # Clamp to avoid 0/0=nan when num_labels_in_seq=0 (padded tokens or fully masked documents)
+        # — those positions also have loss_mask=0 so they correctly contribute 0 to the sum.
+        new_logprobs_mean = (
+            None if num_labels_in_seq is None else (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum()
+        )
 
-        return loss, grad
+        metrics = (
+            grpo_metrics_core(
+                new_log_probs,
+                advantages,
+                old_log_probabilities,
+                loss_mask,
+                num_labels_in_seq,
+                epsilon_low,
+                epsilon_high,
+                _policy_entropy_per_token(logits_norm, exp_logits, sum_exp_logits, group),
+            )
+            if compute_metrics
+            else None
+        )
+
+        if grad_output is None:
+            grad = None
+        else:
+            grad_output = grad_output / divisor * logits_scale_factor
+            # loss[a>=0] = -a * min(x, 1 + epsilon_high)  =>  grad[a>=0] = -a * (x <= 1 + epsilon_high)
+            # loss[a<=0] = a * max(x, 1 - epsilon_low)  =>  grad[a<=0] = a * (x >= 1 - epsilon_low)
+            probability_ratio_grad = (
+                grad_output
+                * (
+                    torch.clamp_min(advantages, 0) * (probability_ratio <= 1 + epsilon_high)
+                    + torch.clamp_max(advantages, 0) * (probability_ratio >= 1 - epsilon_low)
+                )
+                * loss_mask
+            )
+            # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities)
+            # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` so the shared softmax tensors
+            # are not mutated in place, since other losses may reuse them over the same softmax.
+            predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1)
+            grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add(
+                -1,
+                target_masked.unsqueeze(-1),
+                -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32),
+            )
+
+        return loss, grad, (new_logprobs_mean, metrics)
+
+    def register_combinable_extras(self, extra: tuple, kwargs: dict[str, typing.Any], losses: dict | None) -> None:
+        new_logprobs_mean, metrics = extra
+        self._register_new_logprobs(new_logprobs_mean, kwargs, losses)
+        if metrics is not None:
+            self._register_policy_metrics(metrics, kwargs, losses)
 
     def _register_extra_metrics(
         self,
         logits: torch.Tensor,
         kwargs: dict[str, typing.Any],
-        losses: dict,
+        losses: dict | None,
         split_index: int,
     ) -> None:
         metrics = compute_grpo_metrics(
@@ -204,16 +335,79 @@ def _register_extra_metrics(
             self._config.epsilon_high,
             self._logits_scale_factor,
             group=self._parallel_dim.group if self._vocab_parallel else None,
-            compute_entropy=self._config.metrics == PolicyMetricsLevel.with_entropy,
         )
         self._register_policy_metrics(metrics, kwargs, losses)
 
+    def triton_metrics(
+        self,
+        new_log_probs: torch.Tensor,  # flat, from the kernel's shared softmax
+        entropy_per_token: torch.Tensor,  # flat, from the kernel's `Σ exp·logits_norm`
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+    ) -> PolicyMetrics:
+        """GRPO metric family from the triton kernel's shared-softmax outputs, reusing `grpo_metrics_core` so
+        the metrics add no second softmax."""
+        target = self._get_labels(kwargs, split_index)
+        return grpo_metrics_core(
+            new_log_probs.reshape(target.shape),
+            self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index),
+            self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index),
+            target >= 0,
+            self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index),
+            self._config.epsilon_low,
+            self._config.epsilon_high,
+            entropy_per_token.reshape(target.shape),
+        )
+
+    def triton_add_inputs(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> None:
+        (
+            target,
+            advantages,
+            old_log_probabilities,
+            grad_output,
+            divisor,
+            epsilon_low,
+            epsilon_high,
+            num_labels_in_seq,
+            _,
+        ) = self.get_inputs(kwargs, split_index, register)
+        if context.labels is None:
+            context.labels = target
+        if context.divisor is None:
+            context.divisor = divisor
+        context.grpo = (advantages, old_log_probabilities, grad_output, epsilon_low, epsilon_high, num_labels_in_seq)
+
+    def triton_metrics_enabled(self, register: bool) -> bool:
+        return register and self._metrics_level != PolicyMetricsLevel.none
+
+    def triton_finish(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> tuple[torch.Tensor, tuple]:
+        metrics = (
+            self.triton_metrics(context.new_log_probs, context.entropy_per_token, kwargs, split_index)
+            if self.triton_metrics_enabled(register)
+            else None
+        )
+        return context.grpo_loss, (context.grpo_new_logprobs, metrics)
+
     def get_loss_definitions(self) -> list[LossDef]:
         return super().get_loss_definitions() + self._policy_metric_definitions()
 
 
-class LanguageModelGSPOLoss[ConfigType: LanguageModelGSPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]):
-    """GSPO: sequence-level geometric-mean IS-ratio clipping."""
+class LanguageModelGSPOLoss[ConfigType: LanguageModelGSPOLossConfig](
+    CombinableLoss, LanguageModelPolicyGradientLoss[ConfigType]
+):
+    """GSPO: sequence-level geometric-mean IS-ratio clipping.
+
+    Standalone, `_forward_backward` runs the whole three-phase kernel (`fused_gspo_loss_forward_backward` or
+    its Triton twin). Fused into a shared softmax, `fused_core` runs only the forward on that softmax and the
+    segment seam + backward are deferred to `finish`, since the eager `index_add_` segment aggregation can't
+    run inside the compiled boundary."""
+
+    # GSPO's gradient needs the shared softmax before the kernel, so the driver runs the reduced forward first.
+    triton_needs_forward: typing.ClassVar[bool] = True
 
     def __init__(
         self,
@@ -245,6 +439,15 @@ def __init__(
         # aggregation can't recombine across chunks since each call only sees a slice.
         Assert.eq(self._num_splits, 1)
 
+    def _document_index_zero_based(self, kwargs: dict[str, typing.Any], split_index: int) -> torch.Tensor:
+        # `global_document_index_q` is 1-based per the data preprocessor convention; the kernel takes 0-based.
+        return (
+            self._prepare_target(
+                kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False
+            ).long()
+            - 1
+        )
+
     def _forward_backward(
         self,
         logits: "torch.Tensor",
@@ -253,13 +456,7 @@ def _forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        document_index_zero_based = (
-            self._prepare_target(
-                kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False
-            ).long()
-            - 1
-        )
-        # `global_document_index_q` is 1-based per the data preprocessor convention; the kernel takes 0-based.
+        document_index_zero_based = self._document_index_zero_based(kwargs, split_index)
         # `num_documents_in_sequence` is the doc count for this DP rank's batch — identical across
         # SDP/SP ranks within a DP rank, so per-segment buffers are sized consistently for all-reduce.
         num_segments = kwargs[BlockKwargs.num_documents_in_sequence]
@@ -294,7 +491,7 @@ def _forward_backward(
         self._register_new_logprobs(new_logprobs_mean, kwargs, losses)
 
         # Skip the extra softmax pass when there is nothing to register.
-        if losses is not None and self._config.metrics != PolicyMetricsLevel.none:
+        if losses is not None and self._metrics_level != PolicyMetricsLevel.none:
             self._register_extra_metrics(logits, kwargs, losses, split_index, document_index_zero_based, num_segments)
 
         return loss, grad
@@ -303,7 +500,7 @@ def _register_extra_metrics(
         self,
         logits: torch.Tensor,
         kwargs: dict[str, typing.Any],
-        losses: dict,
+        losses: dict | None,
         split_index: int,
         document_index_zero_based: torch.Tensor,
         num_segments: int,
@@ -322,10 +519,207 @@ def _register_extra_metrics(
             group=self._parallel_dim.group if self._vocab_parallel else None,
             sdp_group=self._sequence_data_dim.group if self._sequence_data_active else None,
             sp_group=self._parallel_dim.group if self._sequence_parallel else None,
-            compute_entropy=self._config.metrics == PolicyMetricsLevel.with_entropy,
         )
         self._register_policy_metrics(metrics, kwargs, losses)
 
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
+        return (self._get_labels(kwargs, split_index), register and self._metrics_level != PolicyMetricsLevel.none)
+
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[None, None, tuple]:
+        """GSPO forward over the shared softmax: the per-token log-probs plus the softmax tensors its seam,
+        backward and (when requested) metrics need. Returns `(None, None, forward_state)` — GSPO's loss and
+        gradient can't be produced here (the segment aggregation is eager), so both are deferred to `finish`."""
+        target, compute_metrics = arguments
+        loss_mask = target >= 0
+        predicted_logits, target_masked, target_mask = predicted_logits_from_labels(
+            logits_norm, target, loss_mask, group
+        )
+        new_log_probs = predicted_logits - sum_exp_logits.log()
+        # Reduce the vocab-wide entropy here, inside the compiled boundary, rather than holding `logits_norm`
+        # (a full-vocab tensor) across the eager seam for `finish` to reduce eagerly.
+        entropy_per_token = (
+            _policy_entropy_per_token(logits_norm, exp_logits, sum_exp_logits, group) if compute_metrics else None
+        )
+        return (
+            None,
+            None,
+            (
+                new_log_probs,
+                loss_mask,
+                exp_logits,
+                sum_exp_logits,
+                target_masked,
+                target_mask,
+                entropy_per_token,
+                compute_metrics,
+            ),
+        )
+
+    def _run_segment_seam(
+        self, new_log_probs: torch.Tensor, loss_mask: torch.Tensor, kwargs: dict[str, typing.Any], split_index: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        """Run the eager segment seam from per-token new log-probs, pulling its remaining inputs from `kwargs`.
+        Returns the loss, the `new_logprobs` metric, and the per-token backward coefficient."""
+        return gspo_segment_seam(
+            new_log_probs,
+            loss_mask,
+            self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index),
+            self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index),
+            self._document_index_zero_based(kwargs, split_index),
+            kwargs[BlockKwargs.num_documents_in_sequence],
+            self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index),
+            kwargs[LanguageModelKwargs.num_documents_in_batch],
+            self._get_grad_output(kwargs),
+            self._sequence_data_dim.group if self._sequence_data_active else None,
+            self._parallel_dim.group if self._sequence_parallel else None,
+            self._config.epsilon_low,
+            self._config.epsilon_high,
+            self._logits_scale_factor,
+        )
+
+    def finish(
+        self,
+        loss: torch.Tensor | None,
+        extra: tuple,
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+        grad_logits: torch.Tensor | None,
+        logits_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, tuple, torch.Tensor | None]:
+        """Run the eager segment seam and the compiled backward over the shared softmax deferred by
+        `fused_core`, accumulating GSPO's gradient into `grad_logits`. Returns the loss and the
+        `(new_logprobs, metrics)` extras (metrics from the same shared softmax, `None` when not requested),
+        both registered by `register_combinable_extras`."""
+        (
+            new_log_probs,
+            loss_mask,
+            exp_logits,
+            sum_exp_logits,
+            target_masked,
+            target_mask,
+            entropy_per_token,
+            compute_metrics,
+        ) = extra
+        loss, new_logprobs_mean, effective_grad = self._run_segment_seam(new_log_probs, loss_mask, kwargs, split_index)
+        if effective_grad is not None:
+            grad_logits = gspo_backward_core(
+                exp_logits,
+                sum_exp_logits,
+                target_masked,
+                target_mask,
+                loss_mask,
+                effective_grad,
+                logits_dtype,
+                grad_logits,
+            )
+        metrics = (
+            gspo_metrics_core(
+                new_log_probs,
+                self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index),
+                self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index),
+                loss_mask,
+                self._document_index_zero_based(kwargs, split_index),
+                kwargs[BlockKwargs.num_documents_in_sequence],
+                self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index),
+                self._config.epsilon_low,
+                self._config.epsilon_high,
+                entropy_per_token,
+                self._sequence_data_dim.group if self._sequence_data_active else None,
+                self._parallel_dim.group if self._sequence_parallel else None,
+            )
+            if compute_metrics
+            else None
+        )
+        return loss, (new_logprobs_mean, metrics), grad_logits
+
+    def compute_triton_seam(
+        self,
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+        max_logits: torch.Tensor,  # (n_rows,)
+        sum_exp_logits: torch.Tensor,  # (n_rows,)
+        predicted_logits: torch.Tensor,  # (n_rows,)
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        """GSPO's contribution to the triton monolithic kernel: recover per-token new log-probs from the
+        triton forward pass, run the segment seam, and return the loss, the `new_logprobs` metric, and the
+        flat per-token backward coefficient (`None` when no gradient is requested) the kernel superposes."""
+        target = self._get_labels(kwargs, split_index)
+        loss_mask = target >= 0
+        new_log_probs = (predicted_logits - max_logits - sum_exp_logits.log()).reshape(loss_mask.shape)
+        loss, new_logprobs_mean, effective_grad = self._run_segment_seam(new_log_probs, loss_mask, kwargs, split_index)
+        return loss, new_logprobs_mean, None if effective_grad is None else effective_grad.reshape(-1).contiguous()
+
+    def triton_metrics(
+        self,
+        new_log_probs: torch.Tensor,  # flat, from the kernel's shared softmax
+        entropy_per_token: torch.Tensor,  # flat, from the kernel's `Σ exp·logits_norm`
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+    ) -> PolicyMetrics:
+        """GSPO segment-level metric family from the triton kernel's shared-softmax outputs, reusing
+        `gspo_metrics_core` so the metrics add no second softmax."""
+        target = self._get_labels(kwargs, split_index)
+        return gspo_metrics_core(
+            new_log_probs.reshape(target.shape),
+            self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index),
+            self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index),
+            target >= 0,
+            self._document_index_zero_based(kwargs, split_index),
+            kwargs[BlockKwargs.num_documents_in_sequence],
+            self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index),
+            self._config.epsilon_low,
+            self._config.epsilon_high,
+            entropy_per_token.reshape(target.shape),
+            self._sequence_data_dim.group if self._sequence_data_active else None,
+            self._parallel_dim.group if self._sequence_parallel else None,
+        )
+
+    def triton_add_inputs(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> None:
+        target, _ = self.get_inputs(kwargs, split_index, register)
+        if context.labels is None:
+            context.labels = target
+
+    def triton_seam(
+        self,
+        context: "_TritonContext",
+        softmax: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None],
+        kwargs: dict[str, typing.Any],
+        split_index: int,
+    ) -> None:
+        context.gspo_loss, context.gspo_new_logprobs, context.gspo_coeff = self.compute_triton_seam(
+            kwargs, split_index, *softmax
+        )
+
+    def triton_metrics_enabled(self, register: bool) -> bool:
+        return register and self._metrics_level != PolicyMetricsLevel.none
+
+    def triton_finish(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> tuple[torch.Tensor, tuple]:
+        metrics = (
+            self.triton_metrics(context.new_log_probs, context.entropy_per_token, kwargs, split_index)
+            if self.triton_metrics_enabled(register)
+            else None
+        )
+        return context.gspo_loss, (context.gspo_new_logprobs, metrics)
+
+    def register_combinable_extras(self, extra: tuple, kwargs: dict[str, typing.Any], losses: dict | None) -> None:
+        new_logprobs_mean, metrics = extra
+        self._register_new_logprobs(new_logprobs_mean, kwargs, losses)
+        if metrics is not None:
+            self._register_policy_metrics(metrics, kwargs, losses)
+
     def get_loss_definitions(self) -> list[LossDef]:
         return super().get_loss_definitions() + self._policy_metric_definitions(LossDef(f"{self._name}_num_segments"))
 
@@ -333,32 +727,6 @@ def get_preprocessing_config(self) -> dict[str, typing.Any]:
         return super().get_preprocessing_config() | {"return_document_index": True}
 
 
-def _policy_metrics_log_ratio(
-    logits: torch.Tensor,
-    target: torch.Tensor,
-    old_log_probabilities: torch.Tensor,
-    logits_scale_factor: float,
-    group: torch.distributed.ProcessGroup | None,
-    compute_entropy: bool,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-    """Shared softmax pass for both policy-gradient metric paths. Returns the loss mask, the per-token
-    new/old log-ratio, and (when requested) the per-token policy entropy. `exp_logits` / `logits_norm`
-    are local vocab slices, so the entropy's weighted-logit sum is all-reduced across the tensor-parallel
-    group to recover the global expectation before dividing by the already-global `sum_exp_logits`."""
-    loss_mask = target >= 0
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
-    predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group)
-    log_ratio = predicted_logits - sum_exp_logits.log() - old_log_probabilities
-
-    entropy_per_token: torch.Tensor | None = None
-    if compute_entropy:
-        weighted_logits_sum = (exp_logits * logits_norm).sum(-1)
-        if group is not None:
-            all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group)
-        entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits
-    return loss_mask, log_ratio, entropy_per_token
-
-
 def _policy_metrics_reduce(
     ratio: torch.Tensor,  # per token
     effective_log_ratio: torch.Tensor,  # per token; the log-ratio whose exp gives `ratio`
@@ -370,7 +738,7 @@ def _policy_metrics_reduce(
     variance_weight: torch.Tensor,  # per token: token mask (per-token path) or `document_weight` (per-segment)
     epsilon_low: float,
     epsilon_high: float,
-    entropy_per_token: torch.Tensor | None,
+    entropy_per_token: torch.Tensor,
     num_segments: torch.Tensor | None,
 ) -> PolicyMetrics:
     """Shared metric reduction. Document-weighted sums divided by the document count give per-document
@@ -391,28 +759,39 @@ def _policy_metrics_reduce(
         min_advantage=torch.where(loss_mask, advantages, pos_inf).min(),
         num_tokens=loss_mask.to(document_weight.dtype).sum(),
         num_segments=num_segments,
-        entropy=None if entropy_per_token is None else (entropy_per_token * document_weight).sum(),
+        entropy=(entropy_per_token * document_weight).sum(),
     )
 
 
-@torch.compile
-def compute_grpo_metrics(
-    logits: torch.Tensor,  # (*batch, vocab_local)
-    target: torch.Tensor,  # (*batch,)
+def _policy_entropy_per_token(
+    logits_norm: torch.Tensor,  # (*batch, vocab)
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
+    group: torch.distributed.ProcessGroup | None,
+) -> torch.Tensor:
+    """Per-token policy entropy from a student softmax: `log(Σ exp) - E_softmax[logits_norm]`. `exp_logits`
+    and `logits_norm` are local vocab slices, so the expectation sums over the local slice then all-reduces
+    across the tensor-parallel group before dividing by the already-global `sum_exp_logits`."""
+    weighted_logits_sum = (exp_logits * logits_norm).sum(-1)
+    if group is not None:
+        all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group)
+    return sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits
+
+
+def grpo_metrics_core(
+    new_log_probs: torch.Tensor,  # (*batch,) — predicted_logits - log(sum_exp_logits)
     advantages: torch.Tensor,  # (*batch,)
     old_log_probabilities: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool, == target >= 0
     label_counts: torch.Tensor,  # (*batch,) — global per-sequence count broadcast per token
-    epsilon_low: float = 0.2,
-    epsilon_high: float = 0.2,
-    logits_scale_factor: float = 1.0,
-    group: torch.distributed.ProcessGroup | None = None,
-    compute_entropy: bool = False,
+    epsilon_low: float,
+    epsilon_high: float,
+    entropy_per_token: torch.Tensor,  # (*batch,)
 ) -> PolicyMetrics:
-    """Per-token diagnostics: the importance ratio and its clip / KL are token-level, and the ratio
-    variance is over tokens (`variance_weight` = the token mask)."""
-    loss_mask, log_ratio, entropy_per_token = _policy_metrics_log_ratio(
-        logits, target, old_log_probabilities, logits_scale_factor, group, compute_entropy
-    )
+    """Per-token metric family from per-token new log-probs and a precomputed entropy. The importance ratio's
+    clip / KL are token-level, and the ratio variance is over tokens (`variance_weight` = the token mask).
+    Un-compiled so it inlines into a `@torch.compile` boundary."""
+    log_ratio = new_log_probs - old_log_probabilities
     mask = loss_mask.to(log_ratio.dtype)
     return _policy_metrics_reduce(
         ratio=log_ratio.exp(),
@@ -430,32 +809,57 @@ def compute_grpo_metrics(
     )
 
 
-# Not @torch.compile for the same reason as `fused_gspo_loss_forward_backward`: the Python-int
-# `num_segments` argument trips dynamo. Metrics run only on logging steps, so eager is fine.
-def compute_gspo_metrics(
+@torch.compile
+def compute_grpo_metrics(
     logits: torch.Tensor,  # (*batch, vocab_local)
     target: torch.Tensor,  # (*batch,)
     advantages: torch.Tensor,  # (*batch,)
     old_log_probabilities: torch.Tensor,  # (*batch,)
-    document_index_zero_based: torch.Tensor,  # (*batch,) int — segment ID per token, 0-based
-    num_segments: int,
-    label_counts: torch.Tensor,  # (*batch,) — per-document labeled-token count broadcast per token
+    label_counts: torch.Tensor,  # (*batch,) — global per-sequence count broadcast per token
     epsilon_low: float = 0.2,
     epsilon_high: float = 0.2,
     logits_scale_factor: float = 1.0,
     group: torch.distributed.ProcessGroup | None = None,
+) -> PolicyMetrics:
+    """Standalone per-token diagnostics: one softmax over the logits, then the shared `grpo_metrics_core`."""
+    loss_mask = target >= 0
+    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
+    predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group)
+    new_log_probs = predicted_logits - sum_exp_logits.log()
+    return grpo_metrics_core(
+        new_log_probs,
+        advantages,
+        old_log_probabilities,
+        loss_mask,
+        label_counts,
+        epsilon_low,
+        epsilon_high,
+        _policy_entropy_per_token(logits_norm, exp_logits, sum_exp_logits, group),
+    )
+
+
+# Not @torch.compile for the same reason as `fused_gspo_loss_forward_backward`: the Python-int
+# `num_segments` argument in `index_add_` trips dynamo. Metrics run only on logging steps, so eager is fine.
+def gspo_metrics_core(
+    new_log_probs: torch.Tensor,  # (*batch,) — predicted_logits - log(sum_exp_logits)
+    advantages: torch.Tensor,  # (*batch,)
+    old_log_probabilities: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool, == target >= 0
+    document_index_zero_based: torch.Tensor,  # (*batch,) int — segment ID per token, 0-based
+    num_segments: int,
+    label_counts: torch.Tensor,  # (*batch,) — per-document labeled-token count broadcast per token
+    epsilon_low: float,
+    epsilon_high: float,
+    entropy_per_token: torch.Tensor,  # (*batch,)
     sdp_group: torch.distributed.ProcessGroup | None = None,
     sp_group: torch.distributed.ProcessGroup | None = None,
-    compute_entropy: bool = False,
 ) -> PolicyMetrics:
-    """Segment-level diagnostics (clipping is per document/segment): the ratio is the per-segment
-    geometric mean, broadcast back to tokens, and the ratio variance is over segments
-    (`variance_weight` = `document_weight`). The per-segment log-ratio / advantage are token-weighted
-    by `document_weight` (which sums to 1 per document across SDP/SP ranks) then all-reduced, so they
-    partition correctly across ranks and the token-level reduction matches the per-token loss."""
-    loss_mask, log_ratio, entropy_per_token = _policy_metrics_log_ratio(
-        logits, target, old_log_probabilities, logits_scale_factor, group, compute_entropy
-    )
+    """Segment-level metric family from per-token new log-probs and a precomputed entropy. Clipping is per
+    document/segment: the ratio is the per-segment geometric mean, broadcast back to tokens, and the ratio
+    variance is over segments (`variance_weight` = `document_weight`). The per-segment log-ratio / advantage
+    are token-weighted by `document_weight` (which sums to 1 per document across SDP/SP ranks) then
+    all-reduced, so they partition correctly across ranks and the token-level reduction matches the loss."""
+    log_ratio = new_log_probs - old_log_probabilities
     flat_document_index = document_index_zero_based.reshape(-1).long()
     flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype)
     document_weight = flat_mask / label_counts.reshape(-1).to(log_ratio.dtype).clamp(min=1)
@@ -487,91 +891,230 @@ def compute_gspo_metrics(
         variance_weight=document_weight,
         epsilon_low=epsilon_low,
         epsilon_high=epsilon_high,
-        entropy_per_token=None if entropy_per_token is None else entropy_per_token.reshape(-1),
+        entropy_per_token=entropy_per_token.reshape(-1),
         num_segments=document_weight.sum(),
     )
 
 
 @torch.compile
-def fused_grpo_loss_forward_backward(
-    logits: torch.Tensor,  # (*batch, vocab)
+def _gspo_metrics_forward(
+    logits: torch.Tensor,  # (*batch, vocab_local)
+    target: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool
+    logits_scale_factor: float,
+    group: torch.distributed.ProcessGroup | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Vocab-wide part of the standalone GSPO metrics, fused in one compiled pass: softmax, predicted logits
+    and per-token entropy. The segment `index_add_` keeps `gspo_metrics_core` eager (unlike the fully-compiled
+    `compute_grpo_metrics`), so this pulls the vocab pass into one boundary rather than leaving the entropy as
+    an eager reduction over a vocab-wide temporary."""
+    logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group)
+    predicted_logits, _, _ = predicted_logits_from_labels(logits_norm, target, loss_mask, group)
+    new_log_probs = predicted_logits - sum_exp_logits.log()
+    return new_log_probs, _policy_entropy_per_token(logits_norm, exp_logits, sum_exp_logits, group)
+
+
+def compute_gspo_metrics(
+    logits: torch.Tensor,  # (*batch, vocab_local)
     target: torch.Tensor,  # (*batch,)
     advantages: torch.Tensor,  # (*batch,)
     old_log_probabilities: torch.Tensor,  # (*batch,)
-    grad_logits: torch.Tensor | None = None,
-    grad_output: float | None = None,
-    group: torch.distributed.ProcessGroup | None = None,
+    document_index_zero_based: torch.Tensor,  # (*batch,) int — segment ID per token, 0-based
+    num_segments: int,
+    label_counts: torch.Tensor,  # (*batch,) — per-document labeled-token count broadcast per token
     epsilon_low: float = 0.2,
     epsilon_high: float = 0.2,
     logits_scale_factor: float = 1.0,
-    num_labels_in_seq: (
-        torch.Tensor | None
-    ) = None,  # (*batch,) — response-span length broadcast per token, 0 for non-response
-    divisor: float | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
-    if divisor is None:
-        divisor = logits.shape[:-1].numel()
-    grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
+    group: torch.distributed.ProcessGroup | None = None,
+    sdp_group: torch.distributed.ProcessGroup | None = None,
+    sp_group: torch.distributed.ProcessGroup | None = None,
+) -> PolicyMetrics:
+    """Standalone segment-level diagnostics: a fused softmax + entropy pass, then the shared `gspo_metrics_core`."""
     loss_mask = target >= 0
-
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
-    predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels(
-        logits_norm, target, loss_mask, group
+    new_log_probs, entropy_per_token = _gspo_metrics_forward(logits, target, loss_mask, logits_scale_factor, group)
+    return gspo_metrics_core(
+        new_log_probs,
+        advantages,
+        old_log_probabilities,
+        loss_mask,
+        document_index_zero_based,
+        num_segments,
+        label_counts,
+        epsilon_low,
+        epsilon_high,
+        entropy_per_token,
+        sdp_group,
+        sp_group,
     )
+
+
+@torch.compile
+def _gspo_forward_core(
+    logits: torch.Tensor,  # (*batch, vocab)
+    target: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,), == target >= 0
+    logits_scale_factor: float,
+    group: torch.distributed.ProcessGroup | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """GSPO compiled forward: one softmax + predicted-logit lookup, producing per-token new log-probs.
+    The softmax tensors are handed to the eager segment seam and the compiled backward."""
+    logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group)
+    predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group)
     new_log_probs = predicted_logits - sum_exp_logits.log()
-    probability_ratio = (new_log_probs - old_log_probabilities).exp()
+    return new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask
+
+
+@torch.compile
+def _gspo_segment_weights(
+    new_log_probs: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool
+    advantages: torch.Tensor,  # (*batch,)
+    old_log_probabilities: torch.Tensor,  # (*batch,)
+    num_labels_in_seq: torch.Tensor,  # (*batch,)
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compiled pre-aggregation block: the per-token contributions to the two per-segment sums, ready
+    for `index_add_`. Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across
+    SDP/SP ranks too), regardless of how the doc is sharded. Products stay in `new_log_probs.dtype` (fp32)
+    — casting to a possibly-low input dtype before the segment sum would round each contribution."""
+    log_ratio = (new_log_probs - old_log_probabilities).reshape(-1)
+    mean_token_weight = loss_mask.reshape(-1).to(log_ratio.dtype) / num_labels_in_seq.reshape(-1).to(
+        log_ratio.dtype
+    ).clamp(min=1)
+    return log_ratio * mean_token_weight, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight
+
+
+@torch.compile
+def _gspo_segment_loss(
+    mean_log_ratio_per_segment: torch.Tensor,  # (num_segments,)
+    mean_advantage_per_segment: torch.Tensor,  # (num_segments,)
+    flat_document_index: torch.Tensor,  # (*batch,) int
+    new_log_probs: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool
+    num_labels_in_seq: torch.Tensor,  # (*batch,)
+    epsilon_low: float,
+    epsilon_high: float,
+    compute_grad: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Compiled post-aggregation block: from the reduced per-segment sums to the undivided loss sum, the
+    `new_logprobs` metric, and the unscaled per-token backward coefficient
+    `clip_factor · loss_weight · R_s`. The `/ divisor` and `grad_output` scaling stay eager so those
+    per-step-varying scalars never specialize this graph (`epsilon_*` are fixed per run, so they don't)."""
+    segment_ratio = mean_log_ratio_per_segment.exp()  # (num_segments,) — geometric-mean IS ratio
+    segment_advantage = mean_advantage_per_segment.detach()  # (num_segments,) — no grad through A
+
+    probability_ratio = segment_ratio[flat_document_index].reshape(new_log_probs.shape)
+    advantage_per_token = segment_advantage[flat_document_index].reshape(new_log_probs.shape)
+    loss_weight = loss_mask.to(new_log_probs.dtype)
 
     losses = -torch.min(
-        probability_ratio * advantages,
-        torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantages,
-    )
-    loss = reduce_losses(losses, divisor, loss_mask)
-
-    # Sum of per-sequence mean log-probs, matching pipelinerl's new_logprobs metric:
-    #   sum_sum(new_logprobs / num_labels_in_seq, masks_shifted, segments)
-    # Dividing by num_labels_in_seq (span length broadcast per token) and summing over masked
-    # tokens gives mean logprob per sequence; summing those across sequences matches the deepspeed
-    # convention exactly (segments are redundant once num_labels_in_seq is correct).
-    # Clamp to avoid 0/0=nan when num_labels_in_seq=0 (padded tokens or fully masked documents)
-    # — those positions also have loss_mask=0 so they correctly contribute 0 to the sum.
-    new_logprobs_mean = (
-        None if num_labels_in_seq is None else (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum()
+        probability_ratio * advantage_per_token,
+        torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token,
     )
+    loss_sum = (losses * loss_weight).sum()
+    new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum()
 
-    if grad_output is not None:
-        # loss[a>=0] = -a * min(x, 1 + epsilon_high)  =>  grad[a>=0] = -a * (x <= 1 + epsilon_high)
-        # loss[a<=0] = a * max(x, 1 - epsilon_low)  =>  grad[a<=0] = a * (x >= 1 - epsilon_low)
-        probability_ratio_grad = (
-            grad_output
-            * (
-                torch.clamp_min(advantages, 0) * (probability_ratio <= 1 + epsilon_high)
-                + torch.clamp_max(advantages, 0) * (probability_ratio >= 1 - epsilon_low)
+    if compute_grad:
+        effective_grad_unscaled = (
+            (
+                torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high)
+                + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low)
             )
-            * loss_mask
+            * loss_weight
+            * probability_ratio
         )
+    else:
+        effective_grad_unscaled = None
+    return loss_sum, new_logprobs_mean, effective_grad_unscaled
 
-        # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities)
-        # (Sign absorbed in probability_ratio_grad)
-        predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze_(-1)
-        grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add(
-            -1,
-            target_masked.unsqueeze(-1),
-            -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32),
-        )
-        grad = grad.to(logits.dtype)
 
-        if grad_logits is None:
-            grad_logits = grad
-        else:
-            grad_logits.add_(grad)
+def gspo_segment_seam(
+    new_log_probs: torch.Tensor,  # (*batch,)
+    loss_mask: torch.Tensor,  # (*batch,) bool
+    advantages: torch.Tensor,  # (*batch,)
+    old_log_probabilities: torch.Tensor,  # (*batch,)
+    document_index_zero_based: torch.Tensor,  # (*batch,) int
+    num_segments: int,
+    num_labels_in_seq: torch.Tensor,  # (*batch,)
+    divisor: float,
+    grad_output: float | None,
+    sdp_group: torch.distributed.ProcessGroup | None,
+    sp_group: torch.distributed.ProcessGroup | None,
+    epsilon_low: float,
+    epsilon_high: float,
+    logits_scale_factor: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Eager segment seam between the compiled forward and backward, orchestrating two compiled blocks
+    around the parts that can't compile: the `index_add_` segment aggregation (whose symbolic
+    `num_segments` would trigger per-value recompiles) and the SDP/SP all-reduces. Returns the loss, the
+    `new_logprobs` metric, and the per-token backward coefficient
+    `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` (None when no gradient is requested)."""
+    flat_document_index = document_index_zero_based.reshape(-1).long()
+    weighted_log_ratio, weighted_advantage = _gspo_segment_weights(
+        new_log_probs, loss_mask, advantages, old_log_probabilities, num_labels_in_seq
+    )
+    mean_log_ratio_per_segment = weighted_log_ratio.new_zeros(num_segments).index_add_(
+        0, flat_document_index, weighted_log_ratio
+    )
+    mean_advantage_per_segment = weighted_advantage.new_zeros(num_segments).index_add_(
+        0, flat_document_index, weighted_advantage
+    )
+    for reduce_group in (sdp_group, sp_group):
+        if reduce_group is not None:
+            torch.distributed.all_reduce(
+                mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group
+            )
+            torch.distributed.all_reduce(
+                mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group
+            )
 
-    return loss, grad_logits, new_logprobs_mean
+    loss_sum, new_logprobs_mean, effective_grad_unscaled = _gspo_segment_loss(
+        mean_log_ratio_per_segment,
+        mean_advantage_per_segment,
+        flat_document_index,
+        new_log_probs,
+        loss_mask,
+        num_labels_in_seq,
+        epsilon_low,
+        epsilon_high,
+        grad_output is not None,
+    )
+    loss = loss_sum / divisor
+    effective_grad = (
+        None if grad_output is None else effective_grad_unscaled * (grad_output / divisor * logits_scale_factor)
+    )
+    return loss, new_logprobs_mean, effective_grad
+
+
+@torch.compile
+def gspo_backward_core(
+    exp_logits: torch.Tensor,  # (*batch, vocab)
+    sum_exp_logits: torch.Tensor,  # (*batch,)
+    target_masked: torch.Tensor,  # (*batch,)
+    target_mask: torch.Tensor | None,  # (*batch,) or None (no TP)
+    loss_mask: torch.Tensor,  # (*batch,) bool
+    effective_grad: torch.Tensor,  # (*batch,) — per-token backward coefficient from the seam
+    logits_dtype: torch.dtype,
+    grad_logits: torch.Tensor | None,
+) -> torch.Tensor:
+    """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into one
+    kernel. `sum_exp_logits.unsqueeze` is out-of-place (the standalone eager kernel mutates it)."""
+    predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1)
+    grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add(
+        -1,
+        target_masked.unsqueeze(-1),
+        -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32),
+    )
+    grad = grad.to(logits_dtype)
+    if grad_logits is None:
+        grad_logits = grad
+    else:
+        grad_logits.add_(grad)
+    return grad_logits
 
 
-# Not @torch.compile: dynamo lifts the Python-int `divisor` and `num_segments` args to symbolic
-# ints with no concrete hint, then trips on `grad_output / divisor` during trace evaluation
-# (`ZeroDivisionError` with hint=0). The Triton kernel is the actual GPU perf path; the eager
-# PyTorch fallback doesn't need to be compiled.
+# Orchestrator only: between the compiled forward and backward cores, the segment seam keeps its
+# `index_add_` (with the Python-int `num_segments`) and SDP/SP all-reduces eager, bracketing them with
+# compiled sub-blocks, so `num_segments` never enters a compiled boundary.
 def fused_gspo_loss_forward_backward(
     logits: torch.Tensor,  # (*batch, vocab)
     target: torch.Tensor,  # (*batch,)
@@ -612,76 +1155,35 @@ def fused_gspo_loss_forward_backward(
     computes the same global R_s/A_s. Token-level contributions remain per-rank, so summing the
     kernel loss across SDP/SP via SUM reduction matches the canonical single-rank result.
     """
-    grad_output_scaled = None if grad_output is None else grad_output / divisor * logits_scale_factor
     loss_mask = target >= 0
-
-    logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group)
-    predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels(
-        logits_norm, target, loss_mask, group
+    new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask = _gspo_forward_core(
+        logits, target, loss_mask, logits_scale_factor, group
     )
-    new_log_probs = predicted_logits - sum_exp_logits.log()
-    log_ratio = new_log_probs - old_log_probabilities
-
-    flat_document_index = document_index_zero_based.reshape(-1).long()
-    flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype)
-    # Per-token weight: mask / per-document label count, from the preprocessor.
-    # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across
-    # SDP/SP ranks too), regardless of how the doc is sharded.
-    mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1)
-    # Pre-divide the per-token contributions by the per-doc label count, then sum per segment.
-    # All tokens in a segment share the same N_d, so this is mathematically equivalent to
-    # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction.
-    mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_(
-        0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight
-    )
-    # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype`
-    # before summing would round each token's contribution to a possibly-low input dtype.
-    mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_(
-        0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight
-    )
-    for reduce_group in (sdp_group, sp_group):
-        if reduce_group is not None:
-            torch.distributed.all_reduce(
-                mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group
-            )
-            torch.distributed.all_reduce(
-                mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group
-            )
-
-    segment_ratio = mean_log_ratio_per_segment.exp()  # (num_segments,) — geometric-mean IS ratio
-    segment_advantage = mean_advantage_per_segment.detach()  # (num_segments,) — no grad through A
-
-    probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape)
-    advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape)
-    loss_weight = loss_mask.to(log_ratio.dtype)
-
-    losses = -torch.min(
-        probability_ratio * advantage_per_token,
-        torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token,
+    loss, new_logprobs_mean, effective_grad = gspo_segment_seam(
+        new_log_probs,
+        loss_mask,
+        advantages,
+        old_log_probabilities,
+        document_index_zero_based,
+        num_segments,
+        num_labels_in_seq,
+        divisor,
+        grad_output,
+        sdp_group,
+        sp_group,
+        epsilon_low,
+        epsilon_high,
+        logits_scale_factor,
     )
-    loss = (losses * loss_weight).sum() / divisor
-
-    new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum()
-
-    if grad_output_scaled is not None:
-        probability_ratio_grad = (
-            grad_output_scaled
-            * (
-                torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high)
-                + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low)
-            )
-            * loss_weight
-        )
-        predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze_(-1)
-        grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add(
-            -1,
-            target_masked.unsqueeze(-1),
-            -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32),
+    if effective_grad is not None:
+        grad_logits = gspo_backward_core(
+            exp_logits,
+            sum_exp_logits,
+            target_masked,
+            target_mask,
+            loss_mask,
+            effective_grad,
+            logits.dtype,
+            grad_logits,
         )
-        grad = grad.to(logits.dtype)
-        if grad_logits is None:
-            grad_logits = grad
-        else:
-            grad_logits.add_(grad)
-
     return loss, grad_logits, new_logprobs_mean
diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py
index 2e5f90b1d..475685587 100644
--- a/fast_llm/layers/language_model/loss/z_loss.py
+++ b/fast_llm/layers/language_model/loss/z_loss.py
@@ -3,14 +3,17 @@
 import torch
 
 from fast_llm.functional.config import TritonConfig
-from fast_llm.functional.entropy_loss import fused_softmax_base
+from fast_llm.functional.entropy_loss import z_loss_core
 from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward
 from fast_llm.functional.utils import reduce_losses
 from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig
-from fast_llm.layers.language_model.loss.loss import LanguageModelLoss
+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 LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](LanguageModelLoss[ConfigType]):
+
+class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](CombinableLoss, SingleLoss[ConfigType]):
     def _forward_backward(
         self,
         logits: "torch.Tensor",
@@ -19,19 +22,57 @@ def _forward_backward(
         split_index: int = 0,
         grad_logits: torch.Tensor | None = None,
     ) -> "tuple[torch.Tensor, torch.Tensor | None]":
-        return (
-            triton_z_loss_forward_backward
-            if TritonConfig.enabled(logits.device, self._config.use_triton)
-            else fused_z_loss_forward_backward
-        )(
-            logits,
-            self._get_loss_mask(kwargs, split_index),
-            grad_output=self._get_grad_output(kwargs),
-            group=self._parallel_dim.group if self._vocab_parallel else None,
-            logits_scale_factor=self._logits_scale_factor,
-            grad_logits=grad_logits,
-            divisor=self._get_label_count(kwargs),
-        )
+        arguments = self.get_inputs(kwargs, split_index, losses is not None)
+        group = self._parallel_dim.group if self._vocab_parallel else None
+        if TritonConfig.enabled(logits.device, self._config.use_triton):
+            loss_mask, grad_output, divisor = arguments
+            return triton_z_loss_forward_backward(
+                logits,
+                loss_mask,
+                grad_output=grad_output,
+                group=group,
+                logits_scale_factor=self._logits_scale_factor,
+                grad_logits=grad_logits,
+                divisor=divisor,
+            )
+        loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments)
+        return loss, grad_logits
+
+    def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple:
+        return self._get_loss_mask(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs)
+
+    @staticmethod
+    def fused_core(
+        logits_norm: torch.Tensor,
+        exp_logits: torch.Tensor,
+        sum_exp_logits: torch.Tensor,
+        logits_max: torch.Tensor,
+        group: "torch.distributed.ProcessGroup | None",
+        logits_scale_factor: float,
+        arguments: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, None]:
+        """Post-softmax z-loss over the shared softmax. Returns the loss scalar and the uncast, masked
+        gradient contribution (the caller casts); z-loss emits no extra outputs."""
+        loss_mask, grad_output, divisor = arguments
+        grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
+        loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output)
+        loss = reduce_losses(loss_term, divisor, loss_mask)
+        if grad is not None and loss_mask is not None:
+            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:
+        loss_mask, grad_output, divisor = self.get_inputs(kwargs, split_index, register)
+        if context.divisor is None:
+            context.divisor = divisor
+        context.z = (loss_mask, grad_output)
+
+    def triton_finish(
+        self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool
+    ) -> tuple[torch.Tensor, None]:
+        return context.z_loss, None
 
     def get_preprocessing_config(self) -> dict[str, typing.Any]:
         return {"return_prediction_mask": True}
@@ -49,38 +90,3 @@ def z_loss(
     if loss_mask is not None:
         out = out * loss_mask
     return torch.mean(out)
-
-
-@torch.compile
-def fused_z_loss_forward_backward(
-    logits: torch.Tensor,
-    loss_mask: torch.Tensor | None,
-    grad_logits: torch.Tensor | None = None,
-    grad_output: float | None = None,
-    group: torch.distributed.ProcessGroup | None = None,
-    logits_scale_factor: float = 1.0,
-    divisor: float | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
-    """
-    Z-loss = mean(logsumexp(logits, dim=-1) ** 2)
-    Grad = 2 * log_sum_exp_logits * softmax(logits)
-    """
-    if divisor is None:
-        divisor = logits.shape[:-1].numel()
-    grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor
-    logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group)
-    log_sum_exp_logits = sum_exp_logits.log() + logits_max
-
-    loss = reduce_losses(log_sum_exp_logits**2, divisor, loss_mask)
-
-    if grad_output is not None:
-        grad_base = 2 * grad_output * (log_sum_exp_logits / sum_exp_logits)
-        if loss_mask is not None:
-            grad_base = grad_base * loss_mask
-        grad = (grad_base.unsqueeze(-1) * exp_logits).to(logits.dtype)
-        if grad_logits is None:
-            grad_logits = grad
-        else:
-            grad_logits.add_(grad)
-
-    return loss, grad_logits
diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py
index 6f09cb108..5df2e292d 100644
--- a/tests/layers/test_lm_head.py
+++ b/tests/layers/test_lm_head.py
@@ -6,6 +6,7 @@
 import torch
 
 from fast_llm.engine.config_utils.data_type import DataType
+from fast_llm.functional.triton import triton_available
 from fast_llm.layers.attention.config import AttentionKwargs
 from fast_llm.layers.block.config import BlockKwargs
 from fast_llm.layers.language_model.config import LM_HEAD_LOSS_NAME, LanguageModelKwargs
@@ -13,7 +14,12 @@
 from fast_llm.layers.language_model.loss.config import LanguageModelLossKwargs
 from fast_llm.models.gpt.config import GPTModelConfig
 from fast_llm.utils import Assert
-from tests.layers.test_lm_losses import reference_grpo_loss, reference_gspo_loss
+from tests.layers.test_lm_losses import (
+    reference_grpo_loss,
+    reference_grpo_metrics,
+    reference_gspo_loss,
+    reference_gspo_metrics,
+)
 from tests.utils.utils import get_base_model, get_stage
 
 NUM_TOKENS = 200
@@ -27,6 +33,7 @@ class LMHeadTestConfig:
     name: str
     label_loss: bool | float = False
     distillation_loss: bool | float = False
+    distillation_temperature: float = 1.0
     z_loss: bool | float = False
     grpo_loss: bool | float = False
     gspo_loss: bool | float = False
@@ -39,6 +46,9 @@ class LMHeadTestConfig:
     tied_embedding_weight: bool = False
     num_splits: int = 1
     gspo_document_lengths: tuple[int, ...] | None = None
+    loss_implementation: str = "per_loss"
+    grpo_metrics: str | None = None
+    gspo_metrics: str | None = None
 
     @property
     def actual_label_loss(self):
@@ -68,6 +78,8 @@ def get_config(self) -> GPTModelConfig:
                 losses["label"]["weight"] = self.label_loss
         if self.distillation_loss is not False:
             losses["distillation"] = {"type": "distillation", "reference_model": "distillation"}
+            if self.distillation_temperature != 1.0:
+                losses["distillation"]["temperature"] = self.distillation_temperature
             if isinstance(self.distillation_loss, float):
                 losses["distillation"]["weight"] = self.distillation_loss
         if self.z_loss is not False:
@@ -75,13 +87,31 @@ def get_config(self) -> GPTModelConfig:
             if isinstance(self.z_loss, float):
                 losses["z_loss"]["weight"] = self.z_loss
         if self.grpo_loss is not False:
-            losses["grpo_loss"] = {"type": "grpo"}
+            # Metrics default to `auto` (→ `basic` at `pipeline_parallel == 1`), so pin `none` explicitly
+            # unless the test exercises them.
+            losses["grpo_loss"] = {"type": "grpo", "metrics": self.grpo_metrics or "none"}
             if isinstance(self.grpo_loss, float):
                 losses["grpo_loss"]["weight"] = self.grpo_loss
         if self.gspo_loss is not False:
-            losses["gspo_loss"] = {"type": "gspo"}
+            losses["gspo_loss"] = {"type": "gspo", "metrics": self.gspo_metrics or "none"}
             if isinstance(self.gspo_loss, float):
                 losses["gspo_loss"]["weight"] = self.gspo_loss
+        if self.loss_implementation in ("fused", "fused_triton") and losses:
+            # Wrap the combinable losses in a single `monolithic` loss that shares one softmax; keep the
+            # child keys so the registered metric names match the per-loss configuration. `fused` pins the
+            # compiled path and `fused_triton` the triton path, so both are exercised in every environment.
+            combinable = {
+                name: loss
+                for name, loss in losses.items()
+                if loss["type"] in ("label", "distillation", "z_loss", "grpo", "gspo")
+            }
+            if combinable:
+                losses = {name: loss for name, loss in losses.items() if name not in combinable}
+                losses["monolithic"] = {
+                    "type": "monolithic",
+                    "losses": combinable,
+                    "use_triton": self.loss_implementation == "fused_triton",
+                }
         if losses:
             head_config["losses"] = losses
 
@@ -239,9 +269,12 @@ def get_reference_outputs(
             names_losses_weights.append(("label", label_loss, float(self.actual_label_loss)))
 
         if self.distillation_loss is not False:
+            # Teacher logits are scaled by `logits_scale_factor / temperature` before the softmax, matching the kernel.
+            teacher_logits = kwargs[f"reference_distillation_hidden_states"]["head.logits"].float()
+            teacher_logits = teacher_logits * (self.logits_scale_factor / self.distillation_temperature)
             distillation_loss = torch.nn.functional.cross_entropy(
                 logits,
-                torch.softmax(kwargs[f"reference_distillation_hidden_states"]["head.logits"], -1),
+                torch.softmax(teacher_logits, -1),
                 reduction="mean" if loss_mask is None else "none",
             )
             if loss_mask is not None:
@@ -265,6 +298,27 @@ def get_reference_outputs(
             names_losses_weights.append(("grpo_loss", grpo_loss, float(self.grpo_loss)))
             names_losses_weights.append(("grpo_loss_new_logprobs", new_logprobs, 0.0))
 
+            if self.grpo_metrics is not None:
+                # `logits` is already scaled above, so pass logits_scale_factor=1.0.
+                metrics = reference_grpo_metrics(
+                    logits,
+                    labels,
+                    kwargs[LanguageModelLossKwargs.advantages][head._prediction_distance - 1],
+                    kwargs[LanguageModelLossKwargs.old_log_probabilities][head._prediction_distance - 1],
+                    kwargs[LanguageModelLossKwargs.label_counts][head._prediction_distance - 1],
+                    epsilon_low=0.2,
+                    epsilon_high=0.2,
+                    logits_scale_factor=1.0,
+                )
+                num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch]
+                for attr in ("old_logprobs", "ratio_new_old", "kl_new_old", "clipped_ratio_fraction", "advantage"):
+                    names_losses_weights.append((f"grpo_loss_{attr}", getattr(metrics, attr) / num_documents, 0.0))
+                for attr in ("ratio_new_old_sum", "ratio_new_old_squared_sum", "num_tokens"):
+                    names_losses_weights.append((f"grpo_loss_{attr}", getattr(metrics, attr), 0.0))
+                names_losses_weights.append(("grpo_loss_max_advantage", metrics.max_advantage, 0.0))
+                names_losses_weights.append(("grpo_loss_min_advantage", metrics.min_advantage, 0.0))
+                names_losses_weights.append(("grpo_loss_entropy", metrics.entropy / num_documents, 0.0))
+
         if self.gspo_loss is not False:
             gspo_loss, _ = reference_gspo_loss(
                 logits,
@@ -295,6 +349,30 @@ def get_reference_outputs(
             names_losses_weights.append(("gspo_loss", gspo_loss, float(self.gspo_loss)))
             names_losses_weights.append(("gspo_loss_new_logprobs", new_logprobs, 0.0))
 
+            if self.gspo_metrics is not None:
+                # `logits` is already scaled above, so pass logits_scale_factor=1.0.
+                metrics = reference_gspo_metrics(
+                    logits,
+                    labels,
+                    kwargs[LanguageModelLossKwargs.advantages][head._prediction_distance - 1],
+                    kwargs[LanguageModelLossKwargs.old_log_probabilities][head._prediction_distance - 1],
+                    kwargs[BlockKwargs.global_document_index_q].long() - 1,
+                    kwargs[BlockKwargs.num_documents_in_sequence],
+                    kwargs[LanguageModelLossKwargs.label_counts][head._prediction_distance - 1],
+                    epsilon_low=0.2,
+                    epsilon_high=0.2,
+                    logits_scale_factor=1.0,
+                )
+                num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch]
+                for attr in ("old_logprobs", "ratio_new_old", "kl_new_old", "clipped_ratio_fraction", "advantage"):
+                    names_losses_weights.append((f"gspo_loss_{attr}", getattr(metrics, attr) / num_documents, 0.0))
+                for attr in ("ratio_new_old_sum", "ratio_new_old_squared_sum", "num_tokens"):
+                    names_losses_weights.append((f"gspo_loss_{attr}", getattr(metrics, attr), 0.0))
+                names_losses_weights.append(("gspo_loss_max_advantage", metrics.max_advantage, 0.0))
+                names_losses_weights.append(("gspo_loss_min_advantage", metrics.min_advantage, 0.0))
+                names_losses_weights.append(("gspo_loss_num_segments", metrics.num_segments, 0.0))
+                names_losses_weights.append(("gspo_loss_entropy", metrics.entropy / num_documents, 0.0))
+
         actual_losses = [loss * weight for _, loss, weight in names_losses_weights if weight != 0.0]
         total_loss = sum(actual_losses)
         total_loss.backward()
@@ -359,6 +437,133 @@ def _add_configs(base_name: str, **kwargs):
 _add_configs("label_and_distillation_loss", label_loss=True, distillation_loss=True)
 _add_configs("label_and_z_loss_weighted", label_loss=True, z_loss=0.5)
 _add_configs("label_and_distillation_loss_zero_weight", label_loss=True, distillation_loss=0.0)
+_add_configs("distillation_loss_temperature", distillation_loss=True, distillation_temperature=2.0)
+
+# Monolithic loss type: the combinable losses are wrapped in a single `monolithic` loss that shares one
+# softmax pass; the head treats it as an ordinary loss. These configs must match their per-loss equivalents
+# above (validated against the same independent reference).
+_add_configs("fused", loss_implementation="fused")
+_add_configs("fused_bfloat16", loss_implementation="fused", compute_dtype=DataType.bfloat16)
+_add_configs("fused_logit_scaling", loss_implementation="fused", logits_scale_factor=5.0)
+_add_configs("fused_final_logit_softcap", loss_implementation="fused", final_logit_softcap=2.0)
+_add_configs("fused_tied_embedding_weight", loss_implementation="fused", tied_embedding_weight=True)
+_add_configs("fused_multi_token_prediction", loss_implementation="fused", prediction_heads=2)
+_add_configs("fused_label_and_z_loss_weighted", loss_implementation="fused", label_loss=True, z_loss=0.5)
+_add_configs("fused_distillation_loss", loss_implementation="fused", distillation_loss=True)
+_add_configs("fused_label_and_distillation_loss", loss_implementation="fused", label_loss=True, distillation_loss=True)
+_add_configs(
+    "fused_distillation_loss_temperature",
+    loss_implementation="fused",
+    distillation_loss=True,
+    distillation_temperature=2.0,
+)
+_add_configs("fused_grpo_loss", loss_implementation="fused", grpo_loss=True)
+# Multi-loss combos sharing one softmax pass: a three-way distillation combo and an RL + regularizer combo.
+_add_configs(
+    "fused_label_distillation_z_loss",
+    loss_implementation="fused",
+    label_loss=True,
+    distillation_loss=True,
+    z_loss=0.5,
+)
+_add_configs("fused_grpo_and_z_loss", loss_implementation="fused", grpo_loss=True, z_loss=0.5)
+# GSPO fused into the shared softmax: its eager segment seam runs between the compiled forward and backward,
+# so it defers to `finish` rather than a single-pass `fused_core`. Single-split only (GSPO can't be split);
+# added explicitly since `_add_configs` also emits a split variant. Alone (wrapper only) and sharing the
+# softmax with z-loss (the RL + regularizer combo).
+for _loss_masking in (False, True):
+    _suffix = "_masked" if _loss_masking else ""
+    _lm_head_test_configs.append(
+        LMHeadTestConfig(
+            f"fused_gspo_loss{_suffix}",
+            gspo_loss=True,
+            loss_masking=_loss_masking,
+            loss_implementation="fused",
+        )
+    )
+    _lm_head_test_configs.append(
+        LMHeadTestConfig(
+            f"fused_gspo_and_z_loss{_suffix}",
+            gspo_loss=True,
+            z_loss=0.5,
+            loss_masking=_loss_masking,
+            loss_implementation="fused",
+        )
+    )
+# GRPO/GSPO metric families (`basic` includes entropy) across all three paths — standalone, the compiled
+# shared softmax, and the triton kernel. Single-split only: per-split metric partials reduce across splits,
+# which the whole-sequence reference doesn't model.
+for _loss_implementation in ("per_loss", "fused", "fused_triton"):
+    _prefix = "" if _loss_implementation == "per_loss" else f"{_loss_implementation}_"
+    for _loss_masking in (False, True):
+        _mask_suffix = "_masked" if _loss_masking else ""
+        _lm_head_test_configs.append(
+            LMHeadTestConfig(
+                f"{_prefix}grpo_loss_metrics{_mask_suffix}",
+                grpo_loss=True,
+                grpo_metrics="basic",
+                loss_masking=_loss_masking,
+                loss_implementation=_loss_implementation,
+            )
+        )
+        _lm_head_test_configs.append(
+            LMHeadTestConfig(
+                f"{_prefix}gspo_loss_metrics{_mask_suffix}",
+                gspo_loss=True,
+                gspo_metrics="basic",
+                loss_masking=_loss_masking,
+                loss_implementation=_loss_implementation,
+            )
+        )
+# The metric family co-resides with z-loss in the shared softmax pass, on both fused paths. Single-split
+# (metrics can't be split).
+for _loss_implementation in ("fused", "fused_triton"):
+    for _loss_masking in (False, True):
+        _lm_head_test_configs.append(
+            LMHeadTestConfig(
+                f"{_loss_implementation}_grpo_and_z_loss_metrics{'_masked' if _loss_masking else ''}",
+                grpo_loss=True,
+                z_loss=0.5,
+                grpo_metrics="basic",
+                loss_masking=_loss_masking,
+                loss_implementation=_loss_implementation,
+            )
+        )
+# `auto` metrics resolve to `basic` when pipeline_parallel == 1 (all head tests), covering the default path.
+_lm_head_test_configs.append(LMHeadTestConfig("grpo_loss_metrics_auto", grpo_loss=True, grpo_metrics="auto"))
+_lm_head_test_configs.append(LMHeadTestConfig("gspo_loss_metrics_auto", gspo_loss=True, gspo_metrics="auto"))
+
+# Triton monolithic kernel (`use_triton=True`): the label-based objective set over the shared softmax.
+# Distillation has no triton fused kernel, so it stays on the compiled `fused` configs (policy metrics do).
+_add_configs("fused_triton", loss_implementation="fused_triton")
+_add_configs("fused_triton_z_loss", loss_implementation="fused_triton", z_loss=True)
+_add_configs("fused_triton_bfloat16", loss_implementation="fused_triton", compute_dtype=DataType.bfloat16)
+_add_configs("fused_triton_logit_scaling", loss_implementation="fused_triton", logits_scale_factor=5.0)
+_add_configs("fused_triton_final_logit_softcap", loss_implementation="fused_triton", final_logit_softcap=2.0)
+_add_configs("fused_triton_label_and_z_loss_weighted", loss_implementation="fused_triton", label_loss=True, z_loss=0.5)
+_add_configs("fused_triton_grpo_loss", loss_implementation="fused_triton", grpo_loss=True)
+_add_configs("fused_triton_grpo_and_z_loss", loss_implementation="fused_triton", grpo_loss=True, z_loss=0.5)
+# GSPO on the triton path (its eager segment seam brackets the triton forward and backward). Single-split
+# only; alone and sharing the softmax with z-loss.
+for _loss_masking in (False, True):
+    _suffix = "_masked" if _loss_masking else ""
+    _lm_head_test_configs.append(
+        LMHeadTestConfig(
+            f"fused_triton_gspo_loss{_suffix}",
+            gspo_loss=True,
+            loss_masking=_loss_masking,
+            loss_implementation="fused_triton",
+        )
+    )
+    _lm_head_test_configs.append(
+        LMHeadTestConfig(
+            f"fused_triton_gspo_and_z_loss{_suffix}",
+            gspo_loss=True,
+            z_loss=0.5,
+            loss_masking=_loss_masking,
+            loss_implementation="fused_triton",
+        )
+    )
 
 
 @pytest.mark.slow
@@ -370,6 +575,8 @@ def _add_configs(base_name: str, **kwargs):
     ],
 )
 def test_lm_head(test_config: LMHeadTestConfig):
+    if test_config.loss_implementation == "fused_triton" and not triton_available:
+        pytest.skip("Triton is not available (build the extension or set TRITON_INTERPRET=1).")
     model_config = test_config.get_config()
     model, distributed = get_base_model(model_config)
     input_, kwargs = test_config.get_inputs()
diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py
index 07d42f02a..49118c651 100644
--- a/tests/layers/test_lm_losses.py
+++ b/tests/layers/test_lm_losses.py
@@ -6,26 +6,36 @@
 import torch
 
 from fast_llm.core.ops import split_op
-from fast_llm.engine.config_utils import data_type
 from fast_llm.engine.config_utils.data_type import DataType
-from fast_llm.engine.distributed.config import DistributedBackend
+from fast_llm.engine.distributed.config import DistributedBackend, DistributedConfig
 from fast_llm.functional.config import EntropyLossType, TargetFormat
-from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward, torch_entropy_loss_forward_backward
+from fast_llm.functional.entropy_loss import torch_entropy_loss_forward_backward
 from fast_llm.functional.triton import triton_available
 from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward
 from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward
 from fast_llm.functional.triton.gspo_loss import triton_gspo_loss_forward_backward
+from fast_llm.functional.triton.monolithic_loss import (
+    _monolithic_forward_reduce,
+    triton_monolithic_loss_forward_backward,
+)
 from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward
+from fast_llm.layers.language_model.loss.config import (
+    LanguageModelDistillationLossConfig,
+    LanguageModelGRPOLossConfig,
+    LanguageModelLabelEntropyLossConfig,
+    LanguageModelZLossConfig,
+)
 from fast_llm.layers.language_model.loss.dpo import dpo_loss
 from fast_llm.layers.language_model.loss.loss import loss_forward_backward
+from fast_llm.layers.language_model.loss.monolithic import _monolithic_core
 from fast_llm.layers.language_model.loss.policy_gradient import (
     PolicyMetrics,
     compute_grpo_metrics,
     compute_gspo_metrics,
-    fused_grpo_loss_forward_backward,
     fused_gspo_loss_forward_backward,
+    gspo_segment_seam,
 )
-from fast_llm.layers.language_model.loss.z_loss import fused_z_loss_forward_backward, z_loss
+from fast_llm.layers.language_model.loss.z_loss import z_loss
 from fast_llm.utils import Assert
 from tests.utils.dataset import get_random_spans
 from tests.utils.subtest import DistributedTestContext
@@ -75,8 +85,13 @@ def _compare_losses_and_grads(
     ref_grad: torch.Tensor | None,
     threshold=1e-5,
     group: torch.distributed.ProcessGroup | None = None,
+    loss_min_threshold: float = 1e-6,
 ):
-    Assert.rms_close_relative(loss, ref_loss, threshold, 1e-6)
+    # `loss_min_threshold` raises the absolute floor of the scalar-loss comparison. GSPO's per-segment
+    # geometric-mean reduction accumulates more float32 error than a per-token sum, and its loss can be near
+    # zero, so the fused-vs-reference difference sits at the abs floor rather than the relative band; the
+    # default `1e-6` floor flakes there. The gradient stays tight.
+    Assert.rms_close_relative(loss, ref_loss, threshold, loss_min_threshold)
     if has_grad:
         Assert.rms_close_relative(
             grad, split_op(ref_grad, group, -1), threshold, 1e-8 if grad.dtype == torch.float32 else 1e-7
@@ -137,7 +152,6 @@ def reference_grpo_metrics(
     epsilon_low: float,
     epsilon_high: float,
     logits_scale_factor: float,
-    compute_entropy: bool,
 ) -> PolicyMetrics:
     log_softmax = torch.nn.functional.log_softmax(logits.float() * logits_scale_factor, dim=-1)
     loss_mask = target >= 0
@@ -150,10 +164,8 @@ def reference_grpo_metrics(
     clipped = (ratio < 1.0 - epsilon_low) | (ratio > 1.0 + epsilon_high)
     kl = ratio - log_ratio - 1.0
 
-    entropy = None
-    if compute_entropy:
-        entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1)
-        entropy = (entropy_per_token * masked).sum()
+    entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1)
+    entropy = (entropy_per_token * masked).sum()
 
     return PolicyMetrics(
         old_logprobs=(old_log_probabilities.float() * masked).sum(),
@@ -182,7 +194,6 @@ def reference_gspo_metrics(
     epsilon_low: float,
     epsilon_high: float,
     logits_scale_factor: float,
-    compute_entropy: bool,
 ) -> PolicyMetrics:
     log_softmax = torch.nn.functional.log_softmax(logits.float() * logits_scale_factor, dim=-1)
     loss_mask = target >= 0
@@ -218,11 +229,9 @@ def reference_gspo_metrics(
         old_logprobs = old_logprobs + flat_old[in_segment].sum() / count_float
         num_segments_count = num_segments_count + 1.0
 
-    entropy = None
-    if compute_entropy:
-        entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1).reshape(-1)
-        masked = flat_mask.float() / num_labels_in_seq.reshape(-1).float().clamp(min=1)
-        entropy = (entropy_per_token * masked).sum()
+    entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1).reshape(-1)
+    masked = flat_mask.float() / num_labels_in_seq.reshape(-1).float().clamp(min=1)
+    entropy = (entropy_per_token * masked).sum()
 
     return PolicyMetrics(
         old_logprobs=old_logprobs,
@@ -344,7 +353,6 @@ def _test_entropy_loss(
 ):
     if target_format == TargetFormat.labels and entropy_loss_type == EntropyLossType.reverse_kl:
         pytest.skip(reason="Reverse KL loss not implemented for target labels")
-    # TODO: Test tensor-parallel implementation.
     logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, target_format, batch_shape, dtype)
     local_logits = split_op(logits, group, -1).contiguous()
     local_target = target if target_format == TargetFormat.labels else split_op(target, group, -1).contiguous()
@@ -362,16 +370,19 @@ def _test_entropy_loss(
         previous_grad = torch.randn_like(grad_ref)
         grad_ref = grad_ref + previous_grad
         local_previous_grad = split_op(previous_grad, group, -1).contiguous()
-    out_fused, grad_fused = fused_entropy_loss_forward_backward(
-        logits=local_logits,
-        target=local_target,
-        loss_mask=loss_mask,
-        grad_logits=local_previous_grad.clone() if accumulate else None,
-        grad_output=grad_output,
-        group=group,
-        logits_scale_factor=logits_scale_factor,
-        target_format=target_format,
-        entropy_loss_type=entropy_loss_type,
+    divisor = local_logits.shape[:-1].numel()
+    if target_format == TargetFormat.labels:
+        loss = _combinable_loss(
+            LanguageModelLabelEntropyLossConfig(loss_type=entropy_loss_type), "ce", logits_scale_factor
+        )
+        arguments = (local_target, grad_output, divisor)
+    else:
+        loss = _combinable_loss(
+            LanguageModelDistillationLossConfig(loss_type=entropy_loss_type), "distillation", logits_scale_factor
+        )
+        arguments = (local_target, loss_mask, grad_output, divisor, entropy_loss_type, 1.0)
+    out_fused, grad_fused, _ = loss.combinable_forward_backward(
+        local_logits, group, local_previous_grad.clone() if accumulate else None, arguments
     )
     _compare_losses_and_grads(
         out_fused,
@@ -379,7 +390,7 @@ def _test_entropy_loss(
         grad_output is not None,
         grad_fused,
         grad_ref,
-        threshold=1e-5 if data_type == DataType.float32 else 1e-4,
+        threshold=1e-5 if dtype == DataType.float32 else 1e-4,
         group=group,
     )
 
@@ -403,7 +414,7 @@ def _test_entropy_loss(
         grad_output is not None,
         grad_triton,
         grad_ref,
-        threshold=1e-5 if target_format != TargetFormat.probabilities and data_type == DataType.float32 else 1e-4,
+        threshold=1e-5 if dtype == DataType.float32 else 1e-4,
         group=group,
     )
 
@@ -434,17 +445,12 @@ def _test_grpo_loss(
         previous_grad = torch.randn_like(grad_ref)
         grad_ref = grad_ref + previous_grad
         local_previous_grad = split_op(previous_grad, group, -1).contiguous()
-    out_fused, grad_fused, new_logprobs_fused = fused_grpo_loss_forward_backward(
+    loss = _combinable_loss(LanguageModelGRPOLossConfig(), "grpo", logits_scale_factor)
+    out_fused, grad_fused, (new_logprobs_fused, _) = loss.combinable_forward_backward(
         split_op(logits, group, -1),
-        target,
-        advantages,
-        old_log_probabilities,
-        grad_logits=local_previous_grad.clone() if accumulate else None,
-        grad_output=grad_output,
-        group=group,
-        logits_scale_factor=logits_scale_factor,
-        num_labels_in_seq=num_labels_in_seq,
-        divisor=divisor,
+        group,
+        local_previous_grad.clone() if accumulate else None,
+        (target, advantages, old_log_probabilities, grad_output, divisor, 0.2, 0.2, num_labels_in_seq, False),
     )
     _compare_losses_and_grads(out_fused, out_ref, grad_output is not None, grad_fused, grad_ref, group=group)
 
@@ -532,7 +538,9 @@ def _test_gspo_loss(
         logits_scale_factor=logits_scale_factor,
         num_labels_in_seq=num_labels_in_seq,
     )
-    _compare_losses_and_grads(out_fused, out_ref, grad_output is not None, grad_fused, grad_ref, group=group)
+    _compare_losses_and_grads(
+        out_fused, out_ref, grad_output is not None, grad_fused, grad_ref, group=group, loss_min_threshold=1e-5
+    )
     Assert.rms_close_relative(new_logprobs_fused, ref_new_logprobs, 1e-5, 1e-6)
 
     if not triton_available:
@@ -551,7 +559,9 @@ def _test_gspo_loss(
         group=group,
         logits_scale_factor=logits_scale_factor,
     )
-    _compare_losses_and_grads(out_triton, out_ref, grad_output is not None, grad_triton, grad_ref, group=group)
+    _compare_losses_and_grads(
+        out_triton, out_ref, grad_output is not None, grad_triton, grad_ref, group=group, loss_min_threshold=1e-5
+    )
     Assert.rms_close_relative(new_logprobs_triton, new_logprobs_fused, 1e-5, 1e-6)
 
 
@@ -565,9 +575,7 @@ def _check_policy_metrics(ref: PolicyMetrics, got: PolicyMetrics, threshold: flo
             Assert.rms_close_relative(got_value, ref_value, threshold, 1e-6)
 
 
-def _test_grpo_metrics(
-    batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, group=None
-):
+def _test_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, group=None):
     logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs(
         num_columns, loss_masking, batch_shape, dtype
     )
@@ -585,7 +593,6 @@ def _test_grpo_metrics(
         epsilon_low=0.2,
         epsilon_high=0.2,
         logits_scale_factor=logits_scale_factor,
-        compute_entropy=compute_entropy,
     )
     got = compute_grpo_metrics(
         split_op(logits, group, -1).contiguous(),
@@ -597,14 +604,11 @@ def _test_grpo_metrics(
         epsilon_high=0.2,
         logits_scale_factor=logits_scale_factor,
         group=group,
-        compute_entropy=compute_entropy,
     )
     _check_policy_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4)
 
 
-def _test_gspo_metrics(
-    batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, num_segments, group=None
-):
+def _test_gspo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, num_segments, group=None):
     logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs(
         num_columns, loss_masking, batch_shape, dtype
     )
@@ -631,7 +635,6 @@ def _test_gspo_metrics(
         epsilon_low=0.2,
         epsilon_high=0.2,
         logits_scale_factor=logits_scale_factor,
-        compute_entropy=compute_entropy,
     )
     got = compute_gspo_metrics(
         split_op(logits, group, -1).contiguous(),
@@ -645,11 +648,19 @@ def _test_gspo_metrics(
         epsilon_high=0.2,
         logits_scale_factor=logits_scale_factor,
         group=group,
-        compute_entropy=compute_entropy,
     )
     _check_policy_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4)
 
 
+def _combinable_loss(config, name: str, logits_scale_factor: float):
+    # Build the loss object so its `combinable_forward_backward` method is exercised directly. The tensor-
+    # parallel `group` is passed per call, so a trivial single-rank distributed config suffices even for the
+    # distributed subtests.
+    distributed_config = DistributedConfig()
+    distributed_config.validate()
+    return config.get_layer(distributed_config, name=name, logits_scale_factor=logits_scale_factor)
+
+
 def _test_z_loss(
     batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None
 ):
@@ -666,13 +677,12 @@ def _test_z_loss(
         previous_grad = torch.randn_like(grad_ref)
         grad_ref = grad_ref + previous_grad
         local_previous_grad = split_op(previous_grad, group, -1).contiguous()
-    out_fused, grad_fused = fused_z_loss_forward_backward(
-        logits=local_logits,
-        loss_mask=loss_mask,
-        grad_logits=local_previous_grad.clone() if accumulate else None,
-        grad_output=grad_output,
-        group=group,
-        logits_scale_factor=logits_scale_factor,
+    loss = _combinable_loss(LanguageModelZLossConfig(), "z_loss", logits_scale_factor)
+    out_fused, grad_fused, _ = loss.combinable_forward_backward(
+        local_logits,
+        group,
+        local_previous_grad.clone() if accumulate else None,
+        (loss_mask, grad_output, local_logits.shape[:-1].numel()),
     )
     _compare_losses_and_grads(
         out_fused,
@@ -680,7 +690,7 @@ def _test_z_loss(
         grad_output is not None,
         grad_fused,
         grad_ref,
-        threshold=1e-5 if data_type == DataType.float32 else 1e-4,
+        threshold=1e-5 if dtype == DataType.float32 else 1e-4,
         group=group,
     )
     if not triton_available:
@@ -700,9 +710,319 @@ def _test_z_loss(
         grad_output is not None,
         grad_triton,
         grad_ref,
-        threshold=1e-5 if data_type == DataType.float32 else 1e-4,
+        threshold=1e-5 if dtype == DataType.float32 else 1e-4,
+        group=group,
+    )
+
+
+def _test_monolithic_loss(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None
+):
+    # A cross-entropy (labels) + z-loss composite sharing one softmax, checked against the same two losses run
+    # standalone. This exercises the shared, tensor-parallel-reduced softmax and the fp32 gradient accumulation
+    # against the already-validated single-loss path.
+    logits, target, _ = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype)
+    local_logits = split_op(logits, group, -1).contiguous()
+    divisor = max(int((target >= 0).sum().item()), 1)
+    children = (
+        _combinable_loss(LanguageModelLabelEntropyLossConfig(), "cross_entropy", logits_scale_factor),
+        _combinable_loss(LanguageModelZLossConfig(), "z_loss", logits_scale_factor),
+    )
+    arguments = ((target, grad_output, divisor), (None, grad_output, local_logits.shape[:-1].numel()))
+    previous_grad = torch.randn_like(local_logits) if accumulate else None
+
+    # Reference: run each loss standalone, accumulating into one gradient buffer.
+    grad_ref = previous_grad.clone() if accumulate else None
+    losses_ref = []
+    for child, child_arguments in zip(children, arguments, strict=True):
+        loss_ref, grad_ref, _ = child.combinable_forward_backward(local_logits, group, grad_ref, child_arguments)
+        losses_ref.append(loss_ref)
+
+    results, grad_fused = _monolithic_core(
+        children, local_logits, group, logits_scale_factor, previous_grad.clone() if accumulate else None, arguments
+    )
+
+    threshold = 1e-5 if dtype == DataType.float32 else 1e-4
+    for (loss_fused, _), loss_ref in zip(results, losses_ref, strict=True):
+        Assert.rms_close_relative(loss_fused, loss_ref, threshold, 1e-6)
+    if grad_output is None:
+        assert grad_fused is None and grad_ref is None
+    else:
+        # The composite sums child gradients in fp32 and casts once; the standalone path casts each child
+        # gradient before adding. In fp16 the two differ by up to a rounding step, so allow a wider abs floor.
+        Assert.rms_close_relative(grad_fused, grad_ref, threshold, 1e-8 if grad_fused.dtype == torch.float32 else 1e-6)
+
+
+def _test_monolithic_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None
+):
+    # The triton monolithic kernel (cross-entropy + z-loss over one softmax) against the sum of the standalone
+    # triton kernels. Both use the same label-count divisor, as the head threads the same divisor to each
+    # combinable child. Tensor-parallel `group` is passed per call for the distributed subtest (case B).
+    if not triton_available:
+        return
+    logits, target, _ = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype)
+    local_logits = split_op(logits, group, -1).contiguous()
+    divisor = max(int((target >= 0).sum().item()), 1)
+    previous_grad = torch.randn_like(local_logits) if accumulate else None
+
+    # Reference: standalone triton cross-entropy then z-loss, accumulating into one gradient buffer.
+    grad_ref = previous_grad.clone() if accumulate else None
+    ce_ref, grad_ref = triton_entropy_loss_forward_backward(
+        local_logits,
+        target,
+        None,
+        grad_logits=grad_ref,
+        grad_output=grad_output,
+        group=group,
+        logits_scale_factor=logits_scale_factor,
+        target_format=TargetFormat.labels,
+        divisor=divisor,
+        block_size=block_size,
+    )
+    z_ref, grad_ref = triton_z_loss_forward_backward(
+        local_logits,
+        None,
+        grad_logits=grad_ref,
+        grad_output=grad_output,
+        group=group,
+        logits_scale_factor=logits_scale_factor,
+        divisor=divisor,
+        block_size=block_size,
+    )
+
+    ce_fused, z_fused, _, _, grad_fused, _, _ = triton_monolithic_loss_forward_backward(
+        local_logits,
+        target.contiguous(),
+        previous_grad.clone() if accumulate else None,
+        logits_scale_factor,
+        group,
+        divisor,
+        ce=(grad_output,),
+        z=(None, grad_output),
+        block_size=block_size,
+    )
+
+    threshold = 1e-5 if dtype == DataType.float32 else 1e-4
+    Assert.rms_close_relative(ce_fused, ce_ref, threshold, 1e-6)
+    Assert.rms_close_relative(z_fused, z_ref, threshold, 1e-6)
+    if grad_output is None:
+        assert grad_fused is None and grad_ref is None
+    else:
+        # The kernel superposes both losses' gradients in fp32 before the single cast; the standalone path
+        # rounds cross-entropy's gradient before z-loss adds to it, so the two differ by up to a rounding step.
+        Assert.rms_close_relative(grad_fused, grad_ref, threshold, 1e-8 if grad_fused.dtype == torch.float32 else 1e-6)
+
+
+def _test_monolithic_single_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None
+):
+    # A single cross-entropy loss through the monolithic kernel. With masking this exercises the masked-row
+    # early-return (no z-loss / GSPO / metrics need the softmax on such a row), which must match the standalone
+    # triton cross-entropy exactly — the skip only avoids work that would have summed to zero.
+    if not triton_available:
+        return
+    logits, target, _ = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype)
+    local_logits = split_op(logits, group, -1).contiguous()
+    divisor = max(int((target >= 0).sum().item()), 1)
+    previous_grad = torch.randn_like(local_logits) if accumulate else None
+
+    ce_ref, grad_ref = triton_entropy_loss_forward_backward(
+        local_logits,
+        target,
+        None,
+        grad_logits=previous_grad.clone() if accumulate else None,
+        grad_output=grad_output,
         group=group,
+        logits_scale_factor=logits_scale_factor,
+        target_format=TargetFormat.labels,
+        divisor=divisor,
+        block_size=block_size,
+    )
+    ce_fused, _, _, _, grad_fused, _, _ = triton_monolithic_loss_forward_backward(
+        local_logits,
+        target.contiguous(),
+        previous_grad.clone() if accumulate else None,
+        logits_scale_factor,
+        group,
+        divisor,
+        ce=(grad_output,),
+        block_size=block_size,
+    )
+    threshold = 1e-5 if dtype == DataType.float32 else 1e-4
+    Assert.rms_close_relative(ce_fused, ce_ref, threshold, 1e-6)
+    if grad_output is None:
+        assert grad_fused is None and grad_ref is None
+    else:
+        Assert.rms_close_relative(grad_fused, grad_ref, threshold, 1e-8 if grad_fused.dtype == torch.float32 else 1e-7)
+
+
+def _test_monolithic_grpo_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None
+):
+    # A single GRPO loss through the monolithic kernel's `grpo` slot (loss / gradient / new-log-probs), against
+    # the python reference. The ce+z monolithic tests never pass `grpo=`, so this is the only low-level cover of
+    # that slot; the object-dispatch driver around it is exercised by the head tests.
+    if not triton_available:
+        return
+    logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs(
+        num_columns, loss_masking, batch_shape, dtype
+    )
+    num_labels = int((target >= 0).sum().item())
+    num_labels_in_seq = torch.where(
+        target >= 0,
+        torch.full(batch_shape, num_labels, dtype=torch.int32, device=target.device),
+        torch.zeros(batch_shape, dtype=torch.int32, device=target.device),
+    )
+    divisor = max(num_labels, 1)
+    out_ref, grad_ref = loss_forward_backward(
+        grad_output,
+        lambda *args, **kwargs: reference_grpo_loss(*args, **kwargs)[0],
+        logits,
+        target,
+        advantages,
+        old_log_probabilities,
+        logits_scale_factor=logits_scale_factor,
+    )
+    if accumulate:
+        previous_grad = torch.randn_like(grad_ref)
+        grad_ref = grad_ref + previous_grad
+        local_previous_grad = split_op(previous_grad, group, -1).contiguous()
+    _, _, grpo_fused, _, grad_fused, _, _ = triton_monolithic_loss_forward_backward(
+        split_op(logits, group, -1).contiguous(),
+        target.contiguous(),
+        local_previous_grad.clone() if accumulate else None,
+        logits_scale_factor,
+        group,
+        divisor,
+        grpo=(advantages, old_log_probabilities, grad_output, 0.2, 0.2, num_labels_in_seq),
+        block_size=block_size,
     )
+    _compare_losses_and_grads(grpo_fused, out_ref, grad_output is not None, grad_fused, grad_ref, group=group)
+
+
+def _test_monolithic_gspo_loss_triton(
+    batch_shape,
+    num_columns,
+    grad_output,
+    logits_scale_factor,
+    loss_masking,
+    dtype,
+    num_segments,
+    accumulate,
+    group=None,
+):
+    # A single GSPO loss through the monolithic path: the eager segment seam produces the per-token backward
+    # coefficient, which the kernel's `gspo_coeff` slot then superposes. Checked against the python reference.
+    # The ce+z monolithic tests never pass `gspo_coeff=`.
+    if not triton_available:
+        return
+    logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs(
+        num_columns, loss_masking, batch_shape, dtype
+    )
+    seq_len = batch_shape[-1] if len(batch_shape) > 1 else batch_shape[0]
+    span = max(seq_len // num_segments, 1)
+    document_index = (torch.arange(seq_len, device=target.device) // span).clamp(max=num_segments - 1)
+    document_index = document_index.expand(batch_shape).contiguous()
+    flat_doc = document_index.reshape(-1).long()
+    labels_per_document = torch.zeros(num_segments, dtype=torch.int32, device=target.device).scatter_add(
+        0, flat_doc, (target.reshape(-1) >= 0).to(torch.int32)
+    )
+    num_labels_in_seq = labels_per_document[flat_doc].reshape(target.shape)
+    out_ref, grad_ref = loss_forward_backward(
+        grad_output,
+        lambda *args, **kwargs: reference_gspo_loss(*args, **kwargs)[0],
+        logits,
+        target,
+        advantages,
+        old_log_probabilities,
+        document_index,
+        num_segments,
+        logits_scale_factor=logits_scale_factor,
+    )
+    if accumulate:
+        previous_grad = torch.randn_like(grad_ref)
+        grad_ref = grad_ref + previous_grad
+        local_previous_grad = split_op(previous_grad, group, -1).contiguous()
+    local_logits = split_op(logits, group, -1).contiguous()
+    # Eager pre-kernel phases (forward reduce + segment seam), as the driver runs them, produce the coefficient.
+    softmax = _monolithic_forward_reduce(local_logits, target.contiguous(), group, logits_scale_factor)
+    max_logits, sum_exp_logits, predicted_logits = softmax
+    new_log_probs = (predicted_logits - max_logits - sum_exp_logits.log()).reshape(target.shape)
+    gspo_loss, _, gspo_coeff = gspo_segment_seam(
+        new_log_probs,
+        target >= 0,
+        advantages,
+        old_log_probabilities,
+        document_index,
+        num_segments,
+        num_labels_in_seq,
+        num_segments,  # divisor
+        grad_output,
+        None,  # sdp_group
+        None,  # sp_group: the vocab-parallel path shards the vocab, not the sequence
+        0.2,
+        0.2,
+        logits_scale_factor,
+    )
+    _, _, _, _, grad_fused, _, _ = triton_monolithic_loss_forward_backward(
+        local_logits,
+        target.contiguous(),
+        local_previous_grad.clone() if accumulate else None,
+        logits_scale_factor,
+        group,
+        1.0,  # GSPO pre-scales its coefficient; no in-kernel divisor
+        gspo_coeff=None if gspo_coeff is None else gspo_coeff.reshape(-1).contiguous(),
+        softmax=softmax,
+    )
+    _compare_losses_and_grads(
+        gspo_loss, out_ref, grad_output is not None, grad_fused, grad_ref, group=group, loss_min_threshold=1e-5
+    )
+
+
+def _test_monolithic_metrics_without_grad_triton(batch_shape, num_columns, loss_masking, dtype, group=None):
+    # A zero-weight policy loss (grad_output=None) that still logs metrics: `compute_metrics=True` with no
+    # gradient must emit the same forward loss / new-log-probs / Σ exp·logits_norm as the with-gradient call,
+    # not crash or skip. The kernel runs the backward re-stream for the metrics alone.
+    if not triton_available:
+        return
+    logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs(
+        num_columns, loss_masking, batch_shape, dtype
+    )
+    num_labels = int((target >= 0).sum().item())
+    num_labels_in_seq = torch.where(
+        target >= 0,
+        torch.full(batch_shape, num_labels, dtype=torch.int32, device=target.device),
+        torch.zeros(batch_shape, dtype=torch.int32, device=target.device),
+    )
+    local_logits = split_op(logits, group, -1).contiguous()
+
+    def run(grad_output):
+        return triton_monolithic_loss_forward_backward(
+            local_logits,
+            target.contiguous(),
+            None,
+            1.0,
+            group,
+            max(num_labels, 1),
+            grpo=(advantages, old_log_probabilities, grad_output, 0.2, 0.2, num_labels_in_seq),
+            compute_metrics=True,
+        )
+
+    _, _, loss_grad, new_logprobs_grad, grad_grad, _, weighted_grad = run(1.0)
+    _, _, loss_nograd, new_logprobs_nograd, grad_nograd, _, weighted_nograd = run(None)
+
+    assert grad_grad is not None and grad_nograd is None
+    Assert.rms_close_relative(loss_nograd, loss_grad, 1e-5, 1e-6)
+    Assert.rms_close_relative(new_logprobs_nograd, new_logprobs_grad, 1e-5, 1e-6)
+    Assert.rms_close_relative(weighted_nograd, weighted_grad, 1e-5, 1e-6)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize("loss_masking", (False, True))
+def test_monolithic_metrics_without_grad_triton(batch_shape, loss_masking):
+    _test_monolithic_metrics_without_grad_triton(batch_shape, 500, loss_masking, DataType.float32)
 
 
 @pytest.mark.slow
@@ -711,7 +1031,7 @@ def _test_z_loss(
     ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
     _LOSS_PARAMETERS,
 )
-@pytest.mark.parametrize("target_format", TargetFormat)
+@pytest.mark.parametrize("target_format", (TargetFormat.labels, TargetFormat.logits))
 @pytest.mark.parametrize("entropy_loss_type", EntropyLossType)
 def test_entropy_loss(
     batch_shape,
@@ -753,6 +1073,46 @@ def test_z_loss(
     )
 
 
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize(
+    ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
+    _LOSS_PARAMETERS,
+)
+def test_monolithic_loss(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+):
+    _test_monolithic_loss(batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize(
+    ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
+    _LOSS_PARAMETERS,
+)
+def test_monolithic_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+):
+    _test_monolithic_loss_triton(
+        batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize(
+    ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
+    _LOSS_PARAMETERS,
+)
+def test_monolithic_single_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+):
+    _test_monolithic_single_loss_triton(
+        batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+    )
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
 @pytest.mark.parametrize(
@@ -802,7 +1162,34 @@ def test_gspo_loss(
     ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
     _LOSS_PARAMETERS,
 )
-@pytest.mark.parametrize("compute_entropy", (False, True))
+def test_monolithic_grpo_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+):
+    _test_monolithic_grpo_loss_triton(
+        batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize(
+    ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "num_segments", "accumulate"),
+    _GSPO_PARAMETERS,
+)
+def test_monolithic_gspo_loss_triton(
+    batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate
+):
+    _test_monolithic_gspo_loss_triton(
+        batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES)
+@pytest.mark.parametrize(
+    ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"),
+    _LOSS_PARAMETERS,
+)
 def test_grpo_metrics(
     batch_shape,
     num_columns,
@@ -812,9 +1199,8 @@ def test_grpo_metrics(
     dtype,
     block_size,
     accumulate,
-    compute_entropy,
 ):
-    _test_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy)
+    _test_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype)
 
 
 @pytest.mark.slow
@@ -823,7 +1209,6 @@ def test_grpo_metrics(
     ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "num_segments", "accumulate"),
     _GSPO_PARAMETERS,
 )
-@pytest.mark.parametrize("compute_entropy", (False, True))
 def test_gspo_metrics(
     batch_shape,
     num_columns,
@@ -833,11 +1218,8 @@ def test_gspo_metrics(
     dtype,
     num_segments,
     accumulate,
-    compute_entropy,
 ):
-    _test_gspo_metrics(
-        batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, num_segments
-    )
+    _test_gspo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, num_segments)
 
 
 @pytest.mark.skip(reason="DPO loss is broken")
@@ -866,7 +1248,7 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa
             suffix = f"{num_columns}-{grad_output}-{logits_scale_factor}-{loss_masking}-{dtype}-{block_size}-{accumulate}-{"_".join([str(i) for i in batch_shape])}"
             # Entropy loss
             for entropy_loss_type in EntropyLossType:
-                for target_format in TargetFormat:
+                for target_format in (TargetFormat.labels, TargetFormat.logits):
                     if target_format == TargetFormat.labels and entropy_loss_type == EntropyLossType.reverse_kl:
                         continue
                     with test_context.subtest(
@@ -917,36 +1299,105 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa
                         accumulate,
                         test_context.group,
                     )
+            # GSPO (tensor-parallel vocab path; segment seam runs eagerly per rank)
+            with test_context.subtest(base_path, f"gspo-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_gspo_loss(
+                        batch_shape,
+                        num_columns,
+                        grad_output,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        4,  # num_segments
+                        accumulate,
+                        test_context.group,
+                    )
             # GRPO metrics
-            for compute_entropy in (False, True):
-                with test_context.subtest(base_path, f"grpo_metrics-{compute_entropy}-{suffix}", 2) as subtest:
-                    if subtest.do_run:
-                        torch.manual_seed((seed + hash(subtest.name)) % 2**32)
-                        _test_grpo_metrics(
-                            batch_shape,
-                            num_columns,
-                            logits_scale_factor,
-                            loss_masking,
-                            dtype,
-                            compute_entropy,
-                            test_context.group,
-                        )
+            with test_context.subtest(base_path, f"grpo_metrics-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_grpo_metrics(
+                        batch_shape,
+                        num_columns,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        test_context.group,
+                    )
             # GSPO metrics
             num_segments = 4
-            for compute_entropy in (False, True):
-                with test_context.subtest(base_path, f"gspo_metrics-{compute_entropy}-{suffix}", 2) as subtest:
-                    if subtest.do_run:
-                        torch.manual_seed((seed + hash(subtest.name)) % 2**32)
-                        _test_gspo_metrics(
-                            batch_shape,
-                            num_columns,
-                            logits_scale_factor,
-                            loss_masking,
-                            dtype,
-                            compute_entropy,
-                            num_segments,
-                            test_context.group,
-                        )
+            with test_context.subtest(base_path, f"gspo_metrics-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_gspo_metrics(
+                        batch_shape,
+                        num_columns,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        num_segments,
+                        test_context.group,
+                    )
+            # Monolithic composite: multiple losses share one tensor-parallel-reduced softmax.
+            with test_context.subtest(base_path, f"monolithic-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_monolithic_loss(
+                        batch_shape,
+                        num_columns,
+                        grad_output,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        accumulate,
+                        test_context.group,
+                    )
+            # Triton monolithic composite: the vocab-parallel reduced-softmax path (case B).
+            with test_context.subtest(base_path, f"monolithic-triton-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_monolithic_loss_triton(
+                        batch_shape,
+                        num_columns,
+                        grad_output,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        block_size,
+                        accumulate,
+                        test_context.group,
+                    )
+            # Triton monolithic GRPO / GSPO: the policy-loss slots over the tensor-parallel reduced softmax.
+            with test_context.subtest(base_path, f"monolithic-grpo-triton-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_monolithic_grpo_loss_triton(
+                        batch_shape,
+                        num_columns,
+                        grad_output,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        block_size,
+                        accumulate,
+                        test_context.group,
+                    )
+            with test_context.subtest(base_path, f"monolithic-gspo-triton-{suffix}", 2) as subtest:
+                if subtest.do_run:
+                    torch.manual_seed((seed + hash(subtest.name)) % 2**32)
+                    _test_monolithic_gspo_loss_triton(
+                        batch_shape,
+                        num_columns,
+                        grad_output,
+                        logits_scale_factor,
+                        loss_masking,
+                        dtype,
+                        4,  # num_segments
+                        accumulate,
+                        test_context.group,
+                    )
 
 
 @pytest.mark.slow
@@ -982,15 +1433,15 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path):
         *(
             f"{entropy_loss_type}-{target_format}"
             for entropy_loss_type in EntropyLossType
-            for target_format in TargetFormat
+            for target_format in (TargetFormat.labels, TargetFormat.logits)
             if target_format != TargetFormat.labels or entropy_loss_type != EntropyLossType.reverse_kl
         ),
         "z_loss",
         "grpo",
-        "grpo_metrics-False",
-        "grpo_metrics-True",
-        "gspo_metrics-False",
-        "gspo_metrics-True",
+        "gspo",
+        "grpo_metrics",
+        "gspo_metrics",
+        "monolithic",
     ),
 )
 def test_lm_loss_distributed(
diff --git a/tests/utils/model_configs.py b/tests/utils/model_configs.py
index 4f2da8b54..63b02e66a 100644
--- a/tests/utils/model_configs.py
+++ b/tests/utils/model_configs.py
@@ -708,7 +708,9 @@ def update_and_add_testing_config(
     # Tests mixture of experts, mixtral converter.
     "llama",
     "llama_grpo",
-    updates={("model", "base_model", "head", "losses"): {"grpo": {"type": "grpo"}}},
+    # Metrics default to `auto` (→ `basic` when pipeline_parallel == 1); pin `none` so this loss-mechanics
+    # config doesn't register the metric family (metrics are covered by the subtests in `test_lm_losses`).
+    updates={("model", "base_model", "head", "losses"): {"grpo": {"type": "grpo", "metrics": "none"}}},
     groups={
         ModelTestingGroup.basic: ModelTestingGroupAction.normal,
         ModelTestingGroup.checkpoint: ModelTestingGroupAction.not_implemented,
@@ -724,7 +726,9 @@ def update_and_add_testing_config(
 update_and_add_testing_config(
     "llama",
     "llama_gspo",
-    updates={("model", "base_model", "head", "losses"): {"gspo": {"type": "gspo"}}},
+    # Metrics default to `auto`; pin `none` for the same reason as `llama_grpo` (keep this config focused
+    # on loss mechanics; metrics are covered by `test_lm_losses`).
+    updates={("model", "base_model", "head", "losses"): {"gspo": {"type": "gspo", "metrics": "none"}}},
     # `ms*` (micro_batch_splits>1) and `ce*` (cross_entropy_splits>1) both produce
     # multiple kernel calls per micro-batch; GSPO's per-document geometric mean can't be
     # reconstructed from per-fragment `exp(mean)` values, so we skip these variants.