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