[bugfix] Fix caching of dynamic FlyDSL stage2 tensors

[RolaoDenthu]

Co-authored-by: @1am9trash

Motivation

When launching a Kimi-K3 server with SGLang in our TP8 setup, GPU OOM occurred when --chunked-prefill-size was set to 8192 or higher. This happens after AITER #4642 switched the K3 stage2 kernel from opus_moe2_* to flydsl_moe2_layout_*.

The OOM was caused by the FlyDSL stage2 path passing per-layer dynamic MoE tensors through an LRU-cached conversion helper, which kept their GPU storage alive across layers.

The cached scale conversion helper was originally introduced to reuse uint8 views of static weight scales. It was later also used for dynamic, non-weight MoE tensors in FlyDSL stage2. Since the cache retains references to these tensors, their GPU memory cannot be released after each layer and instead accumulates across layers.

mxfp4_moe_gemm2(
        inter_sorted_quant=_mxfp4_scale_u8(inter_states),
        inter_sorted_shuffled_scale=_mxfp4_scale_u8(a2_scale),

Technical Details

Added an uncached _mxfp4_tensor_u8 helper for dynamic and non-scale MoE tensors. The existing _mxfp4_scale_u8 helper remains cached and is now used only for static weight scales.

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4998 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.