From cb9140ee610819b4244903e496b09c3bbbd9d54a Mon Sep 17 00:00:00 2001
From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Date: Sat, 13 Jun 2026 04:51:50 +0800
Subject: [PATCH] Enable PDL for GPT-OSS tinygemm router (#27941)

---
 python/sglang/srt/models/gpt_oss.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py
index f6010bcbb..4b0174b75 100644
--- a/python/sglang/srt/models/gpt_oss.py
+++ b/python/sglang/srt/models/gpt_oss.py
@@ -26,6 +26,7 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.distributed import (
     get_moe_expert_parallel_rank,
     get_moe_expert_parallel_world_size,
@@ -161,7 +162,7 @@ class TinyGemmLinear(ReplicatedLinear):
             and x.dtype == torch.bfloat16
         ):
             out = x.new_empty((x.shape[0], self.output_size))
-            tinygemm_bf16(x, self.weight, out, self.bias)
+            tinygemm_bf16(x, self.weight, out, self.bias, use_pdl=is_arch_support_pdl())
             return out, None
 
         return super().forward(x)
@@ -1094,7 +1095,6 @@ class GptOssForCausalLM(nn.Module):
         weight_name_mapping: dict,
         other_loaded_param_names=[],
     ):
-        tp_rank = get_tensor_model_parallel_rank()
         if is_nextn:
             logging.warning(
                 "Loading weights for nextn is currently not supported in GptOssForCausalLM. "