Avoid full-prompt embedding in CP MTP prefill

The CP draft shared-KV path only needs this rank's local draft tokens, but the previous compatibility path embedded the full prompt before CP-splitting. For long MTP/EAGLE prefill this recreates the large hidden activation that CP shared KV is trying to avoid.\n\nThis pads local draft input ids to the per-rank max token count recorded in NSA CP metadata, embeds the padded local tensor, then trims back to the true local length. That keeps rank shapes compatible while avoiding full-prompt embedding on every rank. Missing or stale metadata keeps the existing full-embedding fallback.\n\nConstraint: CP ranks can own uneven token counts, so the local embedding path needs a rank-uniform padded shape.\nRejected: Pad local ids to the full prompt length | this preserves compatibility but loses the intended memory reduction.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the full-embedding fallback unless all CP draft metadata producers guarantee max_rank_len for every prefill path.\nTested: g0034 container py_compile for utils.py and deepseek_nextn.py; g0034 container pytest -q test/registered/unit/layers/test_nsa_cp_utils.py => 25 passed, 5 warnings.\nNot-tested: Full distributed E2E with HiCache cache-hit MTP accept-length recovery.

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -12,6 +12,8 @@ from sglang.srt.layers.attention.nsa.utils import (
     cp_split_and_rebuild_1d,
     get_cp_shared_kv_local_out_cache_loc,
     get_cp_shared_kv_local_physical_out_cache_loc,
+    get_cp_local_embedding_padded_token_count,
+    pad_cp_local_input_ids_for_embedding,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -293,6 +295,57 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
 
+    def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
+        from types import SimpleNamespace
+
+        import torch
+
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4096] * 8)
+        )
+
+        self.assertEqual(
+            get_cp_local_embedding_padded_token_count(forward_batch, 4040), 4096
+        )
+        self.assertEqual(
+            get_cp_local_embedding_padded_token_count(forward_batch, 4096), 4096
+        )
+        self.assertEqual(
+            pad_cp_local_input_ids_for_embedding(
+                SimpleNamespace(
+                    nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[6] * 8)
+                ),
+                torch.tensor([11, 12, 13, 14]),
+            ).tolist(),
+            [11, 12, 13, 14, 0, 0],
+        )
+        self.assertEqual(
+            pad_cp_local_input_ids_for_embedding(
+                SimpleNamespace(
+                    nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4] * 8)
+                ),
+                torch.tensor([11, 12, 13, 14]),
+            ).tolist(),
+            [11, 12, 13, 14],
+        )
+
+        missing_metadata = SimpleNamespace(nsa_cp_metadata=None)
+        self.assertIsNone(
+            get_cp_local_embedding_padded_token_count(missing_metadata, 4040)
+        )
+        self.assertIsNone(
+            pad_cp_local_input_ids_for_embedding(
+                missing_metadata, torch.tensor([11, 12, 13, 14])
+            )
+        )
+
+        stale_metadata = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4039] * 8)
+        )
+        self.assertIsNone(
+            get_cp_local_embedding_padded_token_count(stale_metadata, 4040)
+        )
+
     def test_local_out_cache_loc_requires_compute_owner_pages(self):
         import torch
         from types import SimpleNamespace