[SuperGoodGame]

Summary

Resolves #1621.

The Triton FLCE synchronizes the host twice per step on values the computation does not need. With activation offload on a side stream, the forward sync blocks the host until the whole device-to-host backlog drains, so the loss head and the backward start late. This PR drops both syncs where possible. Losses and gradients are unchanged.

• Forward: read the non-ignored count and the target-weight sum on the host only for reduction="mean", the only reduction that divides by them. Read the class-weight sum only when label_smoothing > 0; the kernel multiplies it by eps = label_smoothing / V, so it has no effect otherwise. Token accuracy now divides by the count on the device. For "sum" and "none", the kernel and the cross_entropy_loss_and_grad implementations (Triton, cutedsl) do not read the placeholder values.
• Backward: always launch the scaling kernels and move the grad_output == 1.0 check onto the device. element_mul_kernel gets a SKIP_IF_ONE constexpr. It defaults to False, so every other caller compiles the same kernel as before. When grad_output is 1.0, each program returns right after loading it; otherwise it scales as before.

Details

• reduction="mean" still syncs. Avoiding that needs the count passed to the kernel as a device pointer, so it is left for a follow-up.
• use_token_scaling keeps a data-dependent boolean index that also syncs. It is out of scope here.
• cross_entropy_backward has the same torch.equal check, but there it also sets grad_scaled_in_place, which guards retain_graph=True, so I left it unchanged. The cutedsl, cutile and ascend backends are untouched.
• This overlaps with perf(flce): skip the projection GEMMs for ignored rows #1554, which rewrites the same forward line (.item() -> nonzero()). I'm happy to rebase on whichever lands first. If row compaction is kept, it could run only when the count is needed anyway (reduction="mean"), or behind an opt-in flag, so that "sum" stays sync-free.

Results on one H200:

	main (42dc23c)	this PR
repro from #1621: step time (3 runs)	207.5 / 209.5 / 207.6 ms	199.6 / 199.8 / 202.3 ms
repro from #1621: host time blocked inside the FLCE forward	34.2 ms	3.6 ms
loss.backward() host time with a 2 GB D2H copy queued (BT=8192, H=4096, V=128256, bf16)	43.4 ms	0.24 ms
backward GPU time, grad_output = 1.0	0.284 ms	0.333 ms
backward GPU time, grad_output = 0.5	1.100 ms	0.838 ms

With grad_output = 1.0, launching the early-exit programs costs about 0.05 ms.

Testing Done

• New test_no_host_sync_without_mean_reduction (24 cases): runs forward and backward under torch.cuda.set_sync_debug_mode("error") and checks loss, gradients and token accuracy against PyTorch. It covers:

  • reduction in {sum, none};
  • grad_output in {1.0, 0.37};
  • with and without bias;
  • plain, label smoothing, and class weights.

  All 24 fail on main and pass here.

• Hardware Type: NVIDIA H200 (141 GB); torch 2.14.1 (CUDA 13.3), triton 3.8.0, transformers 5.8.1

• run make test to ensure correctness: 5686 passed, 40 failed. The same 40 fail on main in this environment and are unrelated:

  • 36 × cutlass.cute has no attribute 'ReductionKind' (older CuTe DSL here);
  • 2 × test_rms_norm.py::test_torch_compile_backward_matches_eager;
  • 2 × JSD test_low_probability_vocab_block_is_finite NaN.

  All FLCE tests pass (test/transformers 165, test/ops 171), and so do the other element_mul_kernel users' tests.

• run make checkstyle to ensure code style: ruff check and ruff format --check are clean on the changed files (ruff 0.16.6; pre-commit pins 0.14.11).

• run make test-convergence to ensure convergence: everything passes except failures that are identical on main in this environment:

  • mini_llama4 bf16 loss mismatch;
  • multimodal mini_paligemma / mini_paligemma2 / mini_gemma3 (view size is not compatible);
  • multimodal mini_qwen3_5 / mini_qwen3_5_moe (TileLang layout error).