diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index db10d3ce4..aad266166 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -803,6 +803,16 @@ class FusedMoE(torch.nn.Module): param=param, weight_name=weight_name, ) + elif isinstance(method, Fp8MoEMethod) and ( + get_moe_runner_backend().is_flashinfer_trtllm_routed() + or get_moe_runner_backend().is_flashinfer_trtllm() + ): + # Drop the GPU mxfp8 shuffle-index cache on every reload for mxfp8 trtllm, trtllm_routed + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + clear_mxfp8_shuffle_index_cache, + ) + + clear_mxfp8_shuffle_index_cache() loaded_weight = ( loaded_weight.t().contiguous() @@ -1025,6 +1035,16 @@ class FusedMoE(torch.nn.Module): method = self.quant_method if hasattr(self, "scheme"): method = self.scheme + if isinstance(method, Fp8MoEMethod) and ( + get_moe_runner_backend().is_flashinfer_trtllm_routed() + or get_moe_runner_backend().is_flashinfer_trtllm() + ): + # Drop the GPU mxfp8 shuffle-index cache on every reload for mxfp8 trtllm, trtllm_routed + from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + clear_mxfp8_shuffle_index_cache, + ) + + clear_mxfp8_shuffle_index_cache() loaded_weight = ( loaded_weight.t().contiguous() if ( diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 5c80dc22d..19a1df20f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -113,6 +113,14 @@ _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8: dict[ ] = {} +def clear_mxfp8_shuffle_index_cache() -> None: + """Drop the cached MXFP8 MoE row-index permutations. + The cached index tensors are GPU-resident; sglang reuses the weights-region + memory across weight-update cycles + """ + _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8.clear() + + def _is_gated(layer: Module) -> bool: """Return whether the MoE layer uses a gated activation (default True).""" is_gated = ( diff --git a/test/registered/rl/test_update_weights_from_disk_blackwell.py b/test/registered/rl/test_update_weights_from_disk_blackwell.py index 78fd977d9..d12676e45 100644 --- a/test/registered/rl/test_update_weights_from_disk_blackwell.py +++ b/test/registered/rl/test_update_weights_from_disk_blackwell.py @@ -20,7 +20,7 @@ class UpdateWeightsFromDiskBase: model = None base_url = DEFAULT_URL_FOR_TEST request_timeout = 120 - update_timeout = 240 + update_timeout = 120 idle_timeout = 30 launch_env = None decode_payload = {