[ROCm][Bugfix] Enable the fp32 head_dtype torch.mm fast path on ROCm

[wjabbour]

Purpose

#48390 added a torch.mm(..., out_dtype=torch.float32) fast path so the fp32
head_dtype lm_head projection avoids materializing a full fp32 copy of the
lm_head weight on every step. That fast path was gated on
current_platform.is_cuda(), so ROCm always fell back to the cast path
(F.linear(hidden_states.to(fp32), lm_head.weight.to(fp32), ...)), which
allocates and writes a full fp32 copy of the (large) lm_head weight matrix
every forward step.

ROCm doesn't need that fallback. In PyTorch's aten/src/ATen/native/cuda/Blas.cpp,
the fp32-output/bf16-or-fp16-input case is handled on ROCm too — it's routed
to the classic (non-Lt) hipblasGemmEx path instead of hipBLASLt (which has a
correctness bug for this exact dtype combination), not excluded outright:

#ifdef USE_ROCM
  disable_addmm_cuda_lt = disable_addmm_cuda_lt || is_float_output_with_half_input;
#endif

This PR extends the fast-path gate to include ROCm.

Why this isn't a duplicate

Searched for related open/merged work (head_dtype, out_dtype rocm,
logits_processor rocm). Found two merged follow-ups to #48390 — #48525
(LoRA path support) and #48654 (a ROCm CI test fix for an unrelated
CPU-dispatch issue in the same test file) — neither touches the
current_platform.is_cuda() gate in _apply_head.

Test

• Verified torch.mm(bf16, bf16, out_dtype=torch.float32) produces correct,
  finite results on ROCm hardware (AMD Radeon RX 9070 XT, gfx1201, ROCm 7.1,
  torch 2.13.0.dev) — matches an fp32 reference to ~1e-6.
• Added test_fp32_head_uses_mm_fast_path_on_device to
  tests/v1/sample/test_head_dtype.py, asserting the cast path (F.linear)
  is not called when the fast path is available, and that results match
  the fp32 reference. Confirmed this test fails without the fix (falls
  through to the cast path) and passes with it, on the RDNA4 hardware above.
• pytest tests/v1/sample/test_head_dtype.py -k "not e2e": 6 passed.
• pre-commit run (ruff, mypy, etc.) on both changed files: passed.
• Not yet run on CDNA (MI250/MI300) directly — tests/v1/sample/ is already
  mirrored onto AMD CI's MI250 (gfx90a) and MI300 (gfx942) runners (per
  .buildkite/test-amd.yaml), so CI will validate this on CDNA hardware.
• No model-output/accuracy impact: this only changes which internal GEMM
  call computes the fp32 projection, not the math — same as the cast path,
  just without the extra fp32 weight materialization.

Notes

AI assistance (Claude Code) was used to investigate PyTorch's ROCm dispatch
behavior and draft this change; I reviewed every line and ran the tests above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.