diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index c3b6c89d4..2022c3e8b 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -961,6 +961,11 @@ def get_w8a8_block_fp8_configs( be picked and the associated configuration chosen to invoke the kernel. """ + # Skip config lookup during torch.compile to avoid non-Tensor ops (e.g., device name). + # Returning None forces the caller to use the default config path during compile. + if torch._dynamo.is_compiling(): + return None + # First look up if an optimized configuration is available in the configs # directory device_name = get_device_name().replace(" ", "_") diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index d52d9217d..827f60e0f 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -16,6 +16,7 @@ """Inference-only MiniMax M2 model compatible with HuggingFace weights.""" import logging +from contextlib import nullcontext from typing import Iterable, Optional, Set, Tuple, Union import torch @@ -442,9 +443,14 @@ class MiniMaxM2MoE(nn.Module): hidden_states = state.hidden_states_mlp_input if router_logits is not None: - with get_global_expert_distribution_recorder().with_current_layer( - self.layer_id - ): + ctx = ( + nullcontext() + if get_global_server_args().enable_piecewise_cuda_graph + else get_global_expert_distribution_recorder().with_current_layer( + self.layer_id + ) + ) + with ctx: state.topk_weights_local, state.topk_idx_local, _ = self.topk( hidden_states=hidden_states, router_logits=router_logits, @@ -475,9 +481,14 @@ class MiniMaxM2MoE(nn.Module): def op_dispatch_b(self, state): """Dispatch B operation for TBO - complete async dispatch""" if self.ep_size > 1: - with get_global_expert_distribution_recorder().with_current_layer( - self.layer_id - ): + ctx = ( + nullcontext() + if get_global_server_args().enable_piecewise_cuda_graph + else get_global_expert_distribution_recorder().with_current_layer( + self.layer_id + ) + ) + with ctx: state.dispatch_output = self.experts.deepep_dispatcher.dispatch_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) @@ -896,7 +907,12 @@ class MiniMaxM2Model(nn.Module): ) else: for i in range(self.start_layer, self.end_layer): - with get_global_expert_distribution_recorder().with_current_layer(i): + ctx = ( + nullcontext() + if get_global_server_args().enable_piecewise_cuda_graph + else get_global_expert_distribution_recorder().with_current_layer(i) + ) + with ctx: if i in self.layers_to_capture: aux_hidden_states.append(hidden_states + residual) layer = self.layers[i]