Skip to content

[NPU]: optimize cross entropy kernel gradient computation - #1178

Merged
Tcc0403 merged 1 commit into
linkedin:mainfrom
noemotiovon:cross_entropy_opt
Apr 1, 2026
Merged

[NPU]: optimize cross entropy kernel gradient computation#1178
Tcc0403 merged 1 commit into
linkedin:mainfrom
noemotiovon:cross_entropy_opt

Conversation

@noemotiovon

@noemotiovon noemotiovon commented Apr 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor
  • Hoist loop-invariant scalar computations (z_scale, one_minus_ls, z_deriv) out of the inner loop to avoid redundant recalculation
  • Fuse softmax, z-loss derivative, and smoothing term into a single vector expression in the non-weighted branch
  • Guard tl.where with block-range check (y >= i and y < i + BLOCK_SIZE) to skip unnecessary vector operations when target index is not in the current block

Hardware Type: Type: Atlas 800I A2

  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@noemotiovon
[NPU]: optimize cross entropy kernel gradient computation
16d23e6 
- Hoist loop-invariant scalar computations (z_scale, one_minus_ls,
  z_deriv) out of the inner loop to avoid redundant recalculation
- Fuse softmax, z-loss derivative, and smoothing term into a single
  vector expression in the non-weighted branch
- Guard tl.where with block-range check (y >= i and y < i + BLOCK_SIZE)
  to skip unnecessary vector operations when target index is not in the
  current block
@noemotiovon

noemotiovon commented Apr 1, 2026

Copy link
Copy Markdown
Contributor Author

Accuracy Test:

========================================================== test session starts ===========================================================
platform linux -- Python 3.11.14, pytest-9.0.2, pluggy-1.6.0
rootdir: /home/lichenguang25/github/Liger-Kernel
configfile: pyproject.toml
plugins: anyio-4.12.0, cov-7.0.0, rerunfailures-16.1, xdist-3.8.0, forked-1.6.0, asyncio-1.3.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 177 items                                                                                                                      

test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-sum-2-4096-32000] PASSED                           [  0%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-sum-3-423-32000] PASSED                            [  1%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-mean-2-4096-32000] PASSED                          [  1%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-mean-3-423-32000] PASSED                           [  2%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-none-2-4096-32000] PASSED                          [  2%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype0-1e-08-0.05-none-3-423-32000] PASSED                           [  3%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-sum-2-4096-32000] PASSED                          [  3%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-sum-3-423-32000] PASSED                           [  4%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-mean-2-4096-32000] PASSED                         [  5%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-mean-3-423-32000] PASSED                          [  5%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-none-2-4096-32000] PASSED                         [  6%]
test/transformers/test_cross_entropy.py::test_correctness[1.0-dtype1-1e-08-1e-06-none-3-423-32000] PASSED                          [  6%]
test/transformers/test_cross_entropy.py::test_correctness_functional[1.0-dtype0-1e-08-0.05-2-2-8] PASSED                           [  7%]
test/transformers/test_cross_entropy.py::test_correctness_functional[1.0-dtype0-1e-08-0.05-9-7-41] PASSED                          [  7%]
test/transformers/test_cross_entropy.py::test_correctness_functional[1.0-dtype1-1e-08-1e-06-2-2-8] PASSED                          [  8%]
test/transformers/test_cross_entropy.py::test_correctness_functional[1.0-dtype1-1e-08-1e-06-9-7-41] PASSED                         [  9%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-sum-2-4096-32000-2] PASSED       [  9%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-sum-3-423-32000--123] PASSED     [ 10%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-mean-2-4096-32000-2] PASSED      [ 10%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-mean-3-423-32000--123] PASSED    [ 11%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-none-2-4096-32000-2] PASSED      [ 11%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype0-1e-08-0.05-none-3-423-32000--123] PASSED    [ 12%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-sum-2-4096-32000-2] PASSED      [ 12%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-sum-3-423-32000--123] PASSED    [ 13%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-mean-2-4096-32000-2] PASSED     [ 14%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-mean-3-423-32000--123] PASSED   [ 14%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-none-2-4096-32000-2] PASSED     [ 15%]
test/transformers/test_cross_entropy.py::test_correctness_with_ignore_index[1.0-dtype1-1e-08-1e-06-none-3-423-32000--123] PASSED   [ 15%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_once[1.0-dtype0-1e-08-0.05-2-4096-32000-0.1] PASSED [ 16%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_once[1.0-dtype0-1e-08-0.05-3-423-32000-0.1] PASSED  [ 16%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_once[1.0-dtype1-1e-08-1e-06-2-4096-32000-0.1] PASSED [ 17%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_once[1.0-dtype1-1e-08-1e-06-3-423-32000-0.1] PASSED [ 18%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_with_ignore_index_once[1.0-dtype0-1e-08-0.05-2-4096-32000-1-0.1] PASSED [ 18%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_with_ignore_index_once[1.0-dtype0-1e-08-0.05-3-423-32000--300-0.2] PASSED [ 19%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_with_ignore_index_once[1.0-dtype1-1e-08-1e-06-2-4096-32000-1-0.1] PASSED [ 19%]
test/transformers/test_cross_entropy.py::test_correctness_with_label_smoothing_with_ignore_index_once[1.0-dtype1-1e-08-1e-06-3-423-32000--300-0.2] PASSED [ 20%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-sum-2-4096-32000-30.0] PASSED    [ 20%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-sum-3-423-32000-30.0] PASSED     [ 21%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-mean-2-4096-32000-30.0] PASSED   [ 22%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-mean-3-423-32000-30.0] PASSED    [ 22%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-none-2-4096-32000-30.0] PASSED   [ 23%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype0-1e-08-0.05-none-3-423-32000-30.0] PASSED    [ 23%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-sum-2-4096-32000-30.0] PASSED   [ 24%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-sum-3-423-32000-30.0] PASSED                 [ 24%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-mean-2-4096-32000-30.0] PASSED               [ 25%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-mean-3-423-32000-30.0] PASSED                [ 25%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-none-2-4096-32000-30.0] PASSED                          [ 26%]
test/transformers/test_cross_entropy.py::test_correctness_with_softcap_once[1.0-dtype1-1e-08-1e-06-none-3-423-32000-30.0] PASSED                                       [ 27%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-True-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED                                      [ 27%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-True-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED                                       [ 28%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-True-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED                                     [ 28%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-True-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED                                      [ 29%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-False-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED                                     [ 29%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-False-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED                                      [ 30%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-False-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED                                    [ 31%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[0.0001-False-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED                                     [ 31%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-True-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED                                       [ 32%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-True-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED                                        [ 32%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-True-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED                                      [ 33%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-True-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED                                       [ 33%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-False-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED                                      [ 34%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-False-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED                                       [ 35%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-False-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED                                     [ 35%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_once[1e-05-False-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED                                      [ 36%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-True-0.0001-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED        [ 36%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-True-0.0001-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED         [ 37%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-True-0.0001-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED       [ 37%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-True-0.0001-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED        [ 38%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-False-1e-05-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED        [ 38%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-False-1e-05-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED         [ 39%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-False-1e-05-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED       [ 40%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.1-42-mean-False-1e-05-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED        [ 40%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-True-0.0001-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED        [ 41%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-True-0.0001-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED         [ 41%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-True-0.0001-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED       [ 42%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-True-0.0001-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED        [ 42%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-False-1e-05-1.0-dtype0-1e-08-0.05-2-4096-32000] PASSED        [ 43%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-False-1e-05-1.0-dtype0-1e-08-0.05-3-423-32000] PASSED         [ 44%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-False-1e-05-1.0-dtype1-1e-08-1e-06-2-4096-32000] PASSED       [ 44%]
test/transformers/test_cross_entropy.py::test_correctness_with_z_loss_with_other_params_once[0.2--42-sum-False-1e-05-1.0-dtype1-1e-08-1e-06-3-423-32000] PASSED        [ 45%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-sum-2-4096-32000] PASSED                                              [ 45%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-sum-3-423-32000] PASSED                                               [ 46%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-mean-2-4096-32000] PASSED                                             [ 46%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-mean-3-423-32000] PASSED                                              [ 47%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-none-2-4096-32000] PASSED                                             [ 48%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype0-1e-08-0.05-none-3-423-32000] PASSED                                              [ 48%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-sum-2-4096-32000] PASSED                                             [ 49%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-sum-3-423-32000] PASSED                                              [ 49%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-mean-2-4096-32000] PASSED                                            [ 50%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-mean-3-423-32000] PASSED                                             [ 50%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-none-2-4096-32000] PASSED                                            [ 51%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_once[1.0-dtype1-1e-08-1e-06-none-3-423-32000] PASSED                                             [ 51%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-sum-2-4096-32000] PASSED       [ 52%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-sum-3-423-32000] PASSED        [ 53%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-mean-2-4096-32000] PASSED      [ 53%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-mean-3-423-32000] PASSED       [ 54%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-none-2-4096-32000] PASSED      [ 54%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05--100-0.0001-0.1-30.0-none-3-423-32000] PASSED       [ 55%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-sum-2-4096-32000] PASSED          [ 55%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-sum-3-423-32000] PASSED           [ 56%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-mean-2-4096-32000] PASSED         [ 57%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-mean-3-423-32000] PASSED          [ 57%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-none-2-4096-32000] PASSED         [ 58%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype0-1e-08-0.05-42-1e-05-0.2-40.0-none-3-423-32000] PASSED          [ 58%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-sum-2-4096-32000] PASSED      [ 59%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-sum-3-423-32000] PASSED       [ 59%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-mean-2-4096-32000] PASSED     [ 60%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-mean-3-423-32000] PASSED      [ 61%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-none-2-4096-32000] PASSED     [ 61%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06--100-0.0001-0.1-30.0-none-3-423-32000] PASSED      [ 62%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-sum-2-4096-32000] PASSED         [ 62%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-sum-3-423-32000] PASSED          [ 63%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-mean-2-4096-32000] PASSED        [ 63%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-mean-3-423-32000] PASSED         [ 64%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-none-2-4096-32000] PASSED        [ 64%]
test/transformers/test_cross_entropy.py::test_correctness_with_weight_with_other_params_once[1.0-dtype1-1e-08-1e-06-42-1e-05-0.2-40.0-none-3-423-32000] PASSED         [ 65%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype0-1e-08-0.05-sum-2-4096-32000] PASSED                                                [ 66%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype0-1e-08-0.05-sum-3-423-32000] PASSED                                                 [ 66%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype0-1e-08-0.05-mean-2-4096-32000] PASSED                                               [ 67%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype0-1e-08-0.05-mean-3-423-32000] PASSED                                                [ 67%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype1-1e-08-1e-06-sum-2-4096-32000] PASSED                                               [ 68%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype1-1e-08-1e-06-sum-3-423-32000] PASSED                                                [ 68%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype1-1e-08-1e-06-mean-2-4096-32000] PASSED                                              [ 69%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer[1.0-dtype1-1e-08-1e-06-mean-3-423-32000] PASSED                                               [ 70%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-sum--100-0.0001-0.1-30.0-2-1024-32000] PASSED         [ 70%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-sum--100-0.0001-0.1-30.0-3-423-32000] PASSED          [ 71%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-sum-42-1e-05-0.2-40.0-2-1024-32000] PASSED            [ 71%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-sum-42-1e-05-0.2-40.0-3-423-32000] PASSED             [ 72%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-mean--100-0.0001-0.1-30.0-2-1024-32000] PASSED        [ 72%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-mean--100-0.0001-0.1-30.0-3-423-32000] PASSED         [ 73%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-mean-42-1e-05-0.2-40.0-2-1024-32000] PASSED           [ 74%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype0-1e-08-0.05-mean-42-1e-05-0.2-40.0-3-423-32000] PASSED            [ 74%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-sum--100-0.0001-0.1-30.0-2-1024-32000] PASSED        [ 75%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-sum--100-0.0001-0.1-30.0-3-423-32000] PASSED         [ 75%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-sum-42-1e-05-0.2-40.0-2-1024-32000] PASSED           [ 76%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-sum-42-1e-05-0.2-40.0-3-423-32000] PASSED            [ 76%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-mean--100-0.0001-0.1-30.0-2-1024-32000] PASSED       [ 77%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-mean--100-0.0001-0.1-30.0-3-423-32000] PASSED        [ 77%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-mean-42-1e-05-0.2-40.0-2-1024-32000] PASSED          [ 78%]
test/transformers/test_cross_entropy.py::test_correctness_not_last_layer_with_other_params[1.0-dtype1-1e-08-1e-05-mean-42-1e-05-0.2-40.0-3-423-32000] PASSED           [ 79%]
test/transformers/test_cross_entropy.py::test_float32_internal PASSED                                                                                                  [ 79%]
test/transformers/test_cross_entropy.py::test_correctness_with_out_of_bounds_target_once[2-4096-32000-2] PASSED                                                        [ 80%]
test/transformers/test_cross_entropy.py::test_correctness_with_out_of_bounds_target_once[3-423-32000--123] PASSED                                                      [ 80%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-mean-2-4096-32000--100] PASSED                                    [ 81%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-mean-3-423-32000-2] PASSED                                        [ 81%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-sum-2-4096-32000--100] PASSED                                     [ 82%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-sum-3-423-32000-2] PASSED                                         [ 83%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-none-2-4096-32000--100] PASSED                                    [ 83%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype0-1.0-0.0001-0.0001-none-3-423-32000-2] PASSED                                        [ 84%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-mean-2-4096-32000--100] PASSED                                        [ 84%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-mean-3-423-32000-2] PASSED                                            [ 85%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-sum-2-4096-32000--100] PASSED                                         [ 85%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-sum-3-423-32000-2] PASSED                                             [ 86%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-none-2-4096-32000--100] PASSED                                        [ 87%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype1-1.0-0.01-0.01-none-3-423-32000-2] PASSED                                            [ 87%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-mean-2-4096-32000--100] PASSED                                        [ 88%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-mean-3-423-32000-2] PASSED                                            [ 88%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-sum-2-4096-32000--100] PASSED                                         [ 89%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-sum-3-423-32000-2] PASSED                                             [ 89%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-none-2-4096-32000--100] PASSED                                        [ 90%]
test/transformers/test_cross_entropy.py::test_correctness_with_forward_only[dtype2-1.0-0.01-0.01-none-3-423-32000-2] PASSED                                            [ 90%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[False-False-False] PASSED                                                          [ 91%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[True-False-False] PASSED                                                           [ 92%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[False-True-False] PASSED                                                           [ 92%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[False-False-True] PASSED                                                           [ 93%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[True-True-False] PASSED                                                            [ 93%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[True-False-True] PASSED                                                            [ 94%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[False-True-True] PASSED                                                            [ 94%]
test/transformers/test_cross_entropy.py::test_liger_cross_entropy_structured_output[True-True-True] PASSED                                                             [ 95%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype0--100-2-128-512] PASSED                                                          [ 96%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype0--100-3-47-31] PASSED                                                            [ 96%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype0-2-2-128-512] PASSED                                                             [ 97%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype0-2-3-47-31] PASSED                                                               [ 97%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype1--100-2-128-512] PASSED                                                          [ 98%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype1--100-3-47-31] PASSED                                                            [ 98%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype1-2-2-128-512] PASSED                                                             [ 99%]
test/transformers/test_cross_entropy.py::test_correctness_with_predicted_tokens[dtype1-2-3-47-31] PASSED                                                               [100%]

@noemotiovon

noemotiovon commented Apr 1, 2026

Copy link
Copy Markdown
Contributor Author

Benchmark：

During the testing process, since the device I used only has 32 GB of Global Memory, I reduced the batch size in the benchmark as follows:

"extra_benchmark_configs": [{"B": 4, "T": 2048}]

Before

Speed:
image
Memory:
image

After:

Speed:
image
Memory:
image

@noemotiovon

noemotiovon commented Apr 1, 2026

Copy link
Copy Markdown
Contributor Author

Hi @Tcc0403, could you help review my code?

Tcc0403
Tcc0403 approved these changes Apr 1, 2026
View reviewed changes

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you, LGTM

@Tcc0403
Tcc0403 added this pull request to the merge queue Apr 1, 2026
Merged via the queue into linkedin:main with commit 35cf6a0 Apr 1, 2026
5 of 7 checks passed
@zheliuyu zheliuyu mentioned this pull request Apr 14, 2026
Open
41 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Tcc0403 Tcc0403 Tcc0403 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@noemotiovon @Tcc0403