Skip to content

[Bug] FLCE Triton Fallback - #1396

Open
buffett0323 wants to merge 2 commits into
linkedin:mainfrom
buffett0323:flce-triton-fallback
Open

[Bug] FLCE Triton Fallback#1396
buffett0323 wants to merge 2 commits into
linkedin:mainfrom
buffett0323:flce-triton-fallback

Conversation

@buffett0323

Copy link
Copy Markdown
Contributor

Summary

The CuTe DSL backend's LigerFusedLinearCrossEntropyFunction currently raises a hard RuntimeError on capability misses (e.g., non-SM100 hardware, FP32, or unsupported reductions) instead of falling back to Triton. This breaks fused-linear-cross-entropy training across all models when using the cutedsl backend on non-SM100 hardware (reproducible on H100 via LIGER_KERNEL_IMPL=cutedsl python -m pytest test/convergence/bf16/test_mini_models.py -k "mini_llama3 or mini_qwen2").

This PR makes capability misses fall back to the Triton kernel with a one-time warning. Found while validating #1390 against both backends, not specific to any model. The README already sets this expectation for ops with no CuTe DSL kernel, this PR extends it to an op that has one which declines the current hardware.

Details

Dispatch Implementation

  • LigerFusedLinearCrossEntropyFunction is now a thin dispatcher rather than a torch.autograd.Function. The original autograd logic is renamed to _NativeFusedLinearCrossEntropyFunction.
  • Call choice must occur before autograd records the node. Running Triton's .apply inside an autograd forward() would execute outside the graph and silently drop gradients.
  • Call sites and __module__ assertions remain unchanged.
  • Validation executes before dispatch to guarantee uniform error messages across different GPU architectures.
  • Low-level entry point fused_linear_cross_entropy_forward still raises on misses; only the dispatcher degrades.

Testing Done

Validated across H100 (SM90) and B200 (SM100) to ensure both fallback and native execution paths work properly.

Test Suite H100 (SM90) B200 (SM100)
test/cutedsl FLCE + convergence (mini_llama3, mini_qwen2, muse_glimmer) 26 PASSED 26 PASSED
Full test/cutedsl/test_fused_linear_cross_entropy.py (incl. @sm100_required) 240 PASSED

Test Changes

  • 3 tests now asset Triton fallback instead of a hard raise (test_flce_falls_back_to_triton_on_non_sm100, ..._for_fp32, ..._for_reduction_none).

  • test_flce_uses_native_path_when_supported (@sm100_required) to prevent accidental permanent fallback.

  • test_flce_native_forward_still_raises_on_capability_miss to check the low-level entry point.

  • test_cutedsl_flce_integration_does_not_raise_on_capability_miss to verify the end-to-end integration path via a subprocess.

  • run make test to ensure correctness

  • run make checkstyle to ensure code style

  • run make test-convergence to ensure convergence

@buffett0323

Copy link
Copy Markdown
Contributor Author

Hi @PKUWZP , this is the pr extension for #1390 , please take a look, thanks!
Also, cuTile has a similar constraint (requires SM90/BF16 and raises). Should we address it here or in a follow-up PR?

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