[ROCm] Add mixed precision support to the HIP optimized warp per row TBE backward kernel

[avbokovoy]

Summary

This change enables mixed precision on the AMD optimized warp per row TBE backward path for rowwise adagrad so that the embedding, gradient, and cache element types no longer have to match. It brings the optimized HIP path in line with the combinations the stock generic backward path already supports, excluding FP8. The device kernel was already templated on independent embedding, gradient, and cache types, so the change is mostly about removing an overly strict host side dispatch guard, adding the one missing device side specialization, and fixing a latent load and store layout bug that only became reachable once mixed precision was allowed.

Motivation

Previously the optimized HIP kernel was only selected when the embedding type and the gradient type were identical, because the host dispatch required same precision and excluded BFloat16 gradients. As a result mixed precision workloads silently fell back to the slower generic kernel, and the fallback behavior on ROCm was incorrect. The goal is to make the optimized kernel handle the same emb by grad by cache matrix as the stock warp per row kernel and to fall back cleanly when a valid specialization does not exist.

Changes

The host side dispatch guard in the backward split template was relaxed so that BFloat16 is accepted as a gradient type and the same precision requirement is removed, which lets mixed embedding, gradient, and cache combinations reach the optimized kernel while unsupported cases such as FP8 embeddings, cached tables, mixed dimensions, non CDNA devices, and non 64 wave devices still fall through to the generic path.

The device helpers in the rocm split embeddings common header gained a BFloat16 load specialization implemented as a thin forwarder to the half implementation, since BFloat16 only appears as a gradient type and the load is a bit exact sixteen bit word copy. No BFloat16 store is added because BFloat16 is never an embedding type on this path, so the unsupported store branch remains a clear trap.

The load and store helpers were refactored from roughly thirty five hand written per dimension specializations into a small set of generic over dimension templates, one per element type family, which removes a large amount of duplication while preserving the vectorized access pattern and the range bounded buffer resource that zeroes the out of range tail lanes for dimensions that are not a multiple of the wave size.

A latent correctness bug was fixed by converting the float load and store to float2 contiguous pairs plus a scalar tail so that every element type uses the same lane to element layout. This matters because the optimizer combines the gradient accumulator and the weights element by element per register, so the previous mismatch between the half contiguous pair layout and the float strided layout silently corrupted results whenever the embedding and gradient element sizes differed. This bug was masked before because the old same precision guard never allowed those combinations.

The ROCm backward unit tests were extended to exercise the new combinations by decoupling the output and gradient precision from the weights precision and sweeping the gradient type over FP32, FP16, and BFloat16 for both the optimized and the fallback kernel tests, reusing the existing reference and tolerance harness.

Testing

The optimized and fallback ROCm backward tests pass, and a deterministic comparison confirms that the optimized kernel matches the generic kernel within tolerance across all supported embedding, gradient, and cache combinations. Kernel tracing confirms that the optimized HIP warp per row kernel is actually launched for mixed precision configurations and that the generic kernel is used when the feature gate is disabled.