diff --git a/python/sglang/srt/sampling/penaltylib/repetition_penalty.py b/python/sglang/srt/sampling/penaltylib/repetition_penalty.py index fd03fb2b5..b9ad94f39 100644 --- a/python/sglang/srt/sampling/penaltylib/repetition_penalty.py +++ b/python/sglang/srt/sampling/penaltylib/repetition_penalty.py @@ -1,10 +1,12 @@ import torch from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer -from sglang.srt.utils import get_compiler_backend +from sglang.srt.utils import get_compiler_backend, is_npu + +_is_npu = is_npu() -@torch.compile(dynamic=True, backend=get_compiler_backend()) +@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) def apply_scaling_penalties(logits, scaling_penalties): logits[:] = torch.where( logits < 0,