[tt-crank] fuse cross entropy onto the ttml cross_entropy_fw/bw composites by agobeljicTT · Pull Request #9324 · tenstorrent/tt-mlir

agobeljicTT · 2026-09-15T18:21:56Z

Ticket

Problem description

We have cross_entropy_fw and cross_entropy_bw in ttml, one kernel each way, and we want the loss to use them.

before, rows=1023 C=128256	after
`aten._log_softmax` `[1023, 128256]` f32, 500 MiB	`ttnn.cross_entropy_fw(logits, target)`
`aten.nll_loss_forward`
`aten.nll_loss_backward` `[1023, 128256]`	`ttnn.cross_entropy_bw(grad, logits, target)`
`aten._log_softmax_backward_data` `[1023, 128256]` f32
log probabilities saved for backward	logits, clamped target and row mask saved

What's changed

We inject a custom torch op for F.cross_entropy and lower it onto the ttml kernels.

Note

We swap the F.cross_entropy call itself, not a pattern of decomposed ops. A hand written _log_softmax plus nll_loss has no such call to swap, so it keeps the aten lowerings from #9339 and is not fused.

Checklist

New/Existing tests provide coverage for changes: tests/python/test_cross_entropy.py, swapped and aten paths at OPT_LEVEL 1 and OPT_LEVEL 0 against CPU (mean/sum, ignore_index, off tile shapes, 8k classes, none and f32 fallbacks), a promotion check on the TTIR and TTNN dumps, and edge cases (large logits, a batch where every row is ignored, logits that are not finite in ignored rows, rank 1 rejection). OPT_LEVEL 1 cases are xfail(run=False) under ttsim.

Copilot

🟡 Changes recommended

Unresolved correctness and integration issues remain in the fallback lowerings and partition hook.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds fused cross-entropy forward/backward composites, fallback lowerings, native bindings, and coverage across optimization levels.

Changes:

Adds composite builders and TTNN promotion/decomposition support.
Rewrites supported cross-entropy graphs with fallback handling.
Adds fusion, fallback, and artifact-based tests.

File summaries

File	Reviewed changes	Findings
`tt-crank/tests/python/test_cross_entropy.py`	Cross-entropy fusion and fallback coverage	—
`tt-crank/src/torch/ops/builders.hpp`	Builder exports	—
`tt-crank/src/torch/_native.cpp`	Native builder bindings	—
`tt-crank/src/engine/ttir_module_builder.hpp`	Builder API declarations	—
`tt-crank/src/engine/ttir_module_builder.cpp`	Composite construction and decomposition	—
`tt-crank/python/tt_crank/torch/_compile.py`	Fusion, fallback lowerings, and partition integration	Critical: support or reject N-D inputs correctly. Moderate: use stable log-softmax computation, handle weighted cross-entropy fallback, and preserve the GQA rewrite during partition integration.

Review details

Suppressed comments (4)

tt-crank/python/tt_crank/torch/_compile.py:1012

For mean reduction with every row set to ignore_index, total_weight is zero. Dividing by it and then multiplying the resulting infinity by the all-zero one-hot mask yields NaN gradients, while PyTorch leaves ignored-row gradients at zero. Guard the zero-count case before the division (and preserve the NaN forward loss).

    grad = (
        mb.div(grad_output, total_weight)
        if reduction == _NLL_REDUCTION_MEAN
        else grad_output
    )

tt-crank/python/tt_crank/torch/_compile.py:1103

