[AMD] support ROCm TE grouped expert LoRA

[XinyuJiangCMU]

Co-author-with: @JessicaJiang-123

GroupedExpertLinearAdapter._forward_te_grouped_linear calls TE's private _GroupedLinear.forward. It does not pin a version table: _te_grouped_linear_contract() reads the non_tensor_args field order off the installed TE source, and the call site builds the tuple from the available lookup, raising when a name it cannot supply shows up. ROCm's TE unpacks three fields that CUDA's does not, and the lookup does not carry them, so every grouped-expert LoRA run on a ROCm build stops at:

RuntimeError: transformer_engine's _GroupedLinear.forward expects non_tensor_args fields this shim does not supply: ['m_splits_tensor', 'actual_m_splits', 'unpad_output']. TE changed its private grouped-linear contract; update _forward_te_grouped_linear to match.

Against the CUDA TE this shim is known to work with (2.17/2.18), the ROCm build appends three fields after debug:

build	trailing non_tensor_args fields
NVIDIA TE 2.17/2.18	..., save_original_input, debug
ROCm TE dev (2.17-based)	..., save_original_input, debug, m_splits_tensor, actual_m_splits, unpad_output

actual_m_splits and unpad_output drive ROCm's fused pad-cast-transpose path: the caller may hand TE a padded m_splits and the real row counts alongside it, and ask TE to drop the padding rows before returning. m_splits_tensor is unrelated to padding, it is the device-side group-size tensor ROCm's Triton grouped-GEMM kernel reads. All three are on ROCm/TransformerEngine's dev branch, so this is the platform's contract rather than one wheel's local patch.

Changes

Add the three names to available, using the same defaults TE's own public GroupedLinear.forward declares for them: m_splits_tensor=None, actual_m_splits=None, unpad_output=False. That is what TE passes down when no padding is requested, so the shim keeps the unpadded path it already assumes.

Why a CUDA build is unaffected

field_names comes from the installed TE source and the tuple is built as tuple(available[name] for name in field_names). A TE that does not unpack these names never selects them, so the tuple is element-for-element what it was before and the three entries are dead keys.

Validation

Ran the GPT-OSS 20B MoE-LoRA end-to-end recipe on 4x MI355X with this change and the matching SGLang MoE-LoRA fix, against a ROCm TE 2.14-based wheel that carries the same three fields. Both serving combinations completed (rollout, two training steps, and the post-training LoRA weight update), with the Ray job reporting success for each. Before the change the run stopped at the RuntimeError above.