[AMD] Preserve the AITER expert mask across torch_memory_saver pause/resume

[JessicaJiang-123]

Motivation

On ROCm with AITER, MoriEPDispatcher and DeepEPDispatcher create an expert_mask_gpu at construction time to mark which experts are local to the current EP rank.

The mask is allocated while the model is inside the memory-saver WEIGHTS region, so its GPU pages are managed by torch_memory_saver. However, the dispatcher is not an nn.Module, and expert_mask_gpu is only a plain attribute, so it is not included in model.named_buffers().

SGLang preserves init-time GPU state across release_memory_occupation() / resume_memory_occupation() through _export_static_state() / _import_static_state(), which only operate on registered buffers. As a result, the expert mask is not restored after resume and becomes all zeros.

AITER's fused_moe uses this mask to filter local experts. An all-zero mask makes the rank appear to own no experts, so the expert GEMM output becomes zero and the model silently produces invalid outputs.

On a failing Qwen3-30B-A3B-FP8 run with EP4:

expert_mask  int32     (128,)       nonzero = 0/128    # expected 32/128
gemm_out     bfloat16  (...)        nonzero = 0/...    # all zeros

The model weights themselves remain correct, so normal weight validation does not catch this corruption.

Modifications

Register the dispatcher's expert mask as a non-persistent buffer on FusedMoE. This puts the tensor into named_buffers(), so the existing static-state save/restore mechanism preserves it across memory-saver pause/resume.

persistent=False keeps it out of state_dict(). _import_static_state() restores buffers in place, so the dispatcher continues to reference the same tensor.

Also add expert_mask_gpu to _NON_PERSISTENT_BUFFER_PATTERNS in the RL weight checker, since this rollout-only buffer has no training-side counterpart.

This change only affects ROCm/AITER. The dispatcher only creates the expert mask when:

_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()

On CUDA/NVIDIA, or on ROCm without AITER, expert_mask_gpu remains None, so no buffer is registered and no forward path changes.

The fix covers the constructor-created masks in MoriEPDispatcher and DeepEPDispatcher, which are the paths where the issue was reproduced.

Accuracy Tests

Validated with colocated RL in Miles, using SGLang for rollout and Megatron for training on the same 4×MI355X GPUs.

Configuration:

• Qwen3-30B-A3B-FP8
• TP2 / CP2 / EP4 / ETP1
• --sglang-moe-a2a-backend mori
• --sglang-mem-fraction-static 0.7
• CUDA graphs enabled
• 2 rollouts

metric	before	after	no-offload reference
expert_mask nonzero	0/128	32/128	—
raw_reward	0.0	0.5 → 0.578	0.484
log_probs	-9.43	-0.294 → -0.240	-0.294
rollout_log_probs	-6.18	-0.289 → -0.235	-0.288
total_lengths	8347	6746 → 6637	6752

Before the fix, the rollout engine produces invalid generations and responses consistently hit the 8192-token limit. After the fix, the offload path matches the no-offload reference within normal run-to-run variation, and the Miles RL weight check passes.

Speed Tests and Profiling

N/A

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #31359156884
Latest PR Test (Extra): ❌ Run #31359156799

[HaiShaw]

@JessicaJiang-123 can you make sure it works on NV with DeepEP too?

[JessicaJiang-123]

@JessicaJiang-123 can you make sure it works on NV with DeepEP too?

Confirmed. Both inner dispatchers (deepep.py & moriep.py) initialise the attribute to None and only overwrite it under the AITER gate:

# deepep.py:914-922
        self.expert_mask_gpu = None
        if _use_aiter and num_local_experts is not None:
            expert_mask = torch.zeros(num_local_experts + 1, ...)
            expert_mask[:-1] = 1
            self.expert_mask_gpu = expert_mask

# moriep.py:1074-1084  (ROCm)
        self.expert_mask_gpu = None
        if _use_aiter and num_experts is not None and num_local_experts is not None:
            ...
            self.expert_mask_gpu = expert_mask

On CUDA the attribute value is None, so if expert_mask is not None is always false.

[HaiShaw]

AITER gated.