feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla (#34613)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

python/sglang/srt/mem_cache/memory_pool.py

@@ -1602,13 +1602,15 @@ class KVWriteLoc:
       KERNEL-FACING on every pool: physical by allocation on non-unified
       pools, rebound at ForwardBatch construction (``rebind_write_loc``) on
       the unified pool.
-    - ``swa_loc``: the pre-resolved SWA-sub-pool location for hybrid SWA pools
-      (``None`` otherwise).
-    - ``full_loc``: the full-attention-sub-pool location for the unified
-      memory pool (``None`` otherwise), carried in attention metadata
-      (``ForwardMetadata.out_cache_loc_full_physical``). Since the
-      construction-time rebind it is the SAME id space as ``loc``; the shared
-      full pool writes it directly and never translates.
+    - ``swa_loc``: the SWA-sub-pool location for hybrid SWA pools (``None``
+      otherwise); under the unified pool the translator derives it from the
+      same rebound loc (``sliding_window_write_loc_for``).
+    - ``full_loc``: OPTIONAL full-attention-sub-pool location. Since the
+      construction-time rebind it is the SAME id space as ``loc``, so pools
+      fall back to ``loc`` when it is ``None`` -- only triton's captured path
+      still passes its capture-stable
+      ``ForwardMetadata.out_cache_loc_full_physical`` buffer here (a
+      same-space alias slated for collapse).
 
     ``swa_loc`` and ``full_loc`` are the parallel pair (each a pre-resolved
     loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
@@ -3665,11 +3667,6 @@ class HybridLinearKVPool(KVCache):
         # virtual->physical mamba-slot translate for the HiCache offload path;
         # identity for a static pool, the allocator's `translate` for the unified pool.
         self._mamba_translate = lambda ids: ids
-        # The MLA doors take DIFFERENT id spaces: `get_mla_kv_buffer` gets
-        # ForwardBatch-built read indices (prefix_chunk_kv_indices /
-        # fetch_mha_one_shot_kv_indices), still VIRTUAL, so it translates;
-        # `set_mla_kv_buffer` gets out_cache_loc, already kernel-facing.
-        self._full_translate = lambda ids: ids
         self.use_mla = use_mla
         if full_kv_pool is not None:
             # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
@@ -3967,7 +3964,10 @@ class HybridLinearKVPool(KVCache):
         dst_dtype: Optional[torch.dtype] = None,
     ):
         assert self.use_mla, "get_mla_kv_buffer called when use_mla is False"
-        loc = self._full_translate(loc)
+        # Read door -- same kernel-facing contract as the write door: `loc` is
+        # a read-index tensor already translated at its production site
+        # (fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
+        # pool never translates.
         with self._transfer_id_context(layer):
             return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)