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

Co-authored-by: fei-xx <135589532+fei-xx@users.noreply.github.com>
Co-authored-by: jmswen <jmswen@gmail.com>

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]

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()

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()

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):