Bundle set_kv_buffer write targets into KVWriteLoc (loc + swa_loc) (#27695)

test/registered/attention/unittests/swa/test_swa_out_cache_loc.py

@@ -1,9 +1,9 @@
 """Unit coverage for SWAKVPool.set_kv_buffer with a pre-translated swa_loc.
 
 The attention backend translates out_cache_loc once per forward and passes it
-in via ``swa_loc`` (cached on its forward metadata); set_kv_buffer uses it
-directly for SWA layers and asserts it is provided. The per-backend cuda-graph
-buffer plumbing is covered by the backend SWA integration tests.
+in via a ``KVWriteLoc`` (loc + swa_loc) on the loc_info argument; set_kv_buffer
+uses swa_loc directly for SWA layers and asserts it is provided. The per-backend
+cuda-graph buffer plumbing is covered by the backend SWA integration tests.
 """
 
 import sys
@@ -13,6 +13,7 @@ from types import SimpleNamespace
 
 import torch
 
+from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.test.test_utils import CustomTestCase
 
@@ -48,21 +49,27 @@ class TestSWAKVPoolSetKVBuffer(CustomTestCase):
         swa_loc = torch.tensor([7, 8])
         pool.set_kv_buffer(
             SimpleNamespace(layer_id=1),
-            torch.tensor([3, 4]),
+            KVWriteLoc(torch.tensor([3, 4]), swa_loc),
             None,
             None,
-            swa_loc=swa_loc,
         )
         self.assertIs(recorded["swa_loc"], swa_loc)
 
     def test_swa_layer_requires_swa_loc(self):
         # set_kv_buffer never translates internally; SWA layers must be given a
-        # pre-translated swa_loc.
+        # pre-translated swa_loc (loc_info without swa_loc, or a bare loc).
         pool, _ = self._pool_and_record()
         with self.assertRaises(AssertionError):
             pool.set_kv_buffer(
                 SimpleNamespace(layer_id=1), torch.tensor([3, 4]), None, None
             )
+        with self.assertRaises(AssertionError):
+            pool.set_kv_buffer(
+                SimpleNamespace(layer_id=1),
+                KVWriteLoc(torch.tensor([3, 4])),
+                None,
+                None,
+            )
 
     def test_full_layer_ignores_swa_loc(self):
         pool, recorded = self._pool_and_record()
@@ -70,10 +77,9 @@ class TestSWAKVPoolSetKVBuffer(CustomTestCase):
         # Full layer: swa_loc supplied but ignored; loc is used.
         pool.set_kv_buffer(
             SimpleNamespace(layer_id=0),
-            loc,
+            KVWriteLoc(loc, torch.tensor([99, 99])),
             None,
             None,
-            swa_loc=torch.tensor([99, 99]),
         )
         self.assertIs(recorded["full_loc"], loc)