diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 15aabe904..8ef88491e 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -234,12 +234,13 @@ class RadixAttention(nn.Module): "q_descale", "k_descale", "v_descale", + "mxfp8_norm_rope_positions", ) ): - # A score_mod callable, aux_tensors, rel_bias, or mxfp8 descale - # tensors can't cross the unified_attention_with_output custom-op - # schema; route this backend's extend attention through the plain - # eager path. + # A score_mod callable, aux_tensors, rel_bias, mxfp8 descale + # tensors, or the mxfp8 deferred norm/RoPE operands can't cross + # the unified_attention_with_output custom-op schema; route this + # backend's extend attention through the plain eager path. if is_in_breakable_cuda_graph(): lse = breakable_attention_with_output_extra_kwargs( q, k, v, output, save_kv_cache, self.layer_id, kwargs @@ -601,7 +602,8 @@ def attention_with_output_extra_kwargs( """Breakable/tc_piecewise attention for backends whose forward needs kwargs that cannot cross the ``unified_attention_with_output`` custom-op schema -- a ``score_mod`` callable and/or ``aux_tensors`` (e.g. Inkling's relative-bias - fa4 attention). Plain (not a custom op) so the callable passes through; still + fa4 attention), or the per-token mxfp8 deferred norm/RoPE operands. Plain + (not a custom op) so the callable passes through; still runs eagerly between graph segments under BCG via the wrapper below. Mirrors the real-token narrowing + padded-output write of ``unified_attention_with_output``, and narrows per-token ``aux_tensors`` too. @@ -625,7 +627,14 @@ def attention_with_output_extra_kwargs( aux_tensors = kwargs.get("aux_tensors") if aux_tensors is not None: kwargs["aux_tensors"] = [t[:real_num_tokens] for t in aux_tensors] - for per_token_key in ("rel_bias", "q_descale", "k_descale", "v_descale"): + for per_token_key in ( + "rel_bias", + "q_descale", + "k_descale", + "v_descale", + "mxfp8_norm_rope_positions", + "mxfp8_norm_rope_temp_scale", + ): t = kwargs.get(per_token_key) if t is not None: kwargs[per_token_key] = t[:real_num_tokens] diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 0e0373fd2..31b63617a 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -4,6 +4,7 @@ Multi-modality utils import copy import hashlib +import mmap import os import pickle import sys @@ -1294,10 +1295,28 @@ class ShmPointerMMData: self._materialization_error = f"{type(error).__name__}: {error}" def materialize(self) -> torch.Tensor: - """Clone tensor from shm to owned memory, then release shm handle.""" + """Return independently writable storage, then release the SHM handle. + + On Linux the tensor owns a private copy-on-write mapping. Reading the + pixels needs no clone, and writes remain local to this receiver just + as with the old clone. torch.frombuffer keeps the mapping alive until + the tensor and all derived views are released. Other platforms retain + the clone path. + """ try: if self._materialization_error is not None: raise RuntimeError(self._materialization_error) + if sys.platform == "linux": + owned = mmap.mmap( + self._shm_handle._fd, + self.tensor.numel() * self.tensor.element_size(), + access=mmap.ACCESS_COPY, + ) + try: + return torch.frombuffer(owned, dtype=self.dtype).reshape(self.shape) + except BaseException: + owned.close() + raise return self.tensor.clone() finally: self.close_and_unlink() diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index 2e65999cc..450b83086 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/test/registered/unit/layers/test_radix_attention.py @@ -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() diff --git a/test/registered/unit/managers/test_mm_shm_error_consensus.py b/test/registered/unit/managers/test_mm_shm_error_consensus.py index 51d951709..fb7660cf1 100644 --- a/test/registered/unit/managers/test_mm_shm_error_consensus.py +++ b/test/registered/unit/managers/test_mm_shm_error_consensus.py @@ -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):