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

Carry deferred attention operands and reuse multimodal shared memory (#39870)

Browse Source 
Co-authored-by: fei-xx <135589532+fei-xx@users.noreply.github.com>
Co-authored-by: jmswen <jmswen@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-09-17 01:17:15 -07:00
committed by GitHub
co-authored by fei-xx jmswen
parent 882577451e
commit acfde25d34
4 changed files with 293 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/test_radix_attention.py
+65
View File
@@ -175,6 +175,71 @@ class TestRadixAttentionGraphInterface(CustomTestCase):
self.assertEqual(output.shape, query.shape)
self.assertTrue(torch.all(output == 5))
def test_deferred_norm_rope_operands_follow_real_tokens_on_each_call(self):
layer = self._new_layer()
query = torch.zeros((4, 2, 3))
positions = torch.arange(4)
temp_scale = torch.arange(4, dtype=torch.float32).reshape(4, 1)
norm_weight = torch.ones(3)
operands = {
"mxfp8_norm_rope_positions": positions,
"mxfp8_norm_rope_temp_scale": temp_scale,
"norm_weight": norm_weight,
}
for breakable in (False, True):
with self.subTest(breakable=breakable):
context = self._new_impl_context([layer])
forward_batch = context.forward_batch
forward_batch.forward_mode = ForwardMode.EXTEND
original_cache_loc = forward_batch.out_cache_loc
backend = _RecordingAttentionBackend(return_lse=False)
with (
patch.object(
radix_attention_module,
"get_tc_piecewise_forward_context",
return_value=context,
),
patch.object(
radix_attention_module,
"get_attn_backend",
return_value=backend,
),
patch.object(
radix_attention_module,
"is_in_breakable_cuda_graph",
return_value=breakable,
),
patch.object(
radix_attention_module,
"breakable_attention_with_output_extra_kwargs",
side_effect=radix_attention_module.attention_with_output_extra_kwargs,
) as graph_break,
):
for real_tokens in (2, 4):
forward_batch.global_num_token_non_padded_cpu = real_tokens
positions.add_(10)
result = layer(query, query, query, forward_batch, **operands)
call = backend.calls[-1]
self.assertEqual(call.query.shape[0], real_tokens)
self.assertTrue(
torch.equal(
call.kwargs["mxfp8_norm_rope_positions"],
positions[:real_tokens],
)
)
self.assertTrue(
torch.equal(
call.kwargs["mxfp8_norm_rope_temp_scale"],
temp_scale[:real_tokens],
)
)
self.assertIs(call.kwargs["norm_weight"], norm_weight)
self.assertIs(forward_batch.out_cache_loc, original_cache_loc)
self.assertTrue(torch.all(result[:real_tokens] == 3))
self.assertEqual(positions.shape[0], 4)
self.assertEqual(temp_scale.shape[0], 4)
self.assertEqual(graph_break.call_count, 2 if breakable else 0)
def test_impl_preserves_attention_identity_and_lse(self):
mqa = SimpleNamespace()
mha = SimpleNamespace()
test/registered/unit/managers/test_mm_shm_error_consensus.py
+193 -30
View File
@@ -1,9 +1,14 @@
import gc
import os
import pickle
import sys
import unittest
from array import array
from multiprocessing import shared_memory
from pathlib import Path
from tempfile import TemporaryDirectory
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import torch
import torch.distributed
@@ -19,7 +24,10 @@ from sglang.srt.managers.io_struct import ( # noqa: E402
MMInputsProcessError,
TokenizedEmbeddingReqInput,
)
from sglang.srt.managers.mm_utils import ShmPointerMMData # noqa: E402
from sglang.srt.managers.mm_utils import ( # noqa: E402
ShmPointerMMData,
wrap_shm_features,
)
from sglang.srt.managers.schedule_batch import ( # noqa: E402
Modality,
MultimodalDataItem,
@@ -32,8 +40,9 @@ from sglang.srt.managers.scheduler import ( # noqa: E402
from sglang.srt.managers.scheduler_components.request_receiver import ( # noqa: E402
SchedulerRequestReceiver,
)
from sglang.srt.sampling.sampling_params import SamplingParams # noqa: E402
register_cpu_ci(est_time=23, suite="base-a-test-cpu")
register_cpu_ci(est_time=33, suite="base-a-test-cpu")
class _CloneFailure:
@@ -42,17 +51,14 @@ class _CloneFailure:
class _Handle:
def __init__(self, *, fail_unlink: bool = False):
def __init__(self):
self.closed = False
self.unlinked = False
self.fail_unlink = fail_unlink
def close(self):
self.closed = True
def unlink(self):
if self.fail_unlink:
raise PermissionError("unlink denied")
self.unlinked = True
@@ -69,15 +75,7 @@ def _failed_pointer() -> ShmPointerMMData:
def _successful_pointer() -> ShmPointerMMData:
pointer = object.__new__(ShmPointerMMData)
pointer.shm_name = "unused"
pointer.shape = torch.Size([1])
pointer.dtype = torch.float32
pointer.precomputed_hash = None
pointer._shm_handle = _Handle()
pointer.tensor = torch.ones(1)
pointer._materialization_error = None
return pointer
return pickle.loads(pickle.dumps(ShmPointerMMData(torch.ones(1))))
def _request(feature, rid: str = "vlm-request") -> TokenizedEmbeddingReqInput:
@@ -89,7 +87,7 @@ def _request(feature, rid: str = "vlm-request") -> TokenizedEmbeddingReqInput:
mm_items=[MultimodalDataItem(modality=Modality.IMAGE, feature=feature)]
),
token_type_ids=None,
sampling_params=MagicMock(),
sampling_params=SamplingParams(max_new_tokens=1),
)
@@ -157,8 +155,150 @@ def _run_consensus_rank(rank: int, world_size: int, init_file: str) -> None:
torch.distributed.destroy_process_group()
def _run_image_receiver(rank, init_file, pipe):
torch.set_num_threads(1)
torch.distributed.init_process_group(
backend="gloo", init_method=Path(init_file).as_uri(), rank=rank, world_size=2
)
try:
receiver = _receiver(tp_size=2)
object.__setattr__(receiver, "tp_cpu_group", torch.distributed.group.WORLD)
torch.distributed.barrier()
torch.distributed.all_reduce(torch.zeros(1))
initial_fds = len(os.listdir("/proc/self/fd"))
held = None
pipe.send("ready")
with (
patch(
"sglang.srt.managers.mm_utils._get_is_default_transport",
return_value=False,
),
patch(
"sglang.srt.managers.mm_utils.get_serving",
return_value=SimpleNamespace(skip_tokenizer_init=False),
),
patch(
"sglang.srt.managers.scheduler_components.request_receiver.get_parallel",
return_value=SimpleNamespace(enable_dp_attention=False),
),
):
for base in [30, 90]:
req = pickle.loads(pipe.recv_bytes())
receiver._finalize_shm_features([req])
features = [item.feature for item in req.mm_inputs.mm_items]
assert len(features) == 7
assert all(torch.all(t == base + i) for i, t in enumerate(features))
assert [item.hash for item in req.mm_inputs.mm_items] == list(
range(100, 107)
)
if held is None:
held = features[1][::2, ::2, :]
assert torch.all(held == 31)
# Copy-on-write must preserve the old clone's rank isolation.
if rank == 0:
features[0].zero_()
torch.distributed.barrier()
assert torch.all(features[0] == (0 if rank == 0 else base))
del features, req
pipe.send("exact")
del held
gc.collect()
assert len(os.listdir("/proc/self/fd")) <= initial_fds
finally:
pipe.close()
torch.distributed.destroy_process_group()
class TestShmStorageOwnership(unittest.TestCase):
def test_private_mappings_preserve_dtype_pixels_and_views(self):
for dtype in [torch.uint8, torch.float32, torch.bfloat16]:
with self.subTest(dtype=dtype):
source = torch.arange(3 * 13 * 19).to(dtype).reshape(3, 13, 19)
wire = pickle.dumps(ShmPointerMMData(source))
first, second = pickle.loads(wire), pickle.loads(wire)
left, right = first.materialize(), second.materialize()
self.assertEqual(left.dtype, dtype)
self.assertTrue(torch.equal(left, source))
view = right[:, ::2, ::2]
left.zero_()
self.assertTrue(torch.equal(right, source))
del left, right, first, second
gc.collect()
self.assertTrue(torch.equal(view, source[:, ::2, ::2]))
@unittest.skipUnless(sys.platform == "linux", "requires Linux fd accounting")
def test_seven_input_images_across_two_ranks_and_successive_requests(self):
torch.set_num_threads(1)
ctx = torch.multiprocessing.get_context("spawn")
pipes = [ctx.Pipe() for _ in range(2)]
processes = []
allocated = []
with TemporaryDirectory() as directory:
try:
for rank, (parent, child) in enumerate(pipes):
proc = ctx.Process(
target=_run_image_receiver,
args=(rank, str(Path(directory) / "gloo-init"), child),
)
proc.start()
child.close()
processes.append(proc)
for parent, _ in pipes:
self.assertTrue(parent.poll(120))
self.assertEqual(parent.recv(), "ready")
for base in [30, 90]:
pixels = [
torch.full((641 + i, 643, 3), base + i, dtype=torch.uint8)
for i in range(7)
]
req = _request(None)
req.mm_inputs.mm_items = [
MultimodalDataItem(
modality=Modality.IMAGE, feature=pixel, hash=100 + i
)
for i, pixel in enumerate(pixels)
]
with (
patch(
"sglang.srt.managers.mm_utils._get_is_default_transport",
return_value=False,
),
patch(
"sglang.srt.managers.mm_utils.get_serving",
return_value=SimpleNamespace(skip_tokenizer_init=False),
),
):
wrap_shm_features(req)
names = [item.feature.shm_name for item in req.mm_inputs.mm_items]
allocated.extend(item.feature for item in req.mm_inputs.mm_items)
wire = pickle.dumps(req)
self.assertLess(len(wire), 8192)
for pixel in pixels:
pixel.zero_()
for parent, _ in pipes:
parent.send_bytes(wire)
for parent, _ in pipes:
self.assertTrue(parent.poll(120))
self.assertEqual(parent.recv(), "exact")
for name in names:
with self.assertRaises(FileNotFoundError):
shared_memory.SharedMemory(name=name)
for proc in processes:
proc.join(30)
self.assertEqual(proc.exitcode, 0)
finally:
for proc in processes:
if proc.is_alive():
proc.terminate()
proc.join(30)
for parent, _ in pipes:
parent.close()
for pointer in allocated:
pointer.close_and_unlink()
class TestShmPointerFailureCleanup(unittest.TestCase):
def test_clone_failure_still_unlinks_and_closes(self):
def test_nonlinux_clone_failure_still_unlinks_and_closes(self):
pointer = object.__new__(ShmPointerMMData)
handle = _Handle()
pointer.shm_name = "unused"
@@ -166,7 +306,10 @@ class TestShmPointerFailureCleanup(unittest.TestCase):
pointer.tensor = _CloneFailure()
pointer._materialization_error = None
with self.assertRaisesRegex(RuntimeError, "clone failed"):
with (
patch("sglang.srt.managers.mm_utils.sys.platform", "darwin"),
self.assertRaisesRegex(RuntimeError, "clone failed"),
):
pointer.materialize()
self.assertTrue(handle.unlinked)
@@ -192,18 +335,38 @@ class TestShmPointerFailureCleanup(unittest.TestCase):
pointer.materialize()
def test_cleanup_error_does_not_escape_the_request_boundary(self):
pointer = object.__new__(ShmPointerMMData)
handle = _Handle(fail_unlink=True)
pointer.shm_name = "unused"
pointer._shm_handle = handle
pointer.tensor = torch.ones(1)
pointer._materialization_error = None
pointer = _successful_pointer()
name, handle = pointer.shm_name, pointer._shm_handle
try:
with (
patch.object(
handle, "unlink", side_effect=PermissionError("unlink denied")
),
self.assertLogs("sglang.utils", level="WARNING"),
):
result = pointer.materialize()
self.assertTrue(torch.equal(result, torch.ones(1)))
self.assertEqual(handle._fd, -1)
finally:
segment = shared_memory.SharedMemory(name=name)
segment.unlink()
segment.close()
with self.assertLogs("sglang.utils", level="WARNING"):
result = pointer.materialize()
self.assertTrue(torch.equal(result, torch.ones(1)))
self.assertTrue(handle.closed)
@unittest.skipUnless(sys.platform == "linux", "requires Linux private mappings")
def test_mapping_failure_still_releases_the_segment(self):
pointer = _successful_pointer()
name, handle = pointer.shm_name, pointer._shm_handle
with (
patch(
"sglang.srt.managers.mm_utils.mmap.mmap",
side_effect=OSError("map failed"),
),
self.assertRaisesRegex(OSError, "map failed"),
):
pointer.materialize()
self.assertEqual(handle._fd, -1)
with self.assertRaises(FileNotFoundError):
shared_memory.SharedMemory(name=name)
class TestShmRequestFailureConsensus(unittest.TestCase):