[CPU] add faster KV-cache writes (#25874)

2026-05-25 10:28:52 +08:00
parent e1463bb2c2
commit 821d5f4a5b
6 changed files with 249 additions and 1 deletions
@@ -111,6 +111,16 @@ def _set_kv_buffer_impl(
            row_bytes=row_bytes,
        )

+    if _is_cpu and _cpu_has_amx_support:
+        return torch.ops.sgl_kernel.store_cache_cpu(
+            k,
+            v,
+            k_cache,
+            v_cache,
+            indices,
+            row_dim,
+        )
+
    from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode

    if get_is_capture_mode() and alt_stream is not None: