fix(rocm): route moe_align around Triton 3.8 miscompile

[thedandano]

Problem

On ROCm with Triton 3.8, MoE alignment can put tokens in the wrong expert blocks without raising an error. Decode and prefill can then use the wrong expert weights. This falls under the AMD work in #541.

Cause

moe_align.py builds store masks from tl.histogram. Triton 3.8's AMD backend folds comparisons against that histogram to false. The compiler fix landed in triton-lang/triton#11246, but its 3.8 backport, triton-lang/triton#11577, is still open.

Expected behavior

Each expert gets contiguous, padded token blocks. expert_ids must name the expert for each used block, and num_tokens_post_pad must match the padded count.

Fix

Use a PyTorch path with fixed-size buffers on ROCm with Triton 3.8. Keep the existing Triton path on older ROCm stacks and on CUDA. The fallback handles FreeToken's extra expert slot and invalid token IDs. FREETOKEN_MOE_ALIGN_IMPL lets an operator force either path while diagnosing an installation.

This is the follow-up to the earlier PR against @zihaomu's ROCm branch. Zihao reviewed and tested the fallback there, then suggested sending it upstream after #132 merged. Thanks for the review and for the ROCm foundation, @zihaomu.

Testing

Earlier hardware A/B testing, reported on the first PR: RX 7800 XT (gfx1101), ROCm 10.0, PyTorch 2.11.0+rocm10.0.0, Triton 3.8.0. On the same container and GPU, pytest tests/engine/ tests/moe/ aborted with a GPU memory fault after 14 tests without the fix. With the fix, it finished at 218 passed, 3 unrelated failures, and 9 skipped. tests/kernels/test_moe_align.py had 44 passes. These runs were on the earlier foundation branch, not the current patch against main.

@zihaomu also tested the first PR on an R9700 with Triton 3.7.1: 24 alignment tests passed, including graph capture and replay. Their review led to the version gate and installation fix; current main already allows Triton 3.8.

On Fedora 44 with a Ryzen 7 7800X3D, PyTorch 2.13.0+cu130, and Triton 3.7.1:

• PYTHONPATH=python python -m pytest -q --confcutdir=tests/kernels tests/kernels/test_moe_align.py -> 31 passed, 2 GPU tests skipped because this checkout has no GPU device access.
• python -m ruff check python/freetoken/kernel/triton/moe_align.py tests/kernels/test_moe_align.py -> passed.
• git diff --check -> passed.