The fused mean path has the same all-ignored edge case: grad / total_weight becomes infinite when every row is ignored, so cross_entropy_bw can generate NaNs before the later validity mask zeros the rows. Guard the denominator or feed a zero gradient to the kernel when total_weight == 0, matching PyTorch's zero gradients for ignored targets.

                    if reduction == _NLL_REDUCTION_MEAN:
                        grad = emit(_aten.div.Tensor, grad, total_weight)
                    d_logits = emit(
                        _aten.mul.Tensor,
                        emit(_CROSS_ENTROPY_BW, grad, logits, clamped),

tt-crank/python/tt_crank/torch/_compile.py:975

The PR description says cases outside the fused kernel use the new ATen fallback, but this lowering still raises when weight is provided. A valid weighted F.cross_entropy therefore fails compilation instead of taking the fallback; implement class-weight handling in the forward/backward lowerings or narrow the advertised fallback scope.

@_lowering(_aten.nll_loss_forward.default)
@_skip_prepare(_aten.nll_loss_forward.default)
@_needs_node
def _(mb, node, log_probs, target, weight, reduction, ignore_index):
    if weight is not None:
        raise NotImplementedError("tt-crank nll_loss: class weights are not supported")

tt-crank/python/tt_crank/torch/_compile.py:982

For a rank-greater-than-two fallback, rows and valid retain every non-class dimension, but both reductions here only reduce dimension 0. reduction="sum"/"mean" must produce a scalar and total_weight must count all valid target elements; with NCHW inputs this leaves an HxW tensor and changes the result shape and normalization.

    rows = mb.neg(mb.sum(mb.mul(log_probs, onehot), [1], False))
    total_weight = mb.sum(valid, [0], False)
    if reduction == _NLL_REDUCTION_NONE:
        return rows, total_weight
    loss = mb.sum(rows, [0], False)

Files reviewed: 6/6 changed files
Comments generated: 3
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    classes = mb.arange(0, log_probs.shape[1], 1, _to_runtime_dtype(torch.int32))
+    onehot = mb.typecast(
+        mb.eq(
+            mb.unsqueeze(mb.typecast(target, _to_runtime_dtype(torch.int32)), 1),
+            classes,


+@_lowering(_aten._log_softmax.default)
+def _(mb, x, dim, half_to_float=False):
+    return mb.log(mb.softmax(x, int(dim)))


+    if _fuse_cross_entropy(joint_module.graph):
+        joint_module.recompile()
+    return default_partition(joint_module, joint_inputs, **kwargs)


codecov-commenter · 2026-09-15T18:36:14Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.82%. Comparing base (78b7044) to head (ede8e2e).
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #9324      +/-   ##
==========================================
- Coverage   73.83%   73.82%   -0.01%     
==========================================
  Files         515      515              
  Lines      108129   108129              
==========================================
- Hits        79834    79826       -8     
- Misses      28295    28303       +8

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

…ck) (#9339) ### Ticket Tracking: tenstorrent/tt-blacksmith#685 ### Problem description `F.cross_entropy` decomposes to `aten._log_softmax` + `aten.nll_loss_forward` (backward: `nll_loss_backward` + `_log_softmax_backward_data`). None of them lowered, so single-chip Llama training did not compile. ### What's changed Lowerings for the four ops: `[rows x C]` log-probabilities, integer targets with `ignore_index`, mean/sum/none, no class weights. `_log_softmax` is max-shifted log-sum-exp (plain `log(softmax)` hits -inf past a logit gap of ~87). The mean backward clamps the row count to >= 1 so an all-ignored batch gives zero gradients like torch. Stacked on the max PR for `ModuleBuilder.max`. The fused `cross_entropy_fw/bw` path (#9324) rewrites this pattern onto the ttml kernels and keeps these as its fallback. ### Checklist - [x] New/Existing tests provide coverage for changes: `tests/python/test_cross_entropy.py` (5 reductions/dtypes vs CPU, large-logit finiteness, all-ignored batch, rank-1 rejection)

Emits a ttcore.composite together with a private decomposition function built by the caller through the usual build_* helpers. tt-mlir promotes the composite to its typed op when the name is in the resolver registry and inlines the decomposition otherwise (or at OPT_LEVEL 0), so one builder serves both the kernel and the fallback path. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

agobeljicTT · 2026-09-21T08:46:18Z

    options: dict[CompileOption, str | int | bool] | None = None,
 ):
    """Top-level dynamo backend. Delegates to aot_module_simplified."""
+    rewrite_cross_entropy(gm)


good place to have a wrapper rewrite_custom_ops and group all new thing, but we can let it be like this for now

agobeljicTT · 2026-09-21T08:47:28Z

+            if average:
+                t.copy_(t.div(self.size()))


@pilkicTT @acolicTT I know this deserves a dedicated PR but I am really tired of splitting stuff xd and this is a micro change and can be inlined here

pilkicTT

i don't like the rewriting of the FX graph and believe that this can be done in a more native way to torch, let's sync...

pilkicTT · 2026-09-30T10:50:32Z

    options: dict[CompileOption, str | int | bool] | None = None,
 ):
    """Top-level dynamo backend. Delegates to aot_module_simplified."""
+    rewrite_cross_entropy(gm)


this is not ok - i think we should be able to implement proper AutogradPrivateUse1 kernels and avoid rewriting the FX graph ourselves, but let torch do it... we can sync offline...

Copilot AI lite review requested due to automatic review settings September 15, 2026 18:21

agobeljicTT requested review from acolicTT, pilkicTT and umalesTT as code owners September 15, 2026 18:21

Copilot started reviewing on behalf of agobeljicTT September 15, 2026 18:22 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

agobeljicTT and others added 4 commits September 18, 2026 17:07

[tt-crank] test: same-name composites get distinct decomposition symbols

946377d

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

add matmul + composite test

47ebc59

init

82f4e92

agobeljicTT force-pushed the agobeljicTT/crank-xentropy branch from a97281b to 82f4e92 Compare September 19, 2026 18:02

agobeljicTT changed the base branch from main to agobeljicTT/crank-composites September 19, 2026 18:02

agobeljicTT added 4 commits September 19, 2026 20:53

fix multichip

2f37ff7

fake xentropy

14f23e6

sharding move

1d92feb

split

8213a3b

agobeljicTT commented Sep 21, 2026

View reviewed changes

tests

ed59616

Base automatically changed from agobeljicTT/crank-composites to main September 25, 2026 07:40

agobeljicTT added 2 commits September 29, 2026 11:14

Merge branch 'main' into agobeljicTT/crank-xentropy

f6e8f29

Merge branch 'main' into agobeljicTT/crank-xentropy

0c1fd60

agobeljicTT mentioned this pull request Sep 29, 2026

[Bringup] Single-chip dense LLMs to 20% MFU tenstorrent/tt-blacksmith#685

Closed

pilkicTT requested changes Sep 30, 2026

View reviewed changes

refactor

e7d634d

agobeljicTT added 2 commits October 7, 2026 10:47

Merge branch 'main' into agobeljicTT/crank-xentropy

6b58f1f

fp32

ebb2722

pilkicTT approved these changes Oct 13, 2026

View reviewed changes

agobeljicTT added 2 commits October 13, 2026 13:29

Merge branch 'main' into agobeljicTT/crank-xentropy

836d1bb

fix

ede8e2e

agobeljicTT merged commit 240e14d into main Oct 13, 2026
28 checks passed

agobeljicTT deleted the agobeljicTT/crank-xentropy branch October 13, 2026 12:43

Navigation Menu

Repository navigation

[tt-crank] fuse cross entropy onto the ttml cross_entropy_fw/bw composites - #9324

[tt-crank] fuse cross entropy onto the ttml cross_entropy_fw/bw composites#9324
agobeljicTT merged 16 commits into
mainfrom
agobeljicTT/crank-xentropy

agobeljicTT commented Sep 15, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

codecov-commenter commented Sep 15, 2026 •

edited

Loading

Uh oh!

agobeljicTT Sep 21, 2026

Uh oh!

agobeljicTT Sep 21, 2026

Uh oh!

pilkicTT left a comment

Uh oh!

pilkicTT Sep 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

agobeljicTT commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Ticket

Problem description

What's changed

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (4)

Uh oh!

codecov-commenter commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

agobeljicTT Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

agobeljicTT Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

pilkicTT left a comment

Choose a reason for hiding this comment

Uh oh!

pilkicTT Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

agobeljicTT commented Sep 15, 2026 •

edited

Loading

codecov-commenter commented Sep 15, 2026 •

edited

Loading