Reland: arch-aware FP8 selection via single-instantiation dispatch

[q10]

Summary:
FP8 e4m3 has two encodings. AMD gfx94x (MI300) and gfx90a use the "fnuz" variant;
gfx950 (MI350) and CUDA use the OCP "fn" variant. FBGEMM previously assumed that
any ROCm build uses fnuz, which is wrong on gfx950 and silently produced incorrect
FP8 semantics there. This makes the selection architecture-aware at runtime, so
NFP8 TBE training is correct on gfx950 and unchanged on gfx942, including in a
single fat binary spanning both.

Why the obvious fix does not work

The device-side encoding is already arch-correct: all FP8 device load/store paths
reinterpret through the __nv_fp8_e4m3 alias in utils/float.cuh, which resolves to
the OCP or fnuz HIP type per device-compile pass. Both the fn and fnuz overloads in
vec4.cuh, weight_row.cuh and stochastic_rounding.cuh cast to that same alias, so a
single fnuz kernel instantiation emits correct machine code on both archs. The
at::Float8_e4m3fn vs fnuz C++ type is only a dispatch and tensor-dtype label.

But the label is not free-floating: the C++ template parameter emb_t must EQUAL the
tensor's scalar type. That is enforced in two places:

• fbgemm: TensorAccessorBuilder::checkTensorConstraints
  (include/fbgemm_gpu/utils/tensor_accessor_builder.h), for every
  PackedTensorAccessor*<emb_t, ...> built over an embedding weight
• PyTorch core: TensorBase::data_ptr() -> check_type()
  (aten/src/ATen/templates/TensorMethods.cpp), which FBGEMM cannot relax

So simply adding an at::ScalarType::Float8_e4m3fn dispatch case that routes to the
fnuz instantiation fails at runtime on gfx950:

RuntimeError: Expected tensor 'dev_weights' to have scalar type
Float8_e4m3fnuz, but found Float8_e4m3fn instead!

and, once the FBGEMM-side accessor check is relaxed, fails one layer deeper inside
PyTorch.

Instantiating both FP8 kernel variants on ROCm resolves the coupling but costs
+28.8% .text and +33.2% device code in the TBE training library, which is not
acceptable for large linked binaries. Linker ICF was evaluated as a way to fold the
duplicates and rejected: -Wl,--icf=all recovers only 0.4% of the .text regression,
because the duplicate host stubs relocate against distinct fatbin kernel
descriptors and most of the growth is in the fatbin, which ICF does not touch.

What this does

Keeps a single fnuz kernel instantiation and relabels the tensor at the host
boundary so emb_t matches the label. New helper in embedding_common.h:

relabel_nfp8_for_dispatch(t) -> t.view(at::kFloat8_e4m3fnuz) when t is
Float8_e4m3fn under USE_ROCM, else t unchanged

To be explicit, since the name can be read the wrong way: this does NOT force fnuz
numerics on gfx950. It rewrites only the host tensor's dtype tag, on a view, so the
type system lets the tensor reach the single instantiated kernel. The encoding is
still chosen per-arch inside device code by the __nv_fp8_e4m3 alias, so a gfx950
kernel does OCP fn math either way. Verified end to end: with a weight byte of 0x7A,
TBE forward on gfx950 returns 320.0 (the fn/bias-7 value), not 160.0 (fnuz/bias-8).

It is metadata-only: no copy, no kernel launch, and measured size-independent
(~400 ns at both 25.6M and 204.8M elements, against an ~83 us TBE forward). Tensors
handed back to Python keep their arch-correct fn label; only the kernel-facing view
is relabeled.

Applied at the host dispatch boundaries:

codegen/training/pt2/embedding_split_host_pt2_autograd_template.cpp (fwd, bwd)
codegen/training/backward/embedding_backward_split_host_template.cpp (fwd, bwd)
codegen/training/optimizer/embedding_optimizer_split_template.cu
src/split_embeddings_cache/lfu_cache_populate.cu
src/split_embeddings_cache/lru_cache_populate.cu
src/split_embeddings_cache/lxu_cache.cu
src/split_embeddings_cache/reset_weight_momentum.cu

Deliberately not covered, with reasons:

• lxu_cache_weights and all momentum / prev_iter / row_counter / grad_output
  tensors: cache_t and grad_t never include FP8, and NFP8 TBE pins the cache to
  fp16.
• SSD TBE: tbe/ssd/training.py asserts weights_precision in (FP32, FP16) and has
  no NFP8 support. It shares the non-PT2 host template regardless.
• The op-level forward / backward / indice_weights templates: every reachable
  caller now flows through a relabeling host wrapper. Hardening them would need
  ~13 use-site renames for no reachable benefit.

Note this cannot be done in Python. TBE's forward is torch.jit.script-ed, and
TorchScript rejects every form of the dtype view: t.view(dtype) silently resolves to
the view(int[] size) overload (treating the dtype enum as a shape),
t.view(dtype=...) finds no such variant, and torch.ops.aten.view.dtype fails
attribute lookup. It has to be C++.

Also in this change

• Host scalar type: embedding_common.h adds getNFP8ScalarType(), which queries the
  runtime device through the ATen CUDA hooks and returns kFloat8_e4m3fnuz for
  gfx94x / gfx90a, otherwise kFloat8_e4m3fn. This replaces a compile-time
  HIP_FP8_TYPE_OCP branch that was undefined at the point of use and silently
  forced fnuz on every ROCm arch. It is resolved at runtime, so it is correct for
  multi-arch fat binaries.
• Python: split_embedding_configs.py adds nfp8_dtype(), backed by an arch query
  marked torch.jit.ignore, mirroring the C++ arch-to-encoding mapping. The three
  fnuz-on-any-hip selection sites now use it. SparseType.from_dtype already
  accepted both encodings and is unchanged.
• dispatch_macros.h: adds the Float8_e4m3fn case to
  FBGEMM_DISPATCH_FLOAT_HALF_AND_FP8_CASE on ROCm, and retargets the comment above
  the FP8 dispatch case at the real invariant.
• A static_assert in embedding_backward_split_template.cu pinning FP8 out of the
  ROCm optimized backward. Unlike the generic path, that path converts emb_t
  through the c10 FP8 types, which are label-bound (exponent bias 8 vs 7), so
  admitting FP8 there would be silently wrong on gfx950.
• A benchmark for the relabel overhead, and a regression test
  (test/tbe/training/nfp8_encoding_test.py) that pins the invariant that the
  encoding the device kernel uses matches the encoding the host dtype label
  advertises. It is written arch-agnostically -- it compares the device round
  trip against a host round trip through the tensor's own dtype -- so it is a
  valid check on gfx942 (fnuz), gfx950 (fn) and CUDA (fn) alike, and it fails if
  the relabel ever starts leaking into the math.
• The ROCm NFP8 TBE training tests are re-enabled.

Consistency invariant

The Python nfp8_dtype and the C++ getNFP8ScalarType must use the same
arch-to-encoding mapping (gfx94x and gfx90a map to fnuz, everything else to fn).
Both use identical gfx substring checks; keep them in sync if either changes.

Reviewed By: henrylhtsang

Differential Revision: D114094367

[meta-codesync]

@q10 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114094367.