Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
change
sglang
Commits
c3779233
Unverified
Commit
c3779233
authored
Aug 28, 2025
by
yhyang201
Committed by
GitHub
Aug 28, 2025
Browse files
[feat] Reduce GPU memory overhead by using weakref (#9673)
parent
f84b57c8
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
14 additions
and
2 deletions
+14
-2
python/sglang/srt/sampling/penaltylib/orchestrator.py
python/sglang/srt/sampling/penaltylib/orchestrator.py
+14
-2
No files found.
python/sglang/srt/sampling/penaltylib/orchestrator.py
View file @
c3779233
from
__future__
import
annotations
import
abc
from
typing
import
TYPE_CHECKING
,
Set
,
Type
import
weakref
from
typing
import
TYPE_CHECKING
,
Optional
,
Set
,
Type
import
torch
...
...
@@ -17,7 +18,7 @@ class BatchedPenalizerOrchestrator:
penalizers
:
Set
[
Type
[
"_BatchedPenalizer"
]],
):
self
.
vocab_size
=
vocab_size
self
.
batch
=
batch
self
.
_
batch
_ref
=
weakref
.
ref
(
batch
)
self
.
device
=
batch
.
device
self
.
penalizers
=
{
Penalizer
:
Penalizer
(
self
)
for
Penalizer
in
penalizers
}
...
...
@@ -27,6 +28,17 @@ class BatchedPenalizerOrchestrator:
is_required
|=
pen_is_required
self
.
is_required
=
is_required
@
property
def
batch
(
self
)
->
ScheduleBatch
|
None
:
return
self
.
_batch_ref
()
@
batch
.
setter
def
batch
(
self
,
value
:
Optional
[
ScheduleBatch
]):
if
value
is
None
:
self
.
_batch_ref
=
lambda
:
None
else
:
self
.
_batch_ref
=
weakref
.
ref
(
value
)
def
reqs
(
self
):
return
self
.
batch
.
reqs
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment