Skip to content

fix(cutedsl): preserve FP32 CE output precision - #1372

Open
justinhh4 wants to merge 2 commits into
linkedin:mainfrom
justinhh4:justinhh4/cutedsl-ce-fp32-feature-outputs
Open

fix(cutedsl): preserve FP32 CE output precision#1372
justinhh4 wants to merge 2 commits into
linkedin:mainfrom
justinhh4:justinhh4/cutedsl-ce-fp32-feature-outputs

Conversation

@justinhh4

@justinhh4 justinhh4 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

  • force FP32 CE loss and z-loss stores through explicitly typed FP32 pointers
  • prevent CuTe DSL 4.6 from narrowing feature-enabled FP32 training outputs to bf16
  • preserve the existing single-kernel, in-place-gradient design

Problem

With nvidia-cutlass-dsl==4.6.0, FP32 HAS_GRAD specializations that enable z-loss, softcap, or label smoothing store the public loss outputs at bf16 precision even though the CE math and gradients remain FP32-accurate. For example, the z-loss test produced 8.7725 instead of the Triton/PyTorch result 8.8014; per-token values landed exactly on bf16 quantization points.

Rebuilding the loss and z-loss views with explicit Float32 pointers keeps those stores FP32. FP16/BF16 paths retain their existing typed stores. This does not add launches, buffers, or API changes, and FLCE is unchanged.

Validation

  • python -m pytest --confcutdir=test/cutedsl test/cutedsl/test_cross_entropy.py -q — 158 passed, 1 environment-gated skip
  • LIGER_KERNEL_IMPL=cutedsl python -m pytest --confcutdir=test/cutedsl test/cutedsl/test_cross_entropy.py::test_liger_kernel_impl_cutedsl_selects_cutedsl_ce -q — 1 passed
  • ruff check src/liger_kernel/ops/cutedsl/ops/cross_entropy.py
  • ruff format --check src/liger_kernel/ops/cutedsl/ops/cross_entropy.py

Justin Hu and others added 2 commits August 11, 2026 18:39
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@justinhh4 justinhh4 changed the title fix(cutedsl): preserve FP32 CE feature outputs fix(cutedsl): preserve FP32 CE output precision Aug 11, 2026
@justinhh4
justinhh4 marked this pull request as ready for review August 11, 2026 19:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant