minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Enable PDL for GPT-OSS tinygemm router (#27941)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-06-12 13:51:50 -07:00
committed by GitHub
parent 9d37e710b7
commit cb9140ee61
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/gpt_oss.py
+2 -2
View File
@@ -26,6 +26,7 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_moe_expert_parallel_rank, get_moe_expert_parallel_rank,
get_moe_expert_parallel_world_size, get_moe_expert_parallel_world_size,
@@ -161,7 +162,7 @@ class TinyGemmLinear(ReplicatedLinear):
and x.dtype == torch.bfloat16 and x.dtype == torch.bfloat16
): ):
out = x.new_empty((x.shape[0], self.output_size)) 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 out, None
return super().forward(x) return super().forward(x)
@@ -1094,7 +1095,6 @@ class GptOssForCausalLM(nn.Module):
weight_name_mapping: dict, weight_name_mapping: dict,
other_loaded_param_names=[], other_loaded_param_names=[],
): ):
tp_rank = get_tensor_model_parallel_rank()
if is_nextn: if is_nextn:
logging.warning( logging.warning(
"Loading weights for nextn is currently not supported in GptOssForCausalLM. " "Loading weights for nextn is currently not supported in GptOssForCausalLM. "