optimize mxfp4 cast/transpose

[matthiasdiener]

Description

Implements two optimizations on the cast/transpose kernel, targeting the inefficient paths:
1. Coalesced columnwise store: The plain (backward) columnwise transpose store was a scattered 16-bit-per-thread write, hitting ~35% write efficiency. It's now staged through LDS (smem_col) and flushed as full 128-bit (uint4) contiguous writes. Gated to the plain-colwise instantiations only (COLWISE_TO_LDS = USE_COLWISE && !SHUFFLE_COLWISE_FP4); rowwise and shuffled-forward paths are unaffected.
2. DPP cross-lane reductions: __shfl_xor was lowering to the LDS-path ds_bpermute on gfx950. Replaced the xor1/xor2 lanes in the amax reduction (warp_reduce_max_8_dpp) and the Hadamard butterfly (hadamard16_inplace) with __builtin_amdgcn_update_dpp permute intrinsics, which fuse into the VALU datapath.

Also fixes a missing columnwise usage flag for MXFP4.

Fixes https://github.com/ROCm/frameworks-internal/issues/16628

Performance results are in https://github.com/ROCm/frameworks-internal/issues/16628#issuecomment-4916784356

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[github-actions]

Claude review

Reviewed the two-commit optimization in cast_transpose_mxfp4_shuffled.cuh plus the four-line MXFP4Quantizer fix in layernorm_mlp.py.

• Kernel changes (DPP-based intra-quad __shfl_xor replacement + LDS-staged coalesced flush of the plain columnwise FP4 tile): correctness looks good. The DPP quad_perm masks (0xB1 for xor-1, 0x4E for xor-2) are correct and bit-exact vs. __shfl_xor within a uniformly-active 8-lane group, and the warp_reduce_max_8_dpp sequencing (xor-4 stays as __shfl_xor, xor-2/xor-1 go through DPP) is right. The LDS flush layout, THREADS_PER_COL=4, alignas(16), and the __syncthreads() placement all check out. One nit posted inline: the flush now depends on M % 32 == 0 for uint4 alignment/boundary correctness (previously only M % 2 was required); worth an NVTE_CHECK in the caller.
• layernorm_mlp.py: the added IS_HIP_EXTENSION-guarded set_usage(rowwise=True, columnwise=True) for MXFP4Quantizer correctly patches around the upstream tuple that omits it, without touching CUDA behavior. Existing PR-thread already covers the rationale.

Copyright headers: OK (both files carry 2025-2026 / 2024-2026 AMD lines).

[wenchenvincent]

shuffled means shuffled by Hadamard transform?

[matthiasdiener]

No, shuffled is the swizzling for the Gemm here.

[wenchenvincent]

Oh, it was that the mxfp4 GEMM in AITER need preshuffled B?