[PCG] GPT OSS Triton Kernel Support (#18405)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>

python/sglang/srt/model_executor/model_runner.py

@@ -2126,6 +2126,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         language_model = getattr(self.model, "language_model", self.model)
         self.attention_layers = []
         self.moe_layers = []
+        self.moe_fusions = []
         for layer in language_model.model.layers:
             if hasattr(layer, "self_attn"):
                 if hasattr(layer.self_attn, "attn"):
@@ -2144,15 +2145,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     self.attention_layers.append(layer.attention.attn)
 
             moe_block = None
+            moe_fusion = None
             if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
                 moe_block = layer.mlp.experts
+                moe_fusion = layer.mlp
             if hasattr(layer, "block_sparse_moe") and hasattr(
                 layer.block_sparse_moe, "experts"
             ):
                 moe_block = layer.block_sparse_moe.experts
+                moe_fusion = layer.block_sparse_moe
             if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
                 moe_block = layer.moe.experts
+                moe_fusion = layer.moe
             self.moe_layers.append(moe_block)
+            self.moe_fusions.append(moe_fusion)
 
         if len(self.attention_layers) < self.model_config.num_hidden_layers:
             # TODO(yuwei): support Non-Standard GQA

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -233,6 +233,7 @@ class PiecewiseCudaGraphRunner:
 
         self.attention_layers = self.model_runner.attention_layers
         self.moe_layers = self.model_runner.moe_layers
+        self.moe_fusions = self.model_runner.moe_fusions
 
         if get_global_graph_memory_pool() is None:
             set_global_graph_memory_pool(self.device_module.graph_pool_handle())
@@ -358,7 +359,11 @@ class PiecewiseCudaGraphRunner:
         set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
         set_is_extend_in_batch(False)
         with set_forward_context(
-            forward_batch, self.attention_layers, self.quant_config, self.moe_layers
+            forward_batch,
+            self.attention_layers,
+            self.quant_config,
+            self.moe_layers,
+            self.moe_fusions,
         ):
             _ = self.model_runner.model.forward(
                 forward_batch.input_ids,
@@ -520,7 +525,11 @@ class PiecewiseCudaGraphRunner:
 
             kwargs = {}
             with set_forward_context(
-                forward_batch, self.attention_layers, self.quant_config, self.moe_layers
+                forward_batch,
+                self.attention_layers,
+                self.quant_config,
+                self.moe_layers,
+                self.moe_fusions,
             ):
                 self.model_runner.model.forward(
                     forward_batch.input_ids,
@@ -684,6 +693,7 @@ class PiecewiseCudaGraphRunner:
                 self.attention_layers,
                 self.quant_config,
                 self.moe_layers,
+                self.moe_fusions,
             ):
                 with set_compiled(True):
                     output = self.model_runner.model.forward(