[Rust] Rename mem-cache to sglang-radix-tree (#37290)
This commit is contained in:
Jialin Ouyang
@@ -0,0 +1 @@
|
||||
target/
|
||||
Generated
+905
@@ -0,0 +1,905 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "adler2"
|
||||
version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
|
||||
|
||||
[[package]]
|
||||
name = "aes"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cipher",
|
||||
"cpufeatures",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "allocator-api2"
|
||||
version = "0.2.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.104"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
|
||||
|
||||
[[package]]
|
||||
name = "base64ct"
|
||||
version = "1.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
|
||||
|
||||
[[package]]
|
||||
name = "block-buffer"
|
||||
version = "0.10.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "byteorder"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
|
||||
|
||||
[[package]]
|
||||
name = "bzip2"
|
||||
version = "0.4.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8"
|
||||
dependencies = [
|
||||
"bzip2-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bzip2-sys"
|
||||
version = "0.1.13+1.0.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cc"
|
||||
version = "1.4.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273"
|
||||
dependencies = [
|
||||
"find-msvc-tools",
|
||||
"jobserver",
|
||||
"libc",
|
||||
"shlex",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "cipher"
|
||||
version = "0.4.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
|
||||
dependencies = [
|
||||
"crypto-common",
|
||||
"inout",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "constant_time_eq"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crc32fast"
|
||||
version = "1.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-utils"
|
||||
version = "0.8.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
|
||||
|
||||
[[package]]
|
||||
name = "crunchy"
|
||||
version = "0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
|
||||
|
||||
[[package]]
|
||||
name = "crypto-common"
|
||||
version = "0.1.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890"
|
||||
|
||||
[[package]]
|
||||
name = "flate2"
|
||||
version = "1.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb"
|
||||
dependencies = [
|
||||
"crc32fast",
|
||||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "foldhash"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
|
||||
|
||||
[[package]]
|
||||
name = "generic-array"
|
||||
version = "0.14.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
|
||||
dependencies = [
|
||||
"typenum",
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"wasi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"r-efi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "half"
|
||||
version = "2.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"crunchy",
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
|
||||
dependencies = [
|
||||
"allocator-api2",
|
||||
"equivalent",
|
||||
"foldhash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indoc"
|
||||
version = "2.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
|
||||
dependencies = [
|
||||
"rustversion",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inout"
|
||||
version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
||||
|
||||
[[package]]
|
||||
name = "jobserver"
|
||||
version = "0.1.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.189"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
|
||||
|
||||
[[package]]
|
||||
name = "matrixmultiply"
|
||||
version = "0.3.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
|
||||
|
||||
[[package]]
|
||||
name = "memoffset"
|
||||
version = "0.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "miniz_oxide"
|
||||
version = "0.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c"
|
||||
dependencies = [
|
||||
"adler2",
|
||||
"simd-adler32",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ndarray"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
|
||||
dependencies = [
|
||||
"matrixmultiply",
|
||||
"num-complex",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-complex"
|
||||
version = "0.4.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441"
|
||||
|
||||
[[package]]
|
||||
name = "num-integer"
|
||||
version = "0.1.47"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
||||
|
||||
[[package]]
|
||||
name = "password-hash"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"rand_core",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pbkdf2"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917"
|
||||
dependencies = [
|
||||
"digest",
|
||||
"hmac",
|
||||
"password-hash",
|
||||
"sha2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkg-config"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548"
|
||||
|
||||
[[package]]
|
||||
name = "portable-atomic"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
|
||||
|
||||
[[package]]
|
||||
name = "portable-atomic-util"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618"
|
||||
dependencies = [
|
||||
"portable-atomic",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "powerfmt"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
|
||||
|
||||
[[package]]
|
||||
name = "ppv-lite86"
|
||||
version = "0.2.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
|
||||
dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.107"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3"
|
||||
version = "0.22.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"indoc",
|
||||
"libc",
|
||||
"memoffset",
|
||||
"once_cell",
|
||||
"portable-atomic",
|
||||
"pyo3-build-config",
|
||||
"pyo3-ffi",
|
||||
"pyo3-macros",
|
||||
"unindent",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.22.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
|
||||
dependencies = [
|
||||
"once_cell",
|
||||
"target-lexicon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-ffi"
|
||||
version = "0.22.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"pyo3-build-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros"
|
||||
version = "0.22.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"pyo3-macros-backend",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros-backend"
|
||||
version = "0.22.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"pyo3-build-config",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.47"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "r-efi"
|
||||
version = "6.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"rand_chacha",
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_chacha"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
|
||||
dependencies = [
|
||||
"ppv-lite86",
|
||||
"rand_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_core"
|
||||
version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
|
||||
dependencies = [
|
||||
"getrandom 0.2.17",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||
|
||||
[[package]]
|
||||
name = "rustversion"
|
||||
version = "1.0.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
|
||||
|
||||
[[package]]
|
||||
name = "safetensors"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d93279b86b3de76f820a8854dd06cbc33cfa57a417b19c47f6a25280112fb1df"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_core"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.151"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
"serde",
|
||||
"serde_core",
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sglang-radix-tree"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"hashbrown",
|
||||
"pyo3",
|
||||
"sha2",
|
||||
"tch",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha1"
|
||||
version = "0.10.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
"sha2-asm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2-asm"
|
||||
version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
|
||||
|
||||
[[package]]
|
||||
name = "simd-adler32"
|
||||
version = "0.3.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea"
|
||||
|
||||
[[package]]
|
||||
name = "subtle"
|
||||
version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "2.0.119"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "target-lexicon"
|
||||
version = "0.12.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
|
||||
|
||||
[[package]]
|
||||
name = "tch"
|
||||
version = "0.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d3f84a069d8ba16dbf720b61e8bf131d90ffb8e958a664eae8e4993c5c2fa6f"
|
||||
dependencies = [
|
||||
"half",
|
||||
"lazy_static",
|
||||
"libc",
|
||||
"ndarray",
|
||||
"rand",
|
||||
"safetensors",
|
||||
"thiserror",
|
||||
"torch-sys",
|
||||
"zip",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.69"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "1.0.69"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.55"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134"
|
||||
dependencies = [
|
||||
"deranged",
|
||||
"num-conv",
|
||||
"powerfmt",
|
||||
"serde_core",
|
||||
"time-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time-core"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109"
|
||||
|
||||
[[package]]
|
||||
name = "torch-sys"
|
||||
version = "0.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f4ba78777379cf09aaa79708c63e477cf0f95e021d04360c6821f1a9f56173f7"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"cc",
|
||||
"libc",
|
||||
"zip",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typenum"
|
||||
version = "1.20.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "unindent"
|
||||
version = "0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "wasi"
|
||||
version = "0.11.1+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
|
||||
|
||||
[[package]]
|
||||
name = "zerocopy"
|
||||
version = "0.8.56"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb"
|
||||
dependencies = [
|
||||
"zerocopy-derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zerocopy-derive"
|
||||
version = "0.8.56"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zip"
|
||||
version = "0.6.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
|
||||
dependencies = [
|
||||
"aes",
|
||||
"byteorder",
|
||||
"bzip2",
|
||||
"constant_time_eq",
|
||||
"crc32fast",
|
||||
"crossbeam-utils",
|
||||
"flate2",
|
||||
"hmac",
|
||||
"pbkdf2",
|
||||
"sha1",
|
||||
"time",
|
||||
"zstd",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zmij"
|
||||
version = "1.0.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
|
||||
|
||||
[[package]]
|
||||
name = "zstd"
|
||||
version = "0.11.2+zstd.1.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4"
|
||||
dependencies = [
|
||||
"zstd-safe",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zstd-safe"
|
||||
version = "5.0.2+zstd.1.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"zstd-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zstd-sys"
|
||||
version = "2.0.16+zstd.1.5.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"pkg-config",
|
||||
]
|
||||
@@ -0,0 +1,40 @@
|
||||
[package]
|
||||
name = "sglang-radix-tree"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
description = "Rust tree core for the Unified Radix Cache"
|
||||
license = "Apache-2.0"
|
||||
repository = "https://github.com/sgl-project/sglang"
|
||||
|
||||
[package.metadata.sglang]
|
||||
python-module = "sglang.srt.mem_cache.rust_tree_core.mem_cache"
|
||||
debug = false
|
||||
features = ["python-extension"]
|
||||
torch-compat-header = "torch_2_13_compat.h"
|
||||
|
||||
[lib]
|
||||
name = "mem_cache"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
# Keep this crate in its own workspace until its PyO3 version can be aligned
|
||||
# with the root workspace without conflicting native `links = "python"` crates.
|
||||
[workspace]
|
||||
|
||||
[dependencies]
|
||||
hashbrown = "0.16"
|
||||
pyo3 = { version = "0.22", optional = true }
|
||||
sha2 = "0.10"
|
||||
tch = "=0.24.0"
|
||||
thiserror = "1"
|
||||
|
||||
# Hardware SHA extensions; identical digests.
|
||||
[target.'cfg(all(target_arch = "aarch64", target_os = "linux"))'.dependencies]
|
||||
sha2 = { version = "0.10", features = ["asm"] }
|
||||
|
||||
[features]
|
||||
# Keep the native core as the default so workspace tests do not link a Python
|
||||
# extension. Wheel and source builds select python-extension through package
|
||||
# metadata; shared white-box tests additionally select inspection.
|
||||
default = []
|
||||
python-extension = ["dep:pyo3", "pyo3/extension-module", "tch/python-extension"]
|
||||
inspection = []
|
||||
@@ -0,0 +1,42 @@
|
||||
# sglang-radix-tree
|
||||
|
||||
Rust tree core for the Unified Radix Cache, covering Full attention, sliding window attention, and Mamba components. It implements the tree side of the `UnifiedTreeCoreInterface` split — match/insert walks, node arena, locks, eviction walks, HiCache backup/load-back specs, and KV events — behind a PyO3 binding, while the cache orchestration stays in Python.
|
||||
|
||||
## Usage
|
||||
|
||||
Select the backend with:
|
||||
|
||||
```bash
|
||||
SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=rust
|
||||
```
|
||||
|
||||
SGLang wheels bundle the production extension. A source checkout falls back to
|
||||
the shared fingerprinted Rust-extension cache; it never writes a shared object
|
||||
into the Python package. LibTorch and the Python headers come from the running
|
||||
interpreter's PyTorch install. PyTorch 2.11 through 2.13 are accepted explicitly,
|
||||
and `torch_2_13_compat.h` covers two alignment APIs removed in PyTorch 2.13.
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
# Build (libtorch from the installed torch package):
|
||||
cd rust/sglang-radix-tree
|
||||
LIBTORCH_USE_PYTORCH=1 \
|
||||
LIBTORCH_BYPASS_VERSION_CHECK=1 \
|
||||
CXXFLAGS="-include $PWD/torch_2_13_compat.h" \
|
||||
cargo build --release --locked --features python-extension
|
||||
|
||||
# Native tests do not enable pyo3's extension-module feature:
|
||||
TORCH_ROOT=$(python3 -c 'import pathlib, torch; print(pathlib.Path(torch.__file__).parent)')
|
||||
LIBTORCH_USE_PYTORCH=1 LIBTORCH_BYPASS_VERSION_CHECK=1 \
|
||||
CXXFLAGS="-include $PWD/torch_2_13_compat.h" \
|
||||
LD_LIBRARY_PATH="$TORCH_ROOT/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \
|
||||
cargo test --locked
|
||||
```
|
||||
|
||||
The `inspection` Cargo feature adds white-box methods for the shared Python/Rust
|
||||
cache suite. Production wheels do not enable it.
|
||||
|
||||
Unit tests live in `src/tests/`, mirroring the source layout one file per module (wired via `#[cfg(test)] #[path = ...]`), so implementation files stay free of inline test blocks.
|
||||
|
||||
Supported component sets are `[Full]`, `[Full, SWA]`, `[Full, Mamba]`, and `[Full, SWA, Mamba]`.
|
||||
@@ -0,0 +1,493 @@
|
||||
//! FULL attention component driver: overrides the methods FULL customizes and inherits
|
||||
//! the rest from the `TreeComponent` defaults.
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap, HashSet};
|
||||
|
||||
use tch::{Kind, Tensor};
|
||||
|
||||
use crate::components::TreeComponent;
|
||||
use crate::components::{ComponentType, FULL};
|
||||
use crate::node::ChildKeyType;
|
||||
use crate::node::Node;
|
||||
use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx};
|
||||
use crate::unified_lru_list::PriorityKey;
|
||||
use crate::unified_tree_core::{
|
||||
CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertResult,
|
||||
MatchPrefixParams, MatchResult, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore,
|
||||
};
|
||||
|
||||
/// FULL attention component driver; owns the FULL device/host value slots.
|
||||
pub struct FullComponent;
|
||||
|
||||
impl FullComponent {
|
||||
/// The component's device value slot.
|
||||
pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(FULL);
|
||||
/// The component's host value slot.
|
||||
pub const HOST: ValueSlotIdx = ValueSlotIdx::host(FULL);
|
||||
}
|
||||
|
||||
impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
|
||||
fn component_type(&self) -> ComponentType {
|
||||
FULL
|
||||
}
|
||||
|
||||
fn create_match_validator(
|
||||
&self,
|
||||
_tree_core: &UnifiedTreeCore<K>,
|
||||
match_device_only: bool,
|
||||
) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool> {
|
||||
// Device value present -> always a boundary; otherwise a backuped (host-resident)
|
||||
// node also matches, unless the match is restricted to device.
|
||||
Box::new(move |tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_| {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
node.has_device_value(FULL) || (!match_device_only && node.has_host_value(FULL))
|
||||
})
|
||||
}
|
||||
|
||||
fn finalize_match_result_in_tree_core(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
mut result: MatchResult,
|
||||
params: &MatchPrefixParams<'_, K>,
|
||||
value_chunks: &[Tensor],
|
||||
best_value_len: usize,
|
||||
) -> MatchResult {
|
||||
// Compute Full KV host hit length: walk from last_host_node up to
|
||||
// last_device_node, summing host_value lengths of evicted nodes.
|
||||
let mut kv_host_hit = 0;
|
||||
let mut node_idx = tree_core.arena.resolve(result.best_match_node_id);
|
||||
let last_device_idx = tree_core.arena.resolve(result.last_device_node_id);
|
||||
while node_idx != last_device_idx {
|
||||
let node = tree_core.arena.node(node_idx);
|
||||
let parent = node.try_parent().unwrap_or_else(|| {
|
||||
panic!(
|
||||
"finalize walk from best_match_node {} hit root {} before \
|
||||
last_device_node {}",
|
||||
result.best_match_node_id, node.id, result.last_device_node_id
|
||||
)
|
||||
});
|
||||
kv_host_hit += node.host_value_len(FULL);
|
||||
node_idx = parent;
|
||||
}
|
||||
if kv_host_hit > 0 {
|
||||
result.host_hit_length = result.host_hit_length.max(kv_host_hit);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn redistribute_on_node_split(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
new_parent_id: NodeIdx_,
|
||||
child_id: NodeIdx_,
|
||||
) {
|
||||
let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id);
|
||||
let split_len = new_parent.key.atom_len() as i64;
|
||||
new_parent.copy_device_lock_ref(FULL, child);
|
||||
if child.has_device_value(FULL) {
|
||||
Node::redistribute_child_device_value(new_parent, child, FULL, split_len);
|
||||
}
|
||||
if child.has_host_value(FULL) {
|
||||
Node::redistribute_child_host_value(new_parent, child, FULL, split_len);
|
||||
}
|
||||
}
|
||||
|
||||
fn evict_component(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
target: EvictLayer,
|
||||
) -> (usize, usize) {
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
let mut freed = 0;
|
||||
let mut host_freed = 0;
|
||||
if target.contains(EvictLayer::Device) && node.has_device_value(FULL) {
|
||||
let value = node.device_value(FULL);
|
||||
freed = node.device_value_len(FULL);
|
||||
device_frees
|
||||
.entry(FULL)
|
||||
.or_default()
|
||||
.push(value.shallow_clone());
|
||||
// NOTE: cd.value = None is deferred to _cascade_evict (Full as trigger)
|
||||
// because SWA's free_swa still needs to read Full.value.
|
||||
}
|
||||
if target.contains(EvictLayer::Host) && node.has_host_value(FULL) {
|
||||
host_freed = node.host_value_len(FULL);
|
||||
host_frees
|
||||
.entry(FULL)
|
||||
.or_default()
|
||||
.push(node.take_host_value(FULL));
|
||||
}
|
||||
if freed > 0 {
|
||||
tree_core.dec_evictable_size(FULL, freed);
|
||||
}
|
||||
(freed, host_freed)
|
||||
}
|
||||
|
||||
fn eviction_priority(&self, is_leaf: bool) -> i64 {
|
||||
if is_leaf { 0 } else { 2 }
|
||||
}
|
||||
|
||||
fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize) {
|
||||
tree_core.set_evict_device_start(FULL, request_cnt);
|
||||
tree_core.full_evict_device_heap.clear();
|
||||
let arena = &tree_core.arena;
|
||||
let strategy = &tree_core.eviction_strategy;
|
||||
tree_core.full_evict_device_heap.extend(
|
||||
tree_core
|
||||
.evictable_device_leaves
|
||||
.iter()
|
||||
.map(|id| Reverse((strategy.get_priority(arena.node(id)), id))),
|
||||
);
|
||||
}
|
||||
|
||||
fn evict_device_next_node(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
_device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
_host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) -> Option<NodeIdx_> {
|
||||
let ct = FULL;
|
||||
assert!(
|
||||
tree_core.component_state(FULL).is_evict_device_ongoing,
|
||||
"Full device eviction not started"
|
||||
);
|
||||
// Re-admit the previously returned leaf's parent once it became a
|
||||
// D-leaf; the parent id was captured at return time because the leaf
|
||||
// itself may have been freed by the eviction in between.
|
||||
if let Some(last_node_parent) = tree_core.component_state(FULL).evict_device_cursor
|
||||
&& tree_core.evictable_device_leaves.contains(last_node_parent)
|
||||
{
|
||||
let key = tree_core
|
||||
.eviction_strategy
|
||||
.get_priority(tree_core.arena.node(last_node_parent));
|
||||
tree_core
|
||||
.full_evict_device_heap
|
||||
.push(Reverse((key, last_node_parent)));
|
||||
}
|
||||
tree_core.component_state_mut(FULL).evict_device_cursor = None;
|
||||
// The budget only advances between calls (the driver's evictions fill
|
||||
// the tracker), so it gates the walk once up front.
|
||||
if tracker[&ct] >= tree_core.component_state(FULL).evict_device_request_cnt {
|
||||
return None;
|
||||
}
|
||||
while let Some(Reverse((_, x))) = tree_core.full_evict_device_heap.pop() {
|
||||
if !tree_core.evictable_device_leaves.contains(x) {
|
||||
continue;
|
||||
}
|
||||
let last_node_parent = tree_core.arena.node(x).try_parent();
|
||||
tree_core.component_state_mut(FULL).evict_device_cursor = last_node_parent;
|
||||
return Some(x);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>) {
|
||||
tree_core.set_evict_device_end(FULL);
|
||||
tree_core.full_evict_device_heap.clear();
|
||||
}
|
||||
|
||||
fn reclaim_coexisting_host_values(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
num_tokens: usize,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
for spare_imminent_demotes in [true, false] {
|
||||
if tracker[&FULL] >= num_tokens {
|
||||
break;
|
||||
}
|
||||
let candidates: Vec<NodeIdx_> = tree_core.full_coexisting_host_nodes.iter().collect();
|
||||
for node_id in candidates {
|
||||
if tracker[&FULL] >= num_tokens {
|
||||
break;
|
||||
}
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if !node.has_device_value(FULL) || !node.has_host_value(FULL) {
|
||||
tree_core.full_coexisting_host_nodes.discard(node_id);
|
||||
continue;
|
||||
}
|
||||
if spare_imminent_demotes && tree_core.evictable_device_leaves.contains(node_id) {
|
||||
continue;
|
||||
}
|
||||
if !tree_core.can_reclaim_coexisting_host_value_(node_id, FULL) {
|
||||
continue;
|
||||
}
|
||||
tree_core.release_coexisting_host_value_(
|
||||
node_id,
|
||||
FULL,
|
||||
tracker,
|
||||
device_frees,
|
||||
host_frees,
|
||||
);
|
||||
tree_core.full_coexisting_host_nodes.discard(node_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Evict host leaves to free KV host pool space.
|
||||
fn drive_host_eviction(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
num_tokens: usize,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
let ct = FULL;
|
||||
let arena = &tree_core.arena;
|
||||
let strategy = &tree_core.eviction_strategy;
|
||||
let mut heap: BinaryHeap<Reverse<(PriorityKey, NodeIdx_)>> = tree_core
|
||||
.evictable_host_leaves
|
||||
.iter()
|
||||
.map(|id| Reverse((strategy.get_priority(arena.node(id)), id)))
|
||||
.collect();
|
||||
while tracker[&ct] < num_tokens {
|
||||
let Some(Reverse((_, x))) = heap.pop() else {
|
||||
break;
|
||||
};
|
||||
if !tree_core.evictable_host_leaves.contains(x) {
|
||||
continue;
|
||||
}
|
||||
// The parent id is captured before the eviction frees the leaf.
|
||||
let parent = tree_core.arena.node(x).try_parent();
|
||||
tree_core.evict_host_leaf_(x, tracker, device_frees, host_frees);
|
||||
if let Some(parent) = parent
|
||||
&& tree_core.evictable_host_leaves.contains(parent)
|
||||
{
|
||||
let key = tree_core
|
||||
.eviction_strategy
|
||||
.get_priority(tree_core.arena.node(parent));
|
||||
heap.push(Reverse((key, parent)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn acquire_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
mut result: IncLockRefResult,
|
||||
lock_host: bool,
|
||||
) -> IncLockRefResult {
|
||||
let ct = FULL;
|
||||
|
||||
// Only the last host node needs to be protected.
|
||||
if lock_host {
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
// write_back mode: the anchor may be device-only (no host_value); pin it anyway.
|
||||
if !node.has_host_value(FULL) && !tree_core.is_write_back {
|
||||
return result;
|
||||
}
|
||||
node.inc_host_lock_ref(FULL);
|
||||
tree_core.update_evictable_leaf_sets_(node_id);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Skip the bottom evicted segment, recording it for the matching release.
|
||||
let on_boundary = |node: &Node<K>| node.is_root() || node.has_device_value(FULL);
|
||||
let mut cur = node_id;
|
||||
let mut node = tree_core.arena.node(cur);
|
||||
if !on_boundary(node) {
|
||||
let skip_lock_node_ids = result.skip_lock_node_ids.entry(ct).or_default();
|
||||
loop {
|
||||
skip_lock_node_ids.insert(node.id);
|
||||
cur = node.parent();
|
||||
node = tree_core.arena.node(cur);
|
||||
if on_boundary(node) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Lock the device-on segment up to the root.
|
||||
let mut delta = 0;
|
||||
loop {
|
||||
let node = tree_core.arena.node_mut(cur);
|
||||
if node.is_root() {
|
||||
break;
|
||||
}
|
||||
assert!(
|
||||
node.has_device_value(FULL),
|
||||
"FULL invariant broken: evicted ancestor {cur} above device-on segment"
|
||||
);
|
||||
let parent = node.parent();
|
||||
let newly_locked_len = if node.device_lock_ref(FULL) == 0 {
|
||||
Some(node.device_value_len(FULL))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
node.inc_device_lock_ref(FULL);
|
||||
if let Some(key_len) = newly_locked_len {
|
||||
tree_core.dec_evictable_size(FULL, key_len);
|
||||
tree_core.inc_protected_size(FULL, key_len);
|
||||
delta += key_len;
|
||||
}
|
||||
tree_core.evictable_device_leaves.discard(cur);
|
||||
cur = parent;
|
||||
}
|
||||
result.delta = Some(delta);
|
||||
result
|
||||
}
|
||||
|
||||
fn release_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
params: Option<&DecLockRefParams>,
|
||||
lock_host: bool,
|
||||
) {
|
||||
let ct = FULL;
|
||||
|
||||
if lock_host {
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
if node.host_lock_ref(FULL) == 0 {
|
||||
return;
|
||||
}
|
||||
// Mirror of `acquire`. write_back uses a pure counter.
|
||||
if !node.has_host_value(FULL) && !tree_core.is_write_back {
|
||||
return;
|
||||
}
|
||||
node.dec_host_lock_ref(FULL);
|
||||
tree_core.update_evictable_leaf_sets_(node_id);
|
||||
return;
|
||||
}
|
||||
|
||||
let empty = HashSet::new();
|
||||
let skip_lock_node_ids = params
|
||||
.and_then(|p| p.skip_lock_node_ids.get(&ct))
|
||||
.unwrap_or(&empty);
|
||||
let mut cur = node_id;
|
||||
loop {
|
||||
let node = tree_core.arena.node_mut(cur);
|
||||
if node.is_root() {
|
||||
break;
|
||||
}
|
||||
let parent = node.parent();
|
||||
if skip_lock_node_ids.contains(&node.id) {
|
||||
cur = parent;
|
||||
continue;
|
||||
}
|
||||
assert!(
|
||||
node.has_device_value(FULL),
|
||||
"release_component_lock: node {cur} has no FULL device value"
|
||||
);
|
||||
let old_lock_ref = node.device_lock_ref(FULL);
|
||||
assert!(
|
||||
old_lock_ref > 0,
|
||||
"release_component_lock: node {cur} is not locked"
|
||||
);
|
||||
let newly_unlocked_len = if old_lock_ref == 1 {
|
||||
Some(node.device_value_len(FULL))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
node.dec_device_lock_ref(FULL);
|
||||
if let Some(key_len) = newly_unlocked_len {
|
||||
tree_core.dec_protected_size(FULL, key_len);
|
||||
tree_core.inc_evictable_size(FULL, key_len);
|
||||
tree_core.update_evictable_leaf_sets_(cur);
|
||||
}
|
||||
cur = parent;
|
||||
}
|
||||
}
|
||||
|
||||
fn build_hicache_transfers(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
_mamba_pool_idx: Option<Tensor>,
|
||||
_host_indices: Option<Tensor>,
|
||||
_token_ids: Option<&[i64]>,
|
||||
_prefetch_tokens: usize,
|
||||
_last_hash: Option<&str>,
|
||||
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
|
||||
Ok(match phase {
|
||||
// Full KV backup is handled by the main flow
|
||||
// (cache_controller.write on host_value directly).
|
||||
// No extra PoolTransfer needed.
|
||||
CacheTransferPhase::BackupHost => None,
|
||||
CacheTransferPhase::LoadBack => {
|
||||
// `node` is best_match_node. FULL device evict only from leaves,
|
||||
// so once we hit a device-on node, everything above is also device-on.
|
||||
let mut backed_up: Vec<Tensor> = Vec::new();
|
||||
let mut nodes_to_load: Vec<NodeId> = Vec::new();
|
||||
let mut cur = tree_core.arena.node(node_id);
|
||||
while cur.evicted() {
|
||||
backed_up.push(cur.host_value(FULL).shallow_clone());
|
||||
nodes_to_load.push(cur.id);
|
||||
cur = tree_core.arena.node(cur.parent());
|
||||
}
|
||||
backed_up.reverse();
|
||||
nodes_to_load.reverse();
|
||||
let host_indices = if backed_up.is_empty() {
|
||||
Tensor::empty([0], (Kind::Int64, tch::Device::Cpu))
|
||||
} else {
|
||||
Tensor::cat(&backed_up, 0)
|
||||
};
|
||||
Some(vec![PoolTransfer {
|
||||
name: PoolName::Kv,
|
||||
host_indices: Some(host_indices),
|
||||
nodes_to_load: Some(nodes_to_load),
|
||||
..Default::default()
|
||||
}])
|
||||
}
|
||||
CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
transfers: Vec<PoolTransfer>,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
insert_result: Option<&mut InsertResult>,
|
||||
pool_storage_result: Option<&PoolTransferResult>,
|
||||
) {
|
||||
match phase {
|
||||
CacheTransferPhase::BackupHost => {
|
||||
if let Some(transfer) = transfers.first()
|
||||
&& let Some(host_indices) = &transfer.host_indices
|
||||
{
|
||||
tree_core
|
||||
.arena
|
||||
.set_host_value(node_id, FULL, host_indices.copy());
|
||||
}
|
||||
}
|
||||
CacheTransferPhase::LoadBack => {
|
||||
if let Some(transfer) = transfers.first()
|
||||
&& let Some(device_indices) = &transfer.device_indices
|
||||
{
|
||||
let mut offset = 0i64;
|
||||
for &loaded_id in transfer.nodes_to_load.iter().flatten() {
|
||||
let loaded_idx = tree_core.arena.resolve(loaded_id);
|
||||
let loaded = tree_core.arena.node_mut(loaded_idx);
|
||||
let n_len = loaded.host_value_len(FULL) as i64;
|
||||
loaded
|
||||
.set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy());
|
||||
offset += n_len;
|
||||
// Full uses leaf sets, not LRU.
|
||||
tree_core.inc_evictable_size(FULL, n_len as usize);
|
||||
tree_core.update_evictable_leaf_sets_(loaded_idx);
|
||||
}
|
||||
}
|
||||
tree_core.update_evictable_leaf_sets_(node_id);
|
||||
}
|
||||
// The Full component has no storage-phase commits.
|
||||
CacheTransferPhase::BackupStorage | CacheTransferPhase::Prefetch => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests/components/full.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,751 @@
|
||||
//! Mamba (SSM state) component driver: overrides the methods Mamba customizes
|
||||
//! and inherits the rest from the `TreeComponent` defaults.
|
||||
//! Mamba data is per-leaf single-slot state; sizes count slots, not tokens.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use tch::Tensor;
|
||||
|
||||
use crate::components::TreeComponent;
|
||||
use crate::components::{ComponentType, MAMBA};
|
||||
use crate::node::ChildKeyType;
|
||||
use crate::node::Node;
|
||||
use crate::node::{NodeId, NodeIdx_, TreeCoreRuntimeError, ValueSlotIdx};
|
||||
use crate::unified_tree_core::{
|
||||
CacheAction, CacheInitParams, CacheTransferPhase, DecLockRefParams, EvictLayer,
|
||||
IncLockRefResult, InsertParams, InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult,
|
||||
PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, UnifiedTreeCore,
|
||||
};
|
||||
|
||||
/// Mamba component driver; owns the Mamba device/host value slots.
|
||||
pub struct MambaComponent {
|
||||
/// Joint chunk/tree-page alignment for the mamba branching seqlen.
|
||||
mamba_checkpoint_grid: usize,
|
||||
/// Per-root-path cap on cached Mamba states; None means unlimited.
|
||||
mamba_max_states_per_path: Option<usize>,
|
||||
}
|
||||
|
||||
impl MambaComponent {
|
||||
/// The component's device value slot.
|
||||
pub const DEVICE: ValueSlotIdx = ValueSlotIdx::device(MAMBA);
|
||||
/// The component's host value slot.
|
||||
pub const HOST: ValueSlotIdx = ValueSlotIdx::host(MAMBA);
|
||||
}
|
||||
|
||||
impl MambaComponent {
|
||||
/// Build the driver from the tree's init params.
|
||||
pub fn new(params: &CacheInitParams) -> Self {
|
||||
let mamba_cache_chunk_size = params
|
||||
.mamba_cache_chunk_size
|
||||
.expect("the Mamba component requires mamba_cache_chunk_size");
|
||||
MambaComponent {
|
||||
// A donated checkpoint must land on both the model's chunk grid and
|
||||
// a radix-node boundary. `params.page_size` is already widened by DCP.
|
||||
mamba_checkpoint_grid: least_common_multiple(mamba_cache_chunk_size, params.page_size),
|
||||
mamba_max_states_per_path: params.mamba_max_states_per_path,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn least_common_multiple(lhs: usize, rhs: usize) -> usize {
|
||||
let mut a = lhs;
|
||||
let mut b = rhs;
|
||||
while b != 0 {
|
||||
(a, b) = (b, a % b);
|
||||
}
|
||||
lhs / a * rhs
|
||||
}
|
||||
|
||||
impl MambaComponent {
|
||||
// Tier-selected mamba slot read for the lock paths; `host` picks the host slot.
|
||||
fn has_value<K: ChildKeyType>(node: &Node<K>, host: bool) -> bool {
|
||||
if host {
|
||||
node.has_host_value(MAMBA)
|
||||
} else {
|
||||
node.has_device_value(MAMBA)
|
||||
}
|
||||
}
|
||||
|
||||
/// Defer the path-cap eviction so it runs after the insert's BackupKV.
|
||||
fn emit_excess_path_states_eviction_(
|
||||
&self,
|
||||
tail_node_id: NodeId,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
) {
|
||||
if self.mamba_max_states_per_path.is_none() {
|
||||
return;
|
||||
}
|
||||
cache_actions.push(CacheAction::MambaEvictExcessPathStates { tail_node_id });
|
||||
}
|
||||
}
|
||||
|
||||
impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
|
||||
fn component_type(&self) -> ComponentType {
|
||||
MAMBA
|
||||
}
|
||||
|
||||
fn needs_incremental_backup(&self, tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_) -> bool {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
node.has_device_value(MAMBA) && !node.has_host_value(MAMBA)
|
||||
}
|
||||
|
||||
/// A match consumes only the best-match node's mamba state, so MATCH_END
|
||||
/// touches just that node; new-leaf states enter the LRU at insert commit,
|
||||
/// so WALKDOWN and INSERT_END are no-ops.
|
||||
fn refresh_lru(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
phase: LRURefreshPhase,
|
||||
node_id: NodeIdx_,
|
||||
) {
|
||||
match phase {
|
||||
LRURefreshPhase::Walkdown => {}
|
||||
LRURefreshPhase::MatchEnd => {
|
||||
if tree_core.arena.has_device_value(node_id, MAMBA) {
|
||||
tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id);
|
||||
}
|
||||
}
|
||||
LRURefreshPhase::InsertEnd => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// A per-match predicate accepting nodes that hold mamba data.
|
||||
fn create_match_validator(
|
||||
&self,
|
||||
_tree_core: &UnifiedTreeCore<K>,
|
||||
match_device_only: bool,
|
||||
) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool> {
|
||||
// HiCache: evicted + backuped (host_value present) is also a valid match.
|
||||
Box::new(move |tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_| {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
node.has_device_value(MAMBA) || (!match_device_only && node.has_host_value(MAMBA))
|
||||
})
|
||||
}
|
||||
|
||||
/// The mamba branching seqlen and the host-only hit bump.
|
||||
fn finalize_match_result_in_tree_core(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
mut result: MatchResult,
|
||||
_params: &MatchPrefixParams<'_, K>,
|
||||
_value_chunks: &[Tensor],
|
||||
_best_value_len: usize,
|
||||
) -> MatchResult {
|
||||
let mamba_boundary_len = result.device_indices.size()[0] as usize + result.host_hit_length;
|
||||
|
||||
// Full KV may extend beyond the latest reusable Mamba state. The branching
|
||||
// point is the last checkpoint-grid-aligned position within the Full-KV hit
|
||||
// that lies beyond the current Mamba boundary.
|
||||
let aligned_seqlen =
|
||||
result.full_kv_hit_length / self.mamba_checkpoint_grid * self.mamba_checkpoint_grid;
|
||||
result.mamba_branching_seqlen =
|
||||
(aligned_seqlen > mamba_boundary_len).then_some(aligned_seqlen);
|
||||
|
||||
// HiCache: if mamba was evicted from device but has host backup,
|
||||
// ensure mamba_host_hit_length >= 1 so load_back is triggered.
|
||||
let last_node = tree_core
|
||||
.arena
|
||||
.node(tree_core.arena.resolve(result.best_match_node_id));
|
||||
if !last_node.has_device_value(MAMBA) && last_node.has_host_value(MAMBA) {
|
||||
result.mamba_host_hit_length = result.mamba_host_hit_length.max(1);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Attach the donated mamba slot to the insert target leaf.
|
||||
fn commit_insert_component_data(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
is_new_leaf: bool,
|
||||
params: &InsertParams<'_, K>,
|
||||
result: &mut InsertResult,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
) {
|
||||
let mamba_value = params
|
||||
.mamba_value
|
||||
.as_ref()
|
||||
.expect("mamba insert requires a donated mamba_value");
|
||||
let slot_len = mamba_value.size()[0] as usize;
|
||||
|
||||
if is_new_leaf {
|
||||
tree_core
|
||||
.arena
|
||||
.set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
|
||||
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
|
||||
tree_core.inc_evictable_size(MAMBA, slot_len);
|
||||
self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
|
||||
return;
|
||||
}
|
||||
if !tree_core.arena.has_device_value(node_id, MAMBA) {
|
||||
// Tombstone refill: the node moves from the host LRU to the device LRU.
|
||||
tree_core
|
||||
.arena
|
||||
.set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if host_lru.in_list(Some(node_id)) {
|
||||
host_lru.remove_node(node_id);
|
||||
}
|
||||
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
|
||||
tree_core.inc_evictable_size(MAMBA, slot_len);
|
||||
let tick = tree_core.arena.get_and_bump_access_counter();
|
||||
tree_core.arena.node_mut(node_id).last_access_counter = tick;
|
||||
self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
|
||||
return;
|
||||
}
|
||||
tree_core.device_lru_list_mut(MAMBA).reset_node_mru(node_id);
|
||||
let tick = tree_core.arena.get_and_bump_access_counter();
|
||||
tree_core.arena.node_mut(node_id).last_access_counter = tick;
|
||||
result.mamba_exist = true;
|
||||
}
|
||||
|
||||
/// Mamba data stays on the original leaf; the new prefix node gets none.
|
||||
/// Evict shallow Mamba device checkpoints beyond the per-path cap on the
|
||||
/// tail's root path; Full KV, host backups, the tail, forks, locked nodes,
|
||||
/// and device leaves are preserved (a best-effort soft cap).
|
||||
fn evict_excess_path_states(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
tail_node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
let Some(cap) = self.mamba_max_states_per_path else {
|
||||
return;
|
||||
};
|
||||
// Mamba-value holders on the root path, tail-first.
|
||||
let mut holders: Vec<NodeIdx_> = Vec::new();
|
||||
let mut cursor = Some(tail_node_id);
|
||||
while let Some(node_id) = cursor {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if node.is_root() {
|
||||
break;
|
||||
}
|
||||
if node.has_device_value(MAMBA) {
|
||||
holders.push(node_id);
|
||||
}
|
||||
cursor = node.parent;
|
||||
}
|
||||
let mut excess = holders.len().saturating_sub(cap);
|
||||
if excess == 0 {
|
||||
return;
|
||||
}
|
||||
// Cache-level apply: the counts are not reported, only the frees.
|
||||
let mut tracker: HashMap<ComponentType, usize> = HashMap::new();
|
||||
for &node_id in holders.iter().rev() {
|
||||
if excess == 0 || node_id == tail_node_id {
|
||||
break;
|
||||
}
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if node.device_lock_ref(MAMBA) > 0 || node.children.len() != 1 {
|
||||
continue;
|
||||
}
|
||||
if tree_core.evictable_device_leaves.contains(node_id) {
|
||||
continue;
|
||||
}
|
||||
tree_core.evict_component_and_detach_lru_(
|
||||
node_id,
|
||||
MAMBA,
|
||||
device_frees,
|
||||
host_frees,
|
||||
EvictLayer::Device,
|
||||
Some(&mut tracker),
|
||||
);
|
||||
tree_core.cascade_evict_(
|
||||
node_id,
|
||||
MAMBA,
|
||||
&mut tracker,
|
||||
device_frees,
|
||||
host_frees,
|
||||
EvictLayer::Device,
|
||||
);
|
||||
excess -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
fn redistribute_on_node_split(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
new_parent_id: NodeIdx_,
|
||||
_child_id: NodeIdx_,
|
||||
) {
|
||||
let new_parent = tree_core.arena.node_mut(new_parent_id);
|
||||
if new_parent.has_device_value(MAMBA) {
|
||||
let _ = new_parent.take_device_value(MAMBA);
|
||||
}
|
||||
new_parent.set_lock_ref_(ValueSlotIdx::device(MAMBA), 0);
|
||||
// HiCache: mamba host_value stays on child (mamba = leaf-only data).
|
||||
if new_parent.has_host_value(MAMBA) {
|
||||
let _ = new_parent.take_host_value(MAMBA);
|
||||
}
|
||||
new_parent.set_lock_ref_(ValueSlotIdx::host(MAMBA), 0);
|
||||
}
|
||||
|
||||
/// Free the node's mamba slot on the targeted layer(s).
|
||||
fn evict_component(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
target: EvictLayer,
|
||||
) -> (usize, usize) {
|
||||
let ct = MAMBA;
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
let mut freed = 0;
|
||||
let mut host_freed = 0;
|
||||
|
||||
// Device layer
|
||||
if target.contains(EvictLayer::Device) && node.has_device_value(MAMBA) {
|
||||
freed = node.device_value_len(MAMBA);
|
||||
device_frees
|
||||
.entry(ct)
|
||||
.or_default()
|
||||
.push(node.take_device_value(MAMBA));
|
||||
tree_core.dec_evictable_size(MAMBA, freed);
|
||||
}
|
||||
|
||||
// Host layer
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
if target.contains(EvictLayer::Host) && node.has_host_value(MAMBA) {
|
||||
host_freed = node.host_value_len(MAMBA);
|
||||
host_frees
|
||||
.entry(ct)
|
||||
.or_default()
|
||||
.push(node.take_host_value(MAMBA));
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if host_lru.in_list(Some(node_id)) {
|
||||
host_lru.remove_node(node_id);
|
||||
}
|
||||
}
|
||||
|
||||
// After device tombstone: if only host_value remains, insert into host LRU
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if target == EvictLayer::Device
|
||||
&& !node.has_device_value(MAMBA)
|
||||
&& node.has_host_value(MAMBA)
|
||||
{
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if !host_lru.in_list(Some(node_id)) {
|
||||
host_lru.insert_mru(node_id);
|
||||
}
|
||||
}
|
||||
|
||||
(freed, host_freed)
|
||||
}
|
||||
|
||||
/// Begin the device-eviction walk from this component's LRU cursor.
|
||||
fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize) {
|
||||
tree_core.set_evict_device_start(MAMBA, request_cnt);
|
||||
let cursor = tree_core
|
||||
.device_lru_list(MAMBA)
|
||||
.get_lru_no_lock(&tree_core.arena);
|
||||
tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor;
|
||||
}
|
||||
|
||||
/// Advance one device-eviction step and return a leaf, if selected.
|
||||
///
|
||||
/// An internal tombstone is one complete step so the caller can apply its
|
||||
/// pending frees and recheck allocator capacity before the next mutation.
|
||||
fn evict_device_next_node(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) -> Option<NodeIdx_> {
|
||||
let ct = MAMBA;
|
||||
assert!(
|
||||
tree_core.component_state(MAMBA).is_evict_device_ongoing,
|
||||
"Mamba device eviction not started"
|
||||
);
|
||||
let mut cursor = tree_core.component_state(MAMBA).evict_device_cursor;
|
||||
// The cursor is re-validated (reset to LRU head) if the previous
|
||||
// node's eviction removed it.
|
||||
if cursor.is_some_and(|c| !tree_core.device_lru_list(MAMBA).in_list(Some(c))) {
|
||||
cursor = tree_core
|
||||
.device_lru_list(MAMBA)
|
||||
.get_lru_no_lock(&tree_core.arena);
|
||||
}
|
||||
let next = loop {
|
||||
if tracker[&ct] >= tree_core.component_state(MAMBA).evict_device_request_cnt {
|
||||
break None;
|
||||
}
|
||||
let Some(x) = cursor else {
|
||||
break None;
|
||||
};
|
||||
if !tree_core.device_lru_list(MAMBA).in_list(Some(x)) {
|
||||
break None;
|
||||
}
|
||||
assert!(
|
||||
tree_core.arena.has_device_value(x, MAMBA),
|
||||
"Mamba eviction cursor on a valueless node {x}"
|
||||
);
|
||||
cursor = tree_core
|
||||
.device_lru_list(MAMBA)
|
||||
.get_prev_no_lock(x, &tree_core.arena);
|
||||
// A load-back pin means an in-flight DMA targets this node's slices.
|
||||
if tree_core.arena.node(x).is_load_back_pending() {
|
||||
continue;
|
||||
}
|
||||
if tree_core.evictable_device_leaves.contains(x) {
|
||||
break Some(x);
|
||||
}
|
||||
// Internal nodes are tombstoned inline (no IO).
|
||||
tree_core.evict_component_and_detach_lru_(
|
||||
x,
|
||||
ct,
|
||||
device_frees,
|
||||
host_frees,
|
||||
EvictLayer::Device,
|
||||
Some(tracker),
|
||||
);
|
||||
tree_core.cascade_evict_(x, ct, tracker, device_frees, host_frees, EvictLayer::Device);
|
||||
break None;
|
||||
};
|
||||
tree_core.component_state_mut(MAMBA).evict_device_cursor = cursor;
|
||||
next
|
||||
}
|
||||
|
||||
/// Clear the device-eviction walk cursor state.
|
||||
fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>) {
|
||||
tree_core.set_evict_device_end(MAMBA);
|
||||
}
|
||||
|
||||
/// Single-node mamba lock; host locks also detach from the host LRU.
|
||||
fn acquire_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
mut result: IncLockRefResult,
|
||||
lock_host: bool,
|
||||
) -> IncLockRefResult {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if node.is_root() {
|
||||
return result;
|
||||
}
|
||||
// A node in skip_lock_node_ids was a tombstone when this lock was acquired.
|
||||
if !Self::has_value(node, lock_host) {
|
||||
result
|
||||
.skip_lock_node_ids
|
||||
.entry(MAMBA)
|
||||
.or_default()
|
||||
.insert(node.id);
|
||||
return result;
|
||||
}
|
||||
if lock_host {
|
||||
if node.host_lock_ref(MAMBA) == 0 {
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if host_lru.in_list(Some(node_id)) {
|
||||
host_lru.remove_node(node_id);
|
||||
}
|
||||
}
|
||||
tree_core.arena.inc_host_lock_ref(node_id, MAMBA);
|
||||
} else {
|
||||
let value_len = node.device_value_len(MAMBA);
|
||||
if node.device_lock_ref(MAMBA) == 0 {
|
||||
tree_core.dec_evictable_size(MAMBA, value_len);
|
||||
tree_core.inc_protected_size(MAMBA, value_len);
|
||||
}
|
||||
tree_core.arena.inc_device_lock_ref(node_id, MAMBA);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Single-node mamba unlock; host unlocks reinsert into the host LRU.
|
||||
fn release_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
params: Option<&DecLockRefParams>,
|
||||
lock_host: bool,
|
||||
) {
|
||||
if tree_core.arena.node(node_id).is_root() {
|
||||
return;
|
||||
}
|
||||
if let Some(params) = params
|
||||
&& params
|
||||
.skip_lock_node_ids
|
||||
.get(&MAMBA)
|
||||
.is_some_and(|ids| ids.contains(&tree_core.arena.node(node_id).id))
|
||||
{
|
||||
return;
|
||||
}
|
||||
if lock_host {
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
node.dec_host_lock_ref(MAMBA);
|
||||
if node.host_lock_ref(MAMBA) == 0
|
||||
&& !node.has_device_value(MAMBA)
|
||||
&& node.has_host_value(MAMBA)
|
||||
{
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if !host_lru.in_list(Some(node_id)) {
|
||||
host_lru.insert_mru(node_id);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
let node = tree_core.arena.node(node_id);
|
||||
let device_lock_ref = node.device_lock_ref(MAMBA);
|
||||
if device_lock_ref > 0 {
|
||||
if device_lock_ref == 1 {
|
||||
let value_len = node.device_value_len(MAMBA);
|
||||
tree_core.inc_evictable_size(MAMBA, value_len);
|
||||
tree_core.dec_protected_size(MAMBA, value_len);
|
||||
}
|
||||
tree_core.arena.dec_device_lock_ref(node_id, MAMBA);
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the mamba transfer descriptors for the given phase.
|
||||
fn build_hicache_transfers(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
mamba_pool_idx: Option<Tensor>,
|
||||
host_indices: Option<Tensor>,
|
||||
_token_ids: Option<&[i64]>,
|
||||
_prefetch_tokens: usize,
|
||||
_last_hash: Option<&str>,
|
||||
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
|
||||
Ok(match phase {
|
||||
CacheTransferPhase::BackupHost => {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if node.has_host_value(MAMBA) {
|
||||
return Ok(None);
|
||||
}
|
||||
node.try_device_value(MAMBA).map(|value| {
|
||||
vec![PoolTransfer {
|
||||
name: PoolName::Mamba,
|
||||
device_indices: Some(value.shallow_clone()),
|
||||
..Default::default()
|
||||
}]
|
||||
})
|
||||
}
|
||||
CacheTransferPhase::LoadBack => {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if node.has_device_value(MAMBA) {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut transfers = Vec::new();
|
||||
// restore single node if host_value exists
|
||||
if let Some(host_value) = node.try_host_value(MAMBA) {
|
||||
transfers.push(PoolTransfer {
|
||||
name: PoolName::Mamba,
|
||||
host_indices: Some(host_value.shallow_clone()),
|
||||
nodes_to_load: Some(vec![node.id]),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
// Per-request mamba CoW (H->D copy into the request's device slot,
|
||||
// pre-allocated on caller side).
|
||||
if let (Some(mamba_pool_idx), Some(host_value)) =
|
||||
(mamba_pool_idx, node.try_host_value(MAMBA))
|
||||
{
|
||||
transfers.push(PoolTransfer {
|
||||
name: PoolName::Mamba,
|
||||
host_indices: Some(host_value.shallow_clone()),
|
||||
device_indices: Some(mamba_pool_idx.unsqueeze(0)),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
if transfers.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(transfers)
|
||||
}
|
||||
}
|
||||
CacheTransferPhase::BackupStorage => {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
let Some(host_value) = node.try_host_value(MAMBA) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(hash_value) = node.hash_value.as_ref().filter(|h| !h.is_empty()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
Some(vec![PoolTransfer {
|
||||
name: PoolName::Mamba,
|
||||
host_indices: Some(host_value.shallow_clone()),
|
||||
keys: Some(vec![hash_value[hash_value.len() - 1].clone()]),
|
||||
hit_policy: PoolHitPolicy::TrailingPages,
|
||||
..Default::default()
|
||||
}])
|
||||
}
|
||||
CacheTransferPhase::Prefetch => {
|
||||
let host_indices =
|
||||
host_indices.expect("Mamba PREFETCH build requires host indices");
|
||||
Some(vec![PoolTransfer {
|
||||
name: PoolName::Mamba,
|
||||
host_indices: Some(host_indices),
|
||||
keys: Some(vec!["__placeholder__".to_string()]),
|
||||
hit_policy: PoolHitPolicy::TrailingPages,
|
||||
..Default::default()
|
||||
}])
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Post-transfer mamba bookkeeping for the given phase.
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
transfers: Vec<PoolTransfer>,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
insert_result: Option<&mut InsertResult>,
|
||||
pool_storage_result: Option<&PoolTransferResult>,
|
||||
) {
|
||||
match phase {
|
||||
CacheTransferPhase::BackupHost => {
|
||||
if let Some(transfer) = transfers.first()
|
||||
&& let Some(host_indices) = &transfer.host_indices
|
||||
{
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
if !node.has_host_value(MAMBA) {
|
||||
node.set_host_value(MAMBA, host_indices.copy());
|
||||
}
|
||||
}
|
||||
}
|
||||
CacheTransferPhase::LoadBack => {
|
||||
let Some(transfer) = transfers.first() else {
|
||||
return;
|
||||
};
|
||||
if let Some(device_indices) = &transfer.device_indices {
|
||||
let node = tree_core.arena.node_mut(node_id);
|
||||
node.set_device_value(MAMBA, device_indices.copy());
|
||||
let count = node.device_value_len(MAMBA);
|
||||
// Move from host LRU to device LRU
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if host_lru.in_list(Some(node_id)) {
|
||||
host_lru.remove_node(node_id);
|
||||
}
|
||||
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
|
||||
tree_core.inc_evictable_size(MAMBA, count);
|
||||
}
|
||||
}
|
||||
// The python elif chain has no BACKUP_STORAGE arm.
|
||||
CacheTransferPhase::BackupStorage => {}
|
||||
CacheTransferPhase::Prefetch => {
|
||||
let Some(transfer) = transfers.first() else {
|
||||
return;
|
||||
};
|
||||
let host_indices = transfer.host_indices.as_ref();
|
||||
let loaded = pool_storage_result.is_some_and(|result| {
|
||||
result
|
||||
.extra_pool_hit_pages
|
||||
.get(&PoolName::Mamba)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
>= 1
|
||||
});
|
||||
let target_node_id = insert_result
|
||||
.as_deref()
|
||||
.and_then(|result| result.inserted_host_node)
|
||||
.map(|id| tree_core.arena.resolve(id));
|
||||
let attach_target = match (host_indices, target_node_id) {
|
||||
(Some(_), Some(target))
|
||||
if loaded && !tree_core.arena.has_host_value(target, MAMBA) =>
|
||||
{
|
||||
Some(target)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let Some(target) = attach_target else {
|
||||
// The buffer cannot attach: free it and let the caller keep
|
||||
// its own donated slot bookkeeping.
|
||||
cache_actions.push(CacheAction::FreeComponentHostSlot {
|
||||
component_type: MAMBA,
|
||||
host_indices: host_indices
|
||||
.map(|host| vec![host.shallow_clone()])
|
||||
.unwrap_or_default(),
|
||||
});
|
||||
if let Some(insert_result) = insert_result {
|
||||
insert_result.mamba_exist = true;
|
||||
}
|
||||
return;
|
||||
};
|
||||
let host_indices = host_indices.expect("an attach target implies host indices");
|
||||
tree_core
|
||||
.arena
|
||||
.set_host_value(target, MAMBA, host_indices.copy());
|
||||
if !tree_core.arena.has_device_value(target, MAMBA) {
|
||||
let host_lru = tree_core.host_lru_list_mut(MAMBA);
|
||||
if !host_lru.in_list(Some(target)) {
|
||||
host_lru.insert_mru(target);
|
||||
}
|
||||
}
|
||||
if let Some(insert_result) = insert_result {
|
||||
insert_result.mamba_exist = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Evict mamba host resources: internal nodes tombstone privately, host
|
||||
/// leaves evict atomically.
|
||||
fn drive_host_eviction(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
num_tokens: usize,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
let ct = MAMBA;
|
||||
let mut x = tree_core
|
||||
.host_lru_list(MAMBA)
|
||||
.get_lru_no_lock(&tree_core.arena);
|
||||
loop {
|
||||
if tracker[&ct] >= num_tokens {
|
||||
break;
|
||||
}
|
||||
let Some(cur) = x else {
|
||||
break;
|
||||
};
|
||||
if !tree_core.host_lru_list(MAMBA).in_list(Some(cur)) {
|
||||
break;
|
||||
}
|
||||
let x_next = tree_core
|
||||
.host_lru_list(MAMBA)
|
||||
.get_prev_no_lock(cur, &tree_core.arena);
|
||||
// A load-back pin means an in-flight DMA reads this node's host slices.
|
||||
if tree_core.arena.node(cur).is_load_back_pending() {
|
||||
x = x_next;
|
||||
continue;
|
||||
}
|
||||
if tree_core.evictable_host_leaves.contains(cur) {
|
||||
// Host leaf: atomic eviction (all components host + delete)
|
||||
tree_core.evict_host_leaf_(cur, tracker, device_frees, host_frees);
|
||||
} else {
|
||||
// Internal: tombstone Mamba + cascade
|
||||
assert!(
|
||||
tree_core.arena.has_host_value(cur, MAMBA),
|
||||
"Mamba host LRU member {cur} has no host value"
|
||||
);
|
||||
tree_core.evict_component_and_detach_lru_(
|
||||
cur,
|
||||
ct,
|
||||
device_frees,
|
||||
host_frees,
|
||||
EvictLayer::Host,
|
||||
Some(tracker),
|
||||
);
|
||||
tree_core.cascade_evict_(
|
||||
cur,
|
||||
ct,
|
||||
tracker,
|
||||
device_frees,
|
||||
host_frees,
|
||||
EvictLayer::Host,
|
||||
);
|
||||
}
|
||||
x = x_next;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../tests/components/mamba.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,491 @@
|
||||
//! Per-component drivers; each receives the whole `UnifiedTreeCore` for backward access.
|
||||
#![allow(unused_variables)]
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use tch::Tensor;
|
||||
|
||||
use crate::node::{ChildKeyType, NodeArena, NodeIdx_, TreeCoreRuntimeError};
|
||||
use crate::unified_tree_core::{
|
||||
CacheAction, CacheTransferPhase, DecLockRefParams, EvictLayer, IncLockRefResult, InsertParams,
|
||||
InsertResult, LRURefreshPhase, MatchPrefixParams, MatchResult, PoolTransfer,
|
||||
PoolTransferResult, UnifiedTreeCore,
|
||||
};
|
||||
|
||||
mod full;
|
||||
mod mamba;
|
||||
mod swa;
|
||||
|
||||
pub use full::FullComponent;
|
||||
pub use mamba::MambaComponent;
|
||||
pub use swa::SwaComponent;
|
||||
|
||||
/// Whether `node_id` holds the component's data on `target`, checking its
|
||||
/// device or host slot.
|
||||
pub(crate) fn node_has_component_data<K: ChildKeyType>(
|
||||
arena: &NodeArena<K>,
|
||||
node_id: NodeIdx_,
|
||||
component_type: ComponentType,
|
||||
target: EvictLayer,
|
||||
) -> bool {
|
||||
match target {
|
||||
EvictLayer::Device => arena.has_device_value(node_id, component_type),
|
||||
EvictLayer::Host => arena.has_host_value(node_id, component_type),
|
||||
EvictLayer::All => panic!("node_has_component_data: EvictLayer::All is not a single layer"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every device value of the component across all roots, concatenated.
|
||||
pub(crate) fn all_values_flatten<K: ChildKeyType>(
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
component_type: ComponentType,
|
||||
) -> Tensor {
|
||||
let mut values: Vec<Tensor> = Vec::new();
|
||||
let mut stack: Vec<NodeIdx_> = vec![tree_core.arena.root()];
|
||||
while let Some(node_id) = stack.pop() {
|
||||
let node = tree_core.arena.node(node_id);
|
||||
if let Some(value) = node.try_device_value(component_type) {
|
||||
values.push(value.shallow_clone());
|
||||
}
|
||||
stack.extend(node.children.values().copied());
|
||||
}
|
||||
if values.is_empty() {
|
||||
return tree_core.empty_device_indices.shallow_clone();
|
||||
}
|
||||
Tensor::cat(&values, 0)
|
||||
}
|
||||
|
||||
/// A per-component lock/value/eviction driver over the shared `UnifiedTreeCore`.
|
||||
pub trait TreeComponent<K: ChildKeyType> {
|
||||
/// The component this driver serves.
|
||||
fn component_type(&self) -> ComponentType;
|
||||
|
||||
/// Whether this component has device data that still needs a host backup.
|
||||
fn needs_incremental_backup(
|
||||
&self,
|
||||
_tree_core: &UnifiedTreeCore<K>,
|
||||
_node_id: NodeIdx_,
|
||||
) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Refresh this component's LRU position for `node_id` at the given walk phase.
|
||||
fn refresh_lru(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
phase: LRURefreshPhase,
|
||||
node_id: NodeIdx_,
|
||||
) {
|
||||
// Python reference — tree_component.py::TreeComponent.refresh_lru:
|
||||
// def refresh_lru(
|
||||
// self,
|
||||
// phase: LRURefreshPhase,
|
||||
// node: UnifiedTreeNode,
|
||||
// root_node: UnifiedTreeNode,
|
||||
// ) -> None:
|
||||
// ct = self.component_type
|
||||
// match phase:
|
||||
// case LRURefreshPhase.WALKDOWN:
|
||||
// if node.component_data[ct].value is None:
|
||||
// return
|
||||
// self.tree_core.lru_lists[ct].reset_node_mru(node)
|
||||
// case LRURefreshPhase.MATCH_END:
|
||||
// self.tree_core.lru_lists[ct].reset_node_and_parents_mru(
|
||||
// node, root_node, self.node_has_component_data
|
||||
// )
|
||||
// case LRURefreshPhase.INSERT_END:
|
||||
// # WALKDOWN already refreshed every node on the insert path
|
||||
// # (including the new leaf), so there is nothing more to do.
|
||||
// return
|
||||
// case _:
|
||||
// raise ValueError(f"Unknown LRURefreshPhase: {phase}")
|
||||
unimplemented!("TreeComponent.refresh_lru")
|
||||
}
|
||||
|
||||
/// Return a per-match stateful predicate deciding whether a node is a valid
|
||||
/// match boundary for this component.
|
||||
// Python reference — tree_component.py::TreeComponent.create_match_validator:
|
||||
// @abstractmethod
|
||||
// def create_match_validator(
|
||||
// self, match_device_only: bool = False
|
||||
// ) -> Callable[[UnifiedTreeNode], bool]:
|
||||
// """Return a per-match stateful predicate that decides whether a node
|
||||
// is a valid match boundary for this component.
|
||||
// Called once per match_prefix; the returned closure may carry state.
|
||||
// When match_device_only is true, host-backed nodes must not be accepted
|
||||
// as valid match boundaries.
|
||||
// - Full: returns True if the node has full component data.
|
||||
// - SWA: tracks accumulated length since last gap; returns True only
|
||||
// when the contiguous window reaches swa_sliding_window_size.
|
||||
// - Mamba: returns True iff the node has mamba component data."""
|
||||
// ...
|
||||
fn create_match_validator(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
match_device_only: bool,
|
||||
) -> Box<dyn FnMut(&UnifiedTreeCore<K>, NodeIdx_) -> bool>;
|
||||
|
||||
/// Tree-side post-processing inside the match walk (no cache access).
|
||||
fn finalize_match_result_in_tree_core(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
result: MatchResult,
|
||||
params: &MatchPrefixParams<'_, K>,
|
||||
value_chunks: &[Tensor],
|
||||
best_value_len: usize,
|
||||
) -> MatchResult {
|
||||
result
|
||||
}
|
||||
|
||||
/// Called per-node when an insert's key overlaps an existing node.
|
||||
/// Returns the index within `value_slice` from which this component
|
||||
/// consumed (took ownership of) the underlying KV pool slots.
|
||||
/// Returns `prefix_len` if nothing was consumed (default).
|
||||
/// The insert walk uses this to free only the non-consumed duplicate
|
||||
/// portion: `value_slice[dup_start..consumed_from]`.
|
||||
fn update_component_on_insert_overlap(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
prefix_len: usize,
|
||||
total_prefix_len: usize,
|
||||
value_slice: Tensor,
|
||||
params: &InsertParams<'_, K>,
|
||||
result: &mut InsertResult,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
) -> usize {
|
||||
prefix_len
|
||||
}
|
||||
|
||||
/// Called after `unevict_node_on_insert_` restores the base (Full) value
|
||||
/// on an evicted node. Aux components (e.g. SWA) override this to rebuild
|
||||
/// their own data from the freshly assigned base value when their entry
|
||||
/// is still tombstoned. Default no-op.
|
||||
fn recover_after_unevict(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
prefix_len: usize,
|
||||
total_prefix_len: usize,
|
||||
params: &InsertParams<'_, K>,
|
||||
result: &mut InsertResult,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
) {
|
||||
}
|
||||
|
||||
/// Finalize component data on the target (leaf) node after the insert
|
||||
/// walk completes. Called once per insert.
|
||||
/// - Full: no-op (full data is handled by `add_new_node_`).
|
||||
/// - SWA: for new leaves, checks whether the node straddles the SWA
|
||||
/// eviction boundary (`swa_evicted_seqlen`). If so, splits the node
|
||||
/// via `split_node_` — the parent becomes a tombstone (no SWA) and the
|
||||
/// child (the deeper portion) receives SWA data. If the entire node
|
||||
/// is within the window, sets SWA directly. If entirely outside,
|
||||
/// leaves SWA as None (tombstone).
|
||||
/// - Mamba: sets the mamba component value from params, inserts into the
|
||||
/// mamba LRU list, and increments evictable size. If the node already
|
||||
/// has mamba data, resets its LRU position instead.
|
||||
fn commit_insert_component_data(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
is_new_leaf: bool,
|
||||
params: &InsertParams<'_, K>,
|
||||
result: &mut InsertResult,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
) {
|
||||
}
|
||||
|
||||
/// Evict shallow device checkpoints beyond the per-path state cap on the
|
||||
/// tail's root path; only the Mamba component caps its states.
|
||||
fn evict_excess_path_states(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
tail_node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
}
|
||||
|
||||
/// Redistribute component data between `new_parent` and `child` when a node is
|
||||
/// split; `new_parent` is the newly created prefix node.
|
||||
// Python reference — tree_component.py::TreeComponent.redistribute_on_node_split:
|
||||
// @abstractmethod
|
||||
// def redistribute_on_node_split(
|
||||
// self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode
|
||||
// ):
|
||||
// """Redistribute component data between new_parent and child when a
|
||||
// node is split. new_parent is the newly created prefix node.
|
||||
// - Full: copies child's lock_ref to new_parent.
|
||||
// - SWA: slices (or clones) the swa value for new_parent, copies
|
||||
// lock_ref and component_uuid metadata, then syncs child's swa
|
||||
// value with its (now-trimmed) full_value.
|
||||
// - Mamba: sets new_parent's mamba value to None and lock_ref to 0
|
||||
// (mamba data stays on the original leaf, not on prefix nodes)."""
|
||||
// ...
|
||||
fn redistribute_on_node_split(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
new_parent_id: NodeIdx_,
|
||||
child_id: NodeIdx_,
|
||||
);
|
||||
|
||||
/// Free this component's KV resources on a node being evicted; returns
|
||||
/// (device_freed, host_freed) token counts.
|
||||
// Python reference — tree_component.py::TreeComponent.evict_component:
|
||||
// @abstractmethod
|
||||
// def evict_component(
|
||||
// self,
|
||||
// node: UnifiedTreeNode,
|
||||
// device_frees: dict[ComponentType, list[torch.Tensor]],
|
||||
// host_frees: dict[ComponentType, list[torch.Tensor]],
|
||||
// target: EvictLayer = EvictLayer.DEVICE,
|
||||
// ) -> tuple[int, int]:
|
||||
// """Free this component's KV resources on a node being evicted.
|
||||
//
|
||||
// *target* controls which layer(s) to evict:
|
||||
// - DEVICE: free device memory and tombstone (value = None).
|
||||
// Host data is untouched.
|
||||
// - HOST: free host memory (host_value = None).
|
||||
// Device data is untouched.
|
||||
// - ALL: free both device and host memory.
|
||||
// No tombstone — caller will delete the node.
|
||||
//
|
||||
// Returns (device_freed, host_freed) token counts."""
|
||||
// ...
|
||||
fn evict_component(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
target: EvictLayer,
|
||||
) -> (usize, usize);
|
||||
|
||||
/// Eviction priority on this node type; higher = evicted later, and evicting a
|
||||
/// component cascade-evicts every component of equal or lower priority.
|
||||
fn eviction_priority(&self, is_leaf: bool) -> i64 {
|
||||
0
|
||||
}
|
||||
|
||||
/// Begin this component's device-eviction walk (build its cursor/heap).
|
||||
fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<K>, request_cnt: usize);
|
||||
|
||||
/// Advance one eviction step and return a device leaf, if selected.
|
||||
///
|
||||
/// Implementations must return after one allocator-relevant internal
|
||||
/// mutation so the caller can drain pending frees before continuing.
|
||||
fn evict_device_next_node(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) -> Option<NodeIdx_>;
|
||||
|
||||
/// Clear this component's device-eviction walk state.
|
||||
fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<K>);
|
||||
|
||||
/// Increment component lock refs, protecting nodes from eviction.
|
||||
// Python reference — tree_component.py::TreeComponent.acquire_component_lock:
|
||||
// @abstractmethod
|
||||
// def acquire_component_lock(
|
||||
// self,
|
||||
// node: UnifiedTreeNode,
|
||||
// result: IncLockRefResult,
|
||||
// lock_host: bool = False,
|
||||
// ) -> IncLockRefResult:
|
||||
// """Increment component lock refs, protecting nodes from
|
||||
// eviction. Updates evictable → protected size on first lock.
|
||||
// - Full: path-lock — walks from node up to root, incrementing
|
||||
// lock_ref on every ancestor.
|
||||
// - SWA: path-lock — walks upward collecting swa values until the
|
||||
// sliding window is filled; records a component_uuid at the
|
||||
// boundary for release_component_lock to know where to stop.
|
||||
// - Mamba: single-node lock — only increments lock_ref on the
|
||||
// node itself (mamba state is per-leaf, not per-path).
|
||||
//
|
||||
// When ``lock_host`` is True, the lock applies to host-side state:
|
||||
// - Full: single-node host lock.
|
||||
// - SWA: host window-lock with a dedicated host UUID boundary.
|
||||
// - Mamba: single-node host lock with host LRU detach."""
|
||||
// ...
|
||||
fn acquire_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
result: IncLockRefResult,
|
||||
lock_host: bool,
|
||||
) -> IncLockRefResult;
|
||||
|
||||
/// Decrement component lock refs, un-protecting nodes.
|
||||
// Python reference — tree_component.py::TreeComponent.release_component_lock:
|
||||
// @abstractmethod
|
||||
// def release_component_lock(
|
||||
// self,
|
||||
// node: UnifiedTreeNode,
|
||||
// params: Optional[DecLockRefParams],
|
||||
// lock_host: bool = False,
|
||||
// ) -> None:
|
||||
// """Decrement component lock refs, un-protecting nodes.
|
||||
// Updates protected → evictable size when lock_ref drops to 0.
|
||||
// - Full: path-unlock — walks from node up to root, decrementing
|
||||
// lock_ref on every ancestor.
|
||||
// - SWA: path-unlock — walks upward, stopping at the node whose
|
||||
// component_uuid matches the one recorded during acquire.
|
||||
// - Mamba: single-node unlock — only decrements lock_ref on the
|
||||
// node itself.
|
||||
//
|
||||
// When ``lock_host`` is True, the inverse host-side semantics apply."""
|
||||
// ...
|
||||
fn release_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
params: Option<&DecLockRefParams>,
|
||||
lock_host: bool,
|
||||
);
|
||||
|
||||
/// Early-release the SWA lock along [node, swa_uuid_for_lock] while leaving
|
||||
/// the other components' locks intact; only the SWA component supports it.
|
||||
fn release_window_lock(
|
||||
&self,
|
||||
_tree_core: &mut UnifiedTreeCore<K>,
|
||||
_node_id: NodeIdx_,
|
||||
_swa_uuid_for_lock: Option<i64>,
|
||||
_device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
_host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
unimplemented!("release_window_lock is SWA-only")
|
||||
}
|
||||
|
||||
/// Build transfer descriptors for this component in the given phase; None when
|
||||
/// the component has nothing to transfer.
|
||||
fn build_hicache_transfers(
|
||||
&self,
|
||||
tree_core: &UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
mamba_pool_idx: Option<Tensor>,
|
||||
host_indices: Option<Tensor>,
|
||||
token_ids: Option<&[i64]>,
|
||||
prefetch_tokens: usize,
|
||||
last_hash: Option<&str>,
|
||||
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
|
||||
// Python reference — tree_component.py::TreeComponent.build_hicache_transfers:
|
||||
// def build_hicache_transfers(
|
||||
// self,
|
||||
// node: UnifiedTreeNode,
|
||||
// phase: CacheTransferPhase,
|
||||
// *,
|
||||
// mamba_pool_idx: Optional[torch.Tensor] = None,
|
||||
// host_indices: Optional[torch.Tensor] = None,
|
||||
// token_ids: Optional[Sequence[int]] = None,
|
||||
// prefetch_tokens: int = 0,
|
||||
// last_hash: Optional[str] = None,
|
||||
// ) -> Optional[list[PoolTransfer]]:
|
||||
// """Build transfer descriptors for this component in the given phase.
|
||||
// Returns None if the component has nothing to transfer."""
|
||||
// return None
|
||||
unimplemented!("TreeComponent.build_hicache_transfers")
|
||||
}
|
||||
|
||||
/// Post-transfer bookkeeping: store host indices, update LRU, etc.
|
||||
fn commit_hicache_transfer(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<K>,
|
||||
node_id: NodeIdx_,
|
||||
phase: CacheTransferPhase,
|
||||
transfers: Vec<PoolTransfer>,
|
||||
cache_actions: &mut Vec<CacheAction>,
|
||||
insert_result: Option<&mut InsertResult>,
|
||||
pool_storage_result: Option<&PoolTransferResult>,
|
||||
) {
|
||||
// Python reference — tree_component.py::TreeComponent.commit_hicache_transfer:
|
||||
// def commit_hicache_transfer(
|
||||
// self,
|
||||
// node: UnifiedTreeNode,
|
||||
// phase: CacheTransferPhase,
|
||||
// transfers: list[PoolTransfer] = (),
|
||||
// *,
|
||||
// cache_actions: list[CacheAction | ComponentAction],
|
||||
// insert_result: Optional[InsertResult] = None,
|
||||
// pool_storage_result: Optional[PoolTransferResult] = None,
|
||||
// ) -> None:
|
||||
// """Post-transfer bookkeeping: store host indices, update LRU, etc."""
|
||||
// pass
|
||||
unimplemented!("TreeComponent.commit_hicache_transfer")
|
||||
}
|
||||
|
||||
/// Reclaim host values that coexist with device values before ordinary
|
||||
/// host eviction. Called only under the write-back policy.
|
||||
fn reclaim_coexisting_host_values(
|
||||
&self,
|
||||
_tree_core: &mut UnifiedTreeCore<K>,
|
||||
_num_tokens: usize,
|
||||
_tracker: &mut HashMap<ComponentType, usize>,
|
||||
_device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
_host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
}
|
||||
|
||||
/// Evict from this component's host-side resources.
|
||||
/// Called by HostPoolGroup when the host pool is full.
|
||||
/// Default no-op for components without host storage.
|
||||
fn drive_host_eviction(
|
||||
&self,
|
||||
_tree_core: &mut UnifiedTreeCore<K>,
|
||||
_num_tokens: usize,
|
||||
_tracker: &mut HashMap<ComponentType, usize>,
|
||||
_device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
_host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
}
|
||||
}
|
||||
|
||||
// Tree component types.
|
||||
|
||||
/// The tree components; discriminants define the per-component array indexes.
|
||||
#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
|
||||
pub enum ComponentType {
|
||||
Full = 0,
|
||||
Swa = 1,
|
||||
Mamba = 2,
|
||||
}
|
||||
|
||||
/// Short call-site aliases for the component types.
|
||||
pub const FULL: ComponentType = ComponentType::Full;
|
||||
pub const SWA: ComponentType = ComponentType::Swa;
|
||||
pub const MAMBA: ComponentType = ComponentType::Mamba;
|
||||
|
||||
/// The base component every tree runs; the others are auxiliary.
|
||||
pub const BASE_COMPONENT_TYPE: ComponentType = ComponentType::Full;
|
||||
|
||||
/// Slots per tier — the arrays are sized to this, not the enabled subset.
|
||||
pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1;
|
||||
|
||||
impl ComponentType {
|
||||
/// Index into a per-component array.
|
||||
pub const fn idx(self) -> usize {
|
||||
self as usize
|
||||
}
|
||||
|
||||
/// Whether the component stores one state slot per node (Mamba) instead of
|
||||
/// one row per key atom.
|
||||
pub fn single_value_per_node(self) -> bool {
|
||||
matches!(self, ComponentType::Mamba)
|
||||
}
|
||||
|
||||
/// The component at a per-component array index; panics out of range.
|
||||
pub fn from_idx(idx: usize) -> ComponentType {
|
||||
match idx {
|
||||
0 => ComponentType::Full,
|
||||
1 => ComponentType::Swa,
|
||||
2 => ComponentType::Mamba,
|
||||
_ => panic!("from_idx: {idx} is not a component index"),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
#[path = "../tests/components/base.rs"]
|
||||
mod tests;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,27 @@
|
||||
//! Rust radix tree core for SGLang's KV cache, exposed to Python as `mem_cache`.
|
||||
// TODO(Jialin): Replace recoverable panics with explicit Rust errors and map
|
||||
// them to PyErr at the Python boundary.
|
||||
#![allow(
|
||||
dead_code,
|
||||
unsafe_op_in_unsafe_fn,
|
||||
clippy::unwrap_used,
|
||||
clippy::expect_used,
|
||||
clippy::owned_cow,
|
||||
clippy::panic,
|
||||
clippy::print_stdout,
|
||||
clippy::too_many_arguments,
|
||||
clippy::type_complexity,
|
||||
clippy::unimplemented,
|
||||
clippy::unreachable,
|
||||
clippy::useless_conversion
|
||||
)]
|
||||
|
||||
mod components;
|
||||
mod node;
|
||||
#[cfg(feature = "python-extension")]
|
||||
mod python_bindings;
|
||||
#[cfg(test)]
|
||||
#[path = "tests/test_utils.rs"]
|
||||
pub(crate) mod test_utils;
|
||||
mod unified_lru_list;
|
||||
mod unified_tree_core;
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,160 @@
|
||||
use super::*;
|
||||
use crate::components::FULL;
|
||||
use crate::unified_tree_core::CacheInitParams;
|
||||
|
||||
// Test-only component exercising the trait defaults; abstract hooks stay unimplemented.
|
||||
struct DefaultComponentForTest;
|
||||
|
||||
impl TreeComponent<Vec<i64>> for DefaultComponentForTest {
|
||||
fn component_type(&self) -> ComponentType {
|
||||
FULL
|
||||
}
|
||||
|
||||
fn create_match_validator(
|
||||
&self,
|
||||
_tree_core: &UnifiedTreeCore<Vec<i64>>,
|
||||
match_device_only: bool,
|
||||
) -> Box<dyn FnMut(&UnifiedTreeCore<Vec<i64>>, NodeIdx_) -> bool> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn redistribute_on_node_split(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
|
||||
new_parent_id: NodeIdx_,
|
||||
child_id: NodeIdx_,
|
||||
) {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn evict_component(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
|
||||
node_id: NodeIdx_,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
target: EvictLayer,
|
||||
) -> (usize, usize) {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn evict_device_start(&self, tree_core: &mut UnifiedTreeCore<Vec<i64>>, request_cnt: usize) {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn evict_device_next_node(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) -> Option<NodeIdx_> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn evict_device_end(&self, tree_core: &mut UnifiedTreeCore<Vec<i64>>) {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn acquire_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
|
||||
node_id: NodeIdx_,
|
||||
result: IncLockRefResult,
|
||||
lock_host: bool,
|
||||
) -> IncLockRefResult {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn release_component_lock(
|
||||
&self,
|
||||
tree_core: &mut UnifiedTreeCore<Vec<i64>>,
|
||||
node_id: NodeIdx_,
|
||||
params: Option<&DecLockRefParams>,
|
||||
lock_host: bool,
|
||||
) {
|
||||
unimplemented!()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_overlap_default_consumes_nothing() {
|
||||
let mut tc: UnifiedTreeCore<Vec<i64>> =
|
||||
UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
|
||||
let root = tc.arena.root();
|
||||
let consumed_from = DefaultComponentForTest.update_component_on_insert_overlap(
|
||||
&mut tc,
|
||||
root,
|
||||
/* prefix_len = */ 3,
|
||||
/* total_prefix_len = */ 0,
|
||||
Tensor::from_slice(&[0i64, 1, 2]),
|
||||
&InsertParams {
|
||||
key: &vec![0, 1, 2],
|
||||
namespace: Default::default(),
|
||||
value: Tensor::from_slice(&[0i64, 1, 2]),
|
||||
mamba_value: None,
|
||||
prev_prefix_len: 0,
|
||||
swa_evicted_seqlen: 0,
|
||||
chunked: false,
|
||||
priority: 0,
|
||||
track_adopted_ranges: false,
|
||||
},
|
||||
&mut InsertResult::default(),
|
||||
&mut Vec::new(),
|
||||
);
|
||||
// Nothing consumed: the whole overlap stays freeable as duplicates.
|
||||
assert_eq!(consumed_from, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finalize_match_result_default_returns_result_unchanged() {
|
||||
let tc: UnifiedTreeCore<Vec<i64>> =
|
||||
UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
|
||||
let result = MatchResult {
|
||||
last_device_node_id: 3,
|
||||
best_match_node_id: 7,
|
||||
host_hit_length: 11,
|
||||
..tc.empty_match_result()
|
||||
};
|
||||
let out = DefaultComponentForTest.finalize_match_result_in_tree_core(
|
||||
&tc,
|
||||
result,
|
||||
&MatchPrefixParams {
|
||||
key: &Vec::new(),
|
||||
namespace: Default::default(),
|
||||
},
|
||||
&[],
|
||||
0,
|
||||
);
|
||||
assert_eq!(out.last_device_node_id, 3);
|
||||
assert_eq!(out.best_match_node_id, 7);
|
||||
assert_eq!(out.host_hit_length, 11);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drive_host_eviction_default_is_a_noop() {
|
||||
let mut tc: UnifiedTreeCore<Vec<i64>> =
|
||||
UnifiedTreeCore::new(CacheInitParams::default(), vec![FULL]);
|
||||
let mut tracker = HashMap::from([(FULL, 5usize)]);
|
||||
let mut device_frees = HashMap::new();
|
||||
let mut host_frees = HashMap::new();
|
||||
DefaultComponentForTest.drive_host_eviction(
|
||||
&mut tc,
|
||||
/* num_tokens = */ 100,
|
||||
&mut tracker,
|
||||
&mut device_frees,
|
||||
&mut host_frees,
|
||||
);
|
||||
assert_eq!(tracker[&FULL], 5);
|
||||
assert!(device_frees.is_empty());
|
||||
assert!(host_frees.is_empty());
|
||||
}
|
||||
|
||||
// Component types.
|
||||
|
||||
#[test]
|
||||
fn idx_matches_discriminants() {
|
||||
assert_eq!(ComponentType::Full.idx(), 0);
|
||||
assert_eq!(ComponentType::Swa.idx(), 1);
|
||||
assert_eq!(ComponentType::Mamba.idx(), 2);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,46 @@
|
||||
//! Shared helpers for the crate's unit tests.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use tch::Tensor;
|
||||
|
||||
use crate::components::ComponentType;
|
||||
use crate::unified_tree_core::{CacheAction, EvictionStepResult};
|
||||
|
||||
/// Fold an eviction step into a caller's running accumulators (the Controller
|
||||
/// consumption contract: deltas add, freed tensors append).
|
||||
pub(crate) fn accumulate_step(
|
||||
step: EvictionStepResult,
|
||||
tracker: &mut HashMap<ComponentType, usize>,
|
||||
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
|
||||
) {
|
||||
for (ct, delta) in step.tracker {
|
||||
*tracker.entry(ct).or_insert(0) += delta;
|
||||
}
|
||||
for (ct, tensors) in step.device_frees {
|
||||
device_frees.entry(ct).or_default().extend(tensors);
|
||||
}
|
||||
for (ct, tensors) in step.host_frees {
|
||||
host_frees.entry(ct).or_default().extend(tensors);
|
||||
}
|
||||
}
|
||||
|
||||
/// Short variant names for diagnosing an action sequence's shape.
|
||||
pub(crate) fn action_kinds(actions: &[CacheAction]) -> Vec<&'static str> {
|
||||
actions
|
||||
.iter()
|
||||
.map(|action| match action {
|
||||
CacheAction::FreeDeviceKV(_) => "FreeDeviceKV",
|
||||
CacheAction::FreeDeviceKVFullOnly(_) => "FreeDeviceKVFullOnly",
|
||||
CacheAction::BackupKV(_) => "BackupKV",
|
||||
CacheAction::ReplaceWriteThroughOnNodeSplit { .. } => "ReplaceWriteThroughOnNodeSplit",
|
||||
CacheAction::MambaEvictExcessPathStates { .. } => "MambaEvictExcessPathStates",
|
||||
CacheAction::FreeComponentDeviceSlot { .. } => "FreeComponentDeviceSlot",
|
||||
CacheAction::FreeComponentHostSlot { .. } => "FreeComponentHostSlot",
|
||||
CacheAction::RebuildFullToSwaMapping { .. } => "RebuildFullToSwaMapping",
|
||||
CacheAction::RecoverSwaWithLockedFull { .. } => "RecoverSwaWithLockedFull",
|
||||
CacheAction::SwaRebuild { .. } => "SwaRebuild",
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -0,0 +1,708 @@
|
||||
use super::*;
|
||||
use crate::components::FULL;
|
||||
use crate::node::{NodeArena, NodeIdx_, ValueSlotIdx};
|
||||
|
||||
fn order(list: &UnifiedLRUList) -> Vec<NodeIdx_> {
|
||||
list.iter().collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fresh_list_reads_are_empty() {
|
||||
let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
assert_eq!(list.get_lru_where(|_| true), None);
|
||||
assert_eq!(list.iter().count(), 0);
|
||||
assert_eq!(list.len(), 0);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_mru_orders_most_recent_first() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(20), NodeIdx_(10)]);
|
||||
assert_eq!(list.len(), 3);
|
||||
assert!(list.in_list(Some(NodeIdx_(10))));
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "already in the LRU list")]
|
||||
fn insert_mru_panics_when_already_a_member() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(1));
|
||||
list.insert_mru(NodeIdx_(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_node_updates_membership_immediately() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
list.remove_node(NodeIdx_(20));
|
||||
assert_eq!(order(&list), vec![NodeIdx_(30), NodeIdx_(10)]);
|
||||
assert!(!list.in_list(Some(NodeIdx_(20))));
|
||||
assert_eq!(list.len(), 2);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "index out of bounds")]
|
||||
fn remove_node_panics_when_absent() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.remove_node(NodeIdx_(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "not in the LRU list")]
|
||||
fn remove_node_panics_on_a_node_removed_earlier() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.remove_node(NodeIdx_(10));
|
||||
// The cell is still allocated but reset; membership must gate the removal.
|
||||
list.remove_node(NodeIdx_(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "not in the LRU list")]
|
||||
fn remove_node_panics_on_an_unlisted_cell() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.remove_node(NodeIdx_(10));
|
||||
list.remove_node_(UnifiedLRUList::cell_of_(NodeIdx_(10)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "already in the LRU list")]
|
||||
fn add_node_panics_on_a_linked_cell() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.add_node_(UnifiedLRUList::cell_of_(NodeIdx_(10)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn removed_nodes_can_be_reinserted() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.remove_node(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(20)]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_node_mru_moves_a_member_to_the_front() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
list.reset_node_mru(NodeIdx_(10));
|
||||
assert_eq!(order(&list), vec![NodeIdx_(10), NodeIdx_(30), NodeIdx_(20)]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "index out of bounds")]
|
||||
fn reset_node_mru_panics_on_a_non_member() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.reset_node_mru(NodeIdx_(7));
|
||||
}
|
||||
|
||||
// Arena chain root -> a -> b -> c plus a sibling, two atoms per key.
|
||||
fn arena_chain() -> (
|
||||
NodeArena<Vec<i64>>,
|
||||
NodeIdx_,
|
||||
NodeIdx_,
|
||||
NodeIdx_,
|
||||
NodeIdx_,
|
||||
NodeIdx_,
|
||||
) {
|
||||
let mut arena = NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1);
|
||||
let root = arena.root();
|
||||
let a = arena
|
||||
.alloc_child(
|
||||
root,
|
||||
/* key = */ vec![1, 11],
|
||||
/* priority = */ 0,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
let b = arena
|
||||
.alloc_child(
|
||||
a,
|
||||
/* key = */ vec![2, 22],
|
||||
/* priority = */ 0,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
let c = arena
|
||||
.alloc_child(
|
||||
b,
|
||||
/* key = */ vec![3, 33],
|
||||
/* priority = */ 0,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
let other = arena
|
||||
.alloc_child(
|
||||
root,
|
||||
/* key = */ vec![9, 99],
|
||||
/* priority = */ 0,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
(arena, root, a, b, c, other)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_parents_mru_reranks_included_nodes_deepest_first() {
|
||||
let (arena, _root, a, b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(c);
|
||||
list.insert_mru(other);
|
||||
// b is excluded and skipped; c then a become the MRU run.
|
||||
list.reset_node_and_parents_mru(c, &arena, |node| node.idx != b);
|
||||
assert_eq!(order(&list), vec![c, a, other]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_parents_mru_reranks_ancestors_when_the_deepest_is_excluded() {
|
||||
let (arena, _root, a, _b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(other);
|
||||
// c and b are excluded; a alone becomes the new MRU head.
|
||||
list.reset_node_and_parents_mru(c, &arena, |node| node.idx == a);
|
||||
assert_eq!(order(&list), vec![a, other]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_walks_are_noops_when_node_is_the_root() {
|
||||
let (arena, root, a, _b, _c, _other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.reset_node_and_parents_mru(root, &arena, |_| true);
|
||||
list.reset_node_and_window_ancestors_mru(root, 4, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![a]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "not in the LRU list")]
|
||||
fn reset_parents_mru_panics_on_an_unlisted_included_node() {
|
||||
let (arena, _root, _a, _b, c, _other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(c);
|
||||
list.remove_node(c);
|
||||
list.reset_node_and_parents_mru(c, &arena, |_| true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_window_ancestors_mru_stops_at_the_window() {
|
||||
let (arena, _root, a, b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(b);
|
||||
list.insert_mru(c);
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(other);
|
||||
// A window of 4 atoms covers c and b; a stays put beyond it.
|
||||
list.reset_node_and_window_ancestors_mru(c, 4, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![c, b, other, a]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_window_ancestors_mru_includes_the_straddling_ancestor() {
|
||||
let (arena, _root, a, b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(b);
|
||||
list.insert_mru(c);
|
||||
list.insert_mru(other);
|
||||
// A window of 5 atoms ends mid-a: the straddling ancestor is still included.
|
||||
list.reset_node_and_window_ancestors_mru(c, 5, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![c, b, a, other]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_walks_stop_at_the_salted_chains_root() {
|
||||
let mut arena: NodeArena<Vec<i64>> =
|
||||
NodeArena::new(vec![crate::components::FULL], /* page_size = */ 1);
|
||||
let named = arena.root();
|
||||
let a = arena
|
||||
.alloc_child(
|
||||
named,
|
||||
/* key = */ vec![1, 11],
|
||||
/* priority = */ 0,
|
||||
Some("lora-1"),
|
||||
)
|
||||
.unwrap();
|
||||
let b = arena
|
||||
.alloc_child(
|
||||
a,
|
||||
/* key = */ vec![2, 22],
|
||||
/* priority = */ 0,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(b);
|
||||
// Both walks terminate at the root without visiting it.
|
||||
list.reset_node_and_parents_mru(b, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![b, a]);
|
||||
list.reset_node_and_window_ancestors_mru(b, 100, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![b, a]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_lru_no_lock_returns_the_lru_most_unlocked_member() {
|
||||
let (mut arena, _root, a, b, c, _other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(b);
|
||||
list.insert_mru(c);
|
||||
assert_eq!(list.get_lru_no_lock(&arena), Some(a));
|
||||
// A lock on the list's own slot hides the LRU end from the walker.
|
||||
arena
|
||||
.node_mut(a)
|
||||
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
|
||||
assert_eq!(list.get_lru_no_lock(&arena), Some(b));
|
||||
arena
|
||||
.node_mut(b)
|
||||
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
|
||||
arena
|
||||
.node_mut(c)
|
||||
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
|
||||
assert_eq!(list.get_lru_no_lock(&arena), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_prev_no_lock_skips_locked_members_toward_the_mru_end() {
|
||||
let (mut arena, _root, a, b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(b);
|
||||
list.insert_mru(c);
|
||||
list.insert_mru(other);
|
||||
assert_eq!(list.get_prev_no_lock(a, &arena), Some(b));
|
||||
// The locked b is skipped; from the MRU end there is no predecessor left.
|
||||
arena
|
||||
.node_mut(b)
|
||||
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
|
||||
assert_eq!(list.get_prev_no_lock(a, &arena), Some(c));
|
||||
assert_eq!(list.get_prev_no_lock(other, &arena), None);
|
||||
// A lock on a different slot does not gate this list's walker.
|
||||
arena.node_mut(c).set_lock_ref_(ValueSlotIdx::host(FULL), 1);
|
||||
assert_eq!(list.get_prev_no_lock(a, &arena), Some(c));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_window_accumulation_counts_excluded_nodes() {
|
||||
let (arena, _root, a, b, c, other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(c);
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(other);
|
||||
// b is excluded but its atoms still consume the window, keeping a out of reach.
|
||||
list.reset_node_and_window_ancestors_mru(c, 4, &arena, |node| node.idx != b);
|
||||
assert_eq!(order(&list), vec![c, other, a]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "not in the LRU list")]
|
||||
fn reset_node_mru_panics_on_a_node_removed_earlier() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.remove_node(NodeIdx_(10));
|
||||
// The cell is still allocated but unlisted; the gated read must reject it.
|
||||
list.reset_node_mru(NodeIdx_(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn in_list_is_false_for_none_and_non_members() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
assert!(!list.in_list(None));
|
||||
assert!(!list.in_list(Some(NodeIdx_(5))));
|
||||
list.insert_mru(NodeIdx_(5));
|
||||
assert!(list.in_list(Some(NodeIdx_(5))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_lru_where_walks_from_the_tail() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
assert_eq!(list.get_lru_where(|_| true), Some(NodeIdx_(10)));
|
||||
assert_eq!(
|
||||
list.get_lru_where(|id| id != NodeIdx_(10)),
|
||||
Some(NodeIdx_(20))
|
||||
);
|
||||
assert_eq!(list.get_lru_where(|_| false), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_prev_where_walks_toward_the_head_from_a_member() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
// Order is [30, 20, 10]; 10's predecessors are 20 then 30.
|
||||
assert_eq!(
|
||||
list.get_prev_where(NodeIdx_(10), |_| true),
|
||||
Some(NodeIdx_(20))
|
||||
);
|
||||
assert_eq!(
|
||||
list.get_prev_where(NodeIdx_(10), |id| id != NodeIdx_(20)),
|
||||
Some(NodeIdx_(30))
|
||||
);
|
||||
assert_eq!(list.get_prev_where(NodeIdx_(30), |_| true), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "index out of bounds")]
|
||||
fn get_prev_where_panics_on_a_non_member() {
|
||||
let list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.get_prev_where(NodeIdx_(7), |_| true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "not in the LRU list")]
|
||||
fn get_prev_where_panics_on_a_node_removed_earlier() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.remove_node(NodeIdx_(10));
|
||||
// The cell is still allocated but unlisted; the gated read must reject it.
|
||||
list.get_prev_where(NodeIdx_(10), |_| true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_prev_before_remove_keeps_the_walk_consistent() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.insert_mru(NodeIdx_(30));
|
||||
// The eviction-cursor contract: compute the predecessor, then remove.
|
||||
let next = list.get_prev_where(NodeIdx_(10), |_| true);
|
||||
list.remove_node(NodeIdx_(10));
|
||||
assert_eq!(next, Some(NodeIdx_(20)));
|
||||
assert!(list.in_list(next));
|
||||
assert_eq!(
|
||||
list.get_prev_where(NodeIdx_(20), |_| true),
|
||||
Some(NodeIdx_(30))
|
||||
);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_mru_grows_the_cell_table_one_by_one() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(0));
|
||||
list.insert_mru(NodeIdx_(1));
|
||||
list.insert_mru(NodeIdx_(2));
|
||||
assert_eq!(order(&list), vec![NodeIdx_(2), NodeIdx_(1), NodeIdx_(0)]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "broken prev link")]
|
||||
fn validate_panics_on_a_corrupted_prev_link() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
let cell = UnifiedLRUList::cell_of_(NodeIdx_(10));
|
||||
list.cells[cell.0].prev = cell;
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "membership mismatch")]
|
||||
fn validate_panics_on_a_linked_cell_without_the_flag() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].in_list = false;
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "membership mismatch")]
|
||||
fn validate_panics_on_a_flagged_unlinked_cell() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.insert_mru(NodeIdx_(20));
|
||||
list.remove_node(NodeIdx_(20));
|
||||
// The cell is reset but a stray flag claims membership.
|
||||
list.cells[UnifiedLRUList::cell_of_(NodeIdx_(20)).0].in_list = true;
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_window_ancestors_mru_is_a_noop_on_a_zero_window() {
|
||||
let (arena, _root, a, _b, c, _other) = arena_chain();
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(a);
|
||||
list.insert_mru(c);
|
||||
list.reset_node_and_window_ancestors_mru(c, 0, &arena, |_| true);
|
||||
assert_eq!(order(&list), vec![c, a]);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "out of bounds")]
|
||||
fn validate_panics_on_an_out_of_range_link() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.cells[UnifiedLRUList::cell_of_(NodeIdx_(10)).0].next = CellId(99);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "length mismatch")]
|
||||
fn validate_panics_on_a_desynced_member_counter() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(10));
|
||||
list.len = 2;
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn len_drops_to_zero_after_all_members_removed() {
|
||||
let mut list = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
list.insert_mru(NodeIdx_(1));
|
||||
list.remove_node(NodeIdx_(1));
|
||||
assert_eq!(list.len(), 0);
|
||||
assert_eq!(list.iter().count(), 0);
|
||||
list.validate();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_accepts_a_clean_list() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.insert_mru(NodeIdx_(1));
|
||||
lru.insert_mru(NodeIdx_(2));
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert!(errors.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_a_broken_prev() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.insert_mru(NodeIdx_(1));
|
||||
lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].prev = UnifiedLRUList::cell_of_(NodeIdx_(0));
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert!(errors[0].contains("broken prev at node 0"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_an_unflagged_member() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.insert_mru(NodeIdx_(1));
|
||||
lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = false;
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert!(errors[0].contains("node 0 in list not flagged"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_a_cycle() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.insert_mru(NodeIdx_(1));
|
||||
// 0's next loops back to 1 instead of reaching the tail.
|
||||
lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = UnifiedLRUList::cell_of_(NodeIdx_(1));
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert!(errors.iter().any(|e| e.contains("cycle at node 1")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_a_count_mismatch() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.len = 2;
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert!(errors[0].contains("list=1 != len=2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_an_out_of_bounds_link() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].next = CellId(999);
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert!(errors.iter().any(|e| e.contains("cell 999 out of bounds")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_a_broken_tail_backlink() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.cells[TAIL.0].prev = HEAD;
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert!(errors[0].contains("broken tail backlink"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_linked_list_reports_a_flagged_unreachable_cell() {
|
||||
let mut lru = UnifiedLRUList::new(ValueSlotIdx::device(FULL));
|
||||
lru.insert_mru(NodeIdx_(0));
|
||||
lru.insert_mru(NodeIdx_(1));
|
||||
lru.remove_node(NodeIdx_(0));
|
||||
// Re-flag the unlinked cell without relinking it.
|
||||
lru.cells[UnifiedLRUList::cell_of_(NodeIdx_(0)).0].in_list = true;
|
||||
let mut errors = Vec::new();
|
||||
lru.check_linked_list_("[t]", &mut errors);
|
||||
assert!(
|
||||
errors
|
||||
.iter()
|
||||
.any(|e| e.contains("node 0 flagged but unreachable"))
|
||||
);
|
||||
}
|
||||
|
||||
// Eviction priority keys.
|
||||
|
||||
// A node with distinct field values: last_access 5, creation 7, hits 3, priority 9.
|
||||
fn arena_with_node() -> (NodeArena<Vec<i64>>, NodeIdx_) {
|
||||
let mut arena: NodeArena<Vec<i64>> = NodeArena::new(vec![FULL], /* page_size = */ 1);
|
||||
let root = arena.root();
|
||||
let a = arena
|
||||
.alloc_child(
|
||||
root,
|
||||
/* key = */ vec![1],
|
||||
/* priority = */ 9,
|
||||
/* extra_key = */ None,
|
||||
)
|
||||
.unwrap();
|
||||
let node = arena.node_mut(a);
|
||||
node.last_access_counter = 5;
|
||||
node.creation_counter = 7;
|
||||
node.hit_count = 3;
|
||||
(arena, NodeIdx_(a.0))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn each_strategy_maps_its_node_fields_into_the_key() {
|
||||
let (arena, a) = arena_with_node();
|
||||
let node = arena.node(NodeIdx_(a.0));
|
||||
assert_eq!(LruStrategy.get_priority(node), PriorityKey(5, 0));
|
||||
assert_eq!(LfuStrategy.get_priority(node), PriorityKey(3, 5));
|
||||
assert_eq!(FifoStrategy.get_priority(node), PriorityKey(7, 0));
|
||||
assert_eq!(MruStrategy.get_priority(node), PriorityKey(-5, 0));
|
||||
assert_eq!(FiloStrategy.get_priority(node), PriorityKey(-7, 0));
|
||||
assert_eq!(PriorityStrategy.get_priority(node), PriorityKey(9, 5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slru_segments_on_the_protected_threshold() {
|
||||
let (mut arena, a) = arena_with_node();
|
||||
let slru = SlruStrategy {
|
||||
protected_threshold: 2,
|
||||
};
|
||||
// 3 hits >= threshold 2: protected segment.
|
||||
assert_eq!(
|
||||
slru.get_priority(arena.node(NodeIdx_(a.0))),
|
||||
PriorityKey(1, 5)
|
||||
);
|
||||
// Exactly at the threshold counts as protected.
|
||||
arena.node_mut(NodeIdx_(a.0)).hit_count = 2;
|
||||
assert_eq!(
|
||||
slru.get_priority(arena.node(NodeIdx_(a.0))),
|
||||
PriorityKey(1, 5)
|
||||
);
|
||||
arena.node_mut(NodeIdx_(a.0)).hit_count = 1;
|
||||
assert_eq!(
|
||||
slru.get_priority(arena.node(NodeIdx_(a.0))),
|
||||
PriorityKey(0, 5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_eviction_strategy_resolves_each_policy_name() {
|
||||
let (arena, a) = arena_with_node();
|
||||
let node = arena.node(NodeIdx_(a.0));
|
||||
// Distinct node fields make each policy's key identify its strategy.
|
||||
let cases = [
|
||||
("lru", PriorityKey(5, 0)),
|
||||
("LFU", PriorityKey(3, 5)),
|
||||
("fifo", PriorityKey(7, 0)),
|
||||
("mru", PriorityKey(-5, 0)),
|
||||
("filo", PriorityKey(-7, 0)),
|
||||
("priority", PriorityKey(9, 5)),
|
||||
("slru", PriorityKey(1, 5)),
|
||||
];
|
||||
for (policy, expected) in cases {
|
||||
assert_eq!(
|
||||
get_eviction_strategy::<Vec<i64>>(policy).get_priority(node),
|
||||
expected,
|
||||
"policy {policy}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_policy_names_are_case_insensitive() {
|
||||
let (arena, a) = arena_with_node();
|
||||
let node = arena.node(NodeIdx_(a.0));
|
||||
// Mixed-case names resolve to the same strategies as their lowercase forms.
|
||||
assert_eq!(
|
||||
get_eviction_strategy::<Vec<i64>>("LRU").get_priority(node),
|
||||
PriorityKey(5, 0)
|
||||
);
|
||||
assert_eq!(
|
||||
get_eviction_strategy::<Vec<i64>>("Priority").get_priority(node),
|
||||
PriorityKey(9, 5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_eviction_strategy_slru_default_threshold_is_two() {
|
||||
let (mut arena, a) = arena_with_node();
|
||||
let slru = get_eviction_strategy::<Vec<i64>>("slru");
|
||||
// Exactly 2 hits is protected under the factory default; 1 is not.
|
||||
arena.node_mut(NodeIdx_(a.0)).hit_count = 2;
|
||||
assert_eq!(
|
||||
slru.get_priority(arena.node(NodeIdx_(a.0))),
|
||||
PriorityKey(1, 5)
|
||||
);
|
||||
arena.node_mut(NodeIdx_(a.0)).hit_count = 1;
|
||||
assert_eq!(
|
||||
slru.get_priority(arena.node(NodeIdx_(a.0))),
|
||||
PriorityKey(0, 5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Unknown eviction policy: random. Supported policies:")]
|
||||
fn get_eviction_strategy_panics_on_an_unknown_policy() {
|
||||
get_eviction_strategy::<Vec<i64>>("Random");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn priority_keys_order_lexicographically() {
|
||||
assert!(PriorityKey(0, 9) < PriorityKey(1, 0));
|
||||
assert!(PriorityKey(1, 2) < PriorityKey(1, 3));
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,501 @@
|
||||
//! Self-contained LRU order over `NodeIdx_`s: MRU at the head side, LRU at the
|
||||
//! tail side. Node semantics stay with callers through predicates; the reset
|
||||
//! walks read parent links from the arena.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::node::ChildKeyType;
|
||||
use crate::node::Node;
|
||||
use crate::node::NodeArena;
|
||||
use crate::node::{NodeIdx_, ValueSlotIdx};
|
||||
|
||||
/// Index into the cell table; distinct from `NodeIdx_` so shifted and unshifted
|
||||
/// ids cannot be mixed.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
|
||||
struct CellId(usize);
|
||||
|
||||
/// Head sentinel cell.
|
||||
const HEAD: CellId = CellId(0);
|
||||
/// Tail sentinel cell.
|
||||
const TAIL: CellId = CellId(1);
|
||||
/// Table offset: node ids map to cells after the sentinels.
|
||||
const OFFSET: usize = 2;
|
||||
|
||||
/// One doubly-linked cell; a node's cell lives at `NodeIdx_ + OFFSET`.
|
||||
#[derive(Clone, Copy, Default)]
|
||||
struct Cell {
|
||||
prev: CellId,
|
||||
next: CellId,
|
||||
in_list: bool,
|
||||
}
|
||||
|
||||
/// LRU list over `NodeIdx_`s, with head/tail sentinel cells keeping the link
|
||||
/// operations branchless. External APIs take `NodeIdx_`s; internal (`_`-suffixed)
|
||||
/// interfaces work on `CellId`s.
|
||||
pub struct UnifiedLRUList {
|
||||
/// The (component × tier) value slot whose lock gates this list's walkers.
|
||||
slot: ValueSlotIdx,
|
||||
/// Cell table indexed by `NodeIdx_ + OFFSET`; cells 0/1 are the sentinels.
|
||||
cells: Vec<Cell>,
|
||||
/// Number of member cells, excluding the sentinels.
|
||||
len: usize,
|
||||
}
|
||||
|
||||
impl UnifiedLRUList {
|
||||
pub fn new(slot: ValueSlotIdx) -> Self {
|
||||
UnifiedLRUList {
|
||||
slot,
|
||||
cells: vec![
|
||||
// Sentinels link to each other and stay permanently flagged so
|
||||
// the gated cell accessors admit them.
|
||||
Cell {
|
||||
prev: HEAD,
|
||||
next: TAIL,
|
||||
in_list: true,
|
||||
},
|
||||
Cell {
|
||||
prev: HEAD,
|
||||
next: TAIL,
|
||||
in_list: true,
|
||||
},
|
||||
],
|
||||
len: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// ==== List operations ====
|
||||
|
||||
fn add_node_after_(&mut self, prev: CellId, cell: CellId) {
|
||||
self.new_cell_(cell);
|
||||
let next = self.cell_(prev).next;
|
||||
self.connect_(cell, next);
|
||||
self.connect_(prev, cell);
|
||||
}
|
||||
|
||||
fn add_node_(&mut self, cell: CellId) {
|
||||
self.add_node_after_(HEAD, cell);
|
||||
}
|
||||
|
||||
fn remove_node_(&mut self, cell: CellId) {
|
||||
let Cell { prev, next, .. } = *self.cell_(cell);
|
||||
self.connect_(prev, next);
|
||||
// Unflag the cell; the stale prev/next are never read while unlisted.
|
||||
self.cell_mut_(cell).in_list = false;
|
||||
self.len -= 1;
|
||||
}
|
||||
|
||||
// ==== NodeIdx_ <-> CellId plumbing ====
|
||||
|
||||
/// The node's cell slot; the only `NodeIdx_` -> `CellId` crossing.
|
||||
fn cell_of_(node_id: NodeIdx_) -> CellId {
|
||||
CellId(node_id.0 + OFFSET)
|
||||
}
|
||||
|
||||
/// The cell's node; the only `CellId` -> `NodeIdx_` crossing.
|
||||
fn node_of_(cell: CellId) -> NodeIdx_ {
|
||||
NodeIdx_(cell.0 - OFFSET)
|
||||
}
|
||||
|
||||
/// The cell, asserting it is linked (sentinels always are).
|
||||
#[track_caller]
|
||||
fn cell_(&self, id: CellId) -> &Cell {
|
||||
let cell = &self.cells[id.0];
|
||||
assert!(
|
||||
cell.in_list,
|
||||
"node {} not in the LRU list",
|
||||
Self::node_of_(id)
|
||||
);
|
||||
cell
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn cell_mut_(&mut self, id: CellId) -> &mut Cell {
|
||||
let cell = &mut self.cells[id.0];
|
||||
assert!(
|
||||
cell.in_list,
|
||||
"node {} not in the LRU list",
|
||||
Self::node_of_(id)
|
||||
);
|
||||
cell
|
||||
}
|
||||
|
||||
/// Admit an unlisted cell: grow the table to cover it, then flag and count
|
||||
/// it before any connections.
|
||||
fn new_cell_(&mut self, cell: CellId) {
|
||||
if cell.0 >= self.cells.len() {
|
||||
self.cells.resize(cell.0 + 1, Cell::default());
|
||||
}
|
||||
assert!(
|
||||
!self.cells[cell.0].in_list,
|
||||
"new_cell_: cell {cell:?} already in the LRU list"
|
||||
);
|
||||
self.cells[cell.0].in_list = true;
|
||||
self.len += 1;
|
||||
}
|
||||
|
||||
/// Whether the cell is linked into the list; safe on cells beyond the table.
|
||||
fn in_list_(&self, cell: CellId) -> bool {
|
||||
self.cells.get(cell.0).is_some_and(|cell| cell.in_list)
|
||||
}
|
||||
|
||||
/// Link `a -> b`.
|
||||
fn connect_(&mut self, a: CellId, b: CellId) {
|
||||
self.cell_mut_(a).next = b;
|
||||
self.cell_mut_(b).prev = a;
|
||||
}
|
||||
|
||||
/// Insert a node as the most-recently-used; panics if already a member.
|
||||
pub fn insert_mru(&mut self, node_id: NodeIdx_) {
|
||||
self.add_node_(Self::cell_of_(node_id));
|
||||
}
|
||||
|
||||
/// Remove a member node, resetting its cell; panics if not a member.
|
||||
pub fn remove_node(&mut self, node_id: NodeIdx_) {
|
||||
self.remove_node_(Self::cell_of_(node_id));
|
||||
}
|
||||
|
||||
/// Move a member node back to the most-recently-used position.
|
||||
pub fn reset_node_mru(&mut self, node_id: NodeIdx_) {
|
||||
let cell = Self::cell_of_(node_id);
|
||||
self.remove_node_(cell);
|
||||
self.add_node_(cell);
|
||||
}
|
||||
|
||||
/// Re-rank the `should_include` nodes from `node_id` up to its root
|
||||
/// (exclusive) as the MRU run, deepest first.
|
||||
pub fn reset_node_and_parents_mru<K: ChildKeyType>(
|
||||
&mut self,
|
||||
node_id: NodeIdx_,
|
||||
arena: &NodeArena<K>,
|
||||
mut should_include: impl FnMut(&Node<K>) -> bool,
|
||||
) {
|
||||
let mut prev = HEAD;
|
||||
let mut cur = node_id;
|
||||
loop {
|
||||
let node = arena.node(cur);
|
||||
let Some(parent) = node.try_parent() else {
|
||||
break;
|
||||
};
|
||||
if should_include(node) {
|
||||
let cell = Self::cell_of_(cur);
|
||||
self.remove_node_(cell);
|
||||
self.add_node_after_(prev, cell);
|
||||
prev = cell;
|
||||
}
|
||||
cur = parent;
|
||||
}
|
||||
}
|
||||
|
||||
/// Like `reset_node_and_parents_mru`, stopping once `window_size` atoms
|
||||
/// are covered; excluded ancestors consume the window too.
|
||||
pub fn reset_node_and_window_ancestors_mru<K: ChildKeyType>(
|
||||
&mut self,
|
||||
node_id: NodeIdx_,
|
||||
window_size: usize,
|
||||
arena: &NodeArena<K>,
|
||||
mut should_include: impl FnMut(&Node<K>) -> bool,
|
||||
) {
|
||||
let mut prev = HEAD;
|
||||
let mut accumulated = 0;
|
||||
let mut cur = node_id;
|
||||
while accumulated < window_size {
|
||||
let node = arena.node(cur);
|
||||
let Some(parent) = node.try_parent() else {
|
||||
break;
|
||||
};
|
||||
if should_include(node) {
|
||||
let cell = Self::cell_of_(cur);
|
||||
self.remove_node_(cell);
|
||||
self.add_node_after_(prev, cell);
|
||||
prev = cell;
|
||||
}
|
||||
accumulated += node.key.atom_len();
|
||||
cur = parent;
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the node is a member (`None` is never a member).
|
||||
pub fn in_list(&self, node_id: Option<NodeIdx_>) -> bool {
|
||||
node_id.is_some_and(|id| self.in_list_(Self::cell_of_(id)))
|
||||
}
|
||||
|
||||
/// The nearest predecessor of `cell` satisfying `pred`, walking toward the
|
||||
/// head; `cell` itself is excluded.
|
||||
fn get_prev_where_(
|
||||
&self,
|
||||
cell: CellId,
|
||||
mut pred: impl FnMut(NodeIdx_) -> bool,
|
||||
) -> Option<NodeIdx_> {
|
||||
let mut cell = self.cell_(cell).prev;
|
||||
while cell != HEAD {
|
||||
let node = Self::node_of_(cell);
|
||||
if pred(node) {
|
||||
return Some(node);
|
||||
}
|
||||
cell = self.cell_(cell).prev;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// The nearest predecessor of a member satisfying `pred`; panics if
|
||||
/// `node_id` is not a member.
|
||||
pub fn get_prev_where(
|
||||
&self,
|
||||
node_id: NodeIdx_,
|
||||
pred: impl FnMut(NodeIdx_) -> bool,
|
||||
) -> Option<NodeIdx_> {
|
||||
self.get_prev_where_(Self::cell_of_(node_id), pred)
|
||||
}
|
||||
|
||||
/// The least-recent member whose lock on the list's own slot is free.
|
||||
pub fn get_lru_no_lock<K: ChildKeyType>(&self, arena: &NodeArena<K>) -> Option<NodeIdx_> {
|
||||
self.get_lru_where(|id| arena.node(id).lock_ref_(self.slot) == 0)
|
||||
}
|
||||
|
||||
/// The nearest more-recent member whose lock on the list's own slot is
|
||||
/// free, from `node_id`.
|
||||
pub fn get_prev_no_lock<K: ChildKeyType>(
|
||||
&self,
|
||||
node_id: NodeIdx_,
|
||||
arena: &NodeArena<K>,
|
||||
) -> Option<NodeIdx_> {
|
||||
self.get_prev_where(node_id, |id| arena.node(id).lock_ref_(self.slot) == 0)
|
||||
}
|
||||
|
||||
/// The least-recently-used member satisfying `pred`.
|
||||
pub fn get_lru_where(&self, pred: impl FnMut(NodeIdx_) -> bool) -> Option<NodeIdx_> {
|
||||
self.get_prev_where_(TAIL, pred)
|
||||
}
|
||||
|
||||
/// Number of member cells, excluding the sentinels.
|
||||
pub fn len(&self) -> usize {
|
||||
self.len
|
||||
}
|
||||
|
||||
/// Materialize the current members from most to least recent.
|
||||
///
|
||||
/// Inspection callers need an owned snapshot across the Python boundary;
|
||||
/// the linked-list iterator itself never escapes the Rust core.
|
||||
pub(crate) fn snapshot_node_ids(&self) -> Vec<NodeIdx_> {
|
||||
let mut node_ids = Vec::with_capacity(self.len);
|
||||
let mut cell = self.cell_(HEAD).next;
|
||||
while cell != TAIL {
|
||||
node_ids.push(Self::node_of_(cell));
|
||||
cell = self.cell_(cell).next;
|
||||
}
|
||||
node_ids
|
||||
}
|
||||
|
||||
// ==== Test-only conveniences ====
|
||||
|
||||
/// The members, MRU to LRU.
|
||||
#[cfg(test)]
|
||||
pub fn iter(&self) -> impl Iterator<Item = NodeIdx_> + '_ {
|
||||
let mut cell = self.cell_(HEAD).next;
|
||||
std::iter::from_fn(move || {
|
||||
if cell == TAIL {
|
||||
return None;
|
||||
}
|
||||
let node = Self::node_of_(cell);
|
||||
cell = self.cell_(cell).next;
|
||||
Some(node)
|
||||
})
|
||||
}
|
||||
|
||||
/// Panics if the links, membership flags, or member counter are inconsistent.
|
||||
/// Reads cells raw: it inspects possibly-inconsistent state that the gated
|
||||
/// accessors would reject.
|
||||
#[cfg(test)]
|
||||
pub fn validate(&self) {
|
||||
let mut count = 0;
|
||||
let mut prev = HEAD;
|
||||
let mut cell = self.cells[HEAD.0].next;
|
||||
while cell != TAIL {
|
||||
assert!(
|
||||
cell.0 >= OFFSET && cell.0 < self.cells.len(),
|
||||
"validate: cell {cell:?} out of bounds"
|
||||
);
|
||||
assert_eq!(
|
||||
self.cells[cell.0].prev, prev,
|
||||
"validate: broken prev link at cell {cell:?}"
|
||||
);
|
||||
assert!(
|
||||
self.cells[cell.0].in_list,
|
||||
"validate: membership mismatch at cell {cell:?}"
|
||||
);
|
||||
count += 1;
|
||||
assert!(count <= self.len, "validate: cycle detected");
|
||||
prev = cell;
|
||||
cell = self.cells[cell.0].next;
|
||||
}
|
||||
assert_eq!(self.cells[TAIL.0].prev, prev, "validate: broken tail link");
|
||||
assert_eq!(count, self.len, "validate: length mismatch");
|
||||
let flagged = self
|
||||
.cells
|
||||
.iter()
|
||||
.skip(OFFSET)
|
||||
.filter(|cell| cell.in_list)
|
||||
.count();
|
||||
assert_eq!(flagged, self.len, "validate: membership mismatch");
|
||||
}
|
||||
/// Test-only: desynchronize `len` to force integrity errors.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn bump_len_for_test(&mut self) {
|
||||
self.len += 1;
|
||||
}
|
||||
|
||||
/// Walk a LRU doubly-linked list, collect integrity errors.
|
||||
pub(crate) fn check_linked_list_(&self, label: &str, errors: &mut Vec<String>) {
|
||||
let mut visited: HashSet<usize> = HashSet::new();
|
||||
let mut prev = HEAD;
|
||||
let mut x = self.cells[HEAD.0].next;
|
||||
while x != TAIL {
|
||||
if x.0 < OFFSET {
|
||||
errors.push(format!("{label} broken chain: link points at a sentinel"));
|
||||
break;
|
||||
}
|
||||
let Some(cell) = self.cells.get(x.0) else {
|
||||
errors.push(format!("{label} broken chain: cell {} out of bounds", x.0));
|
||||
break;
|
||||
};
|
||||
if cell.prev != prev {
|
||||
errors.push(format!("{label} broken prev at node {}", Self::node_of_(x)));
|
||||
}
|
||||
if !cell.in_list {
|
||||
errors.push(format!(
|
||||
"{label} node {} in list not flagged",
|
||||
Self::node_of_(x)
|
||||
));
|
||||
}
|
||||
if !visited.insert(x.0) {
|
||||
errors.push(format!("{label} cycle at node {}", Self::node_of_(x)));
|
||||
break;
|
||||
}
|
||||
prev = x;
|
||||
x = cell.next;
|
||||
}
|
||||
// The tail backlink closes the list onto the last visited member.
|
||||
if x == TAIL && self.cells[TAIL.0].prev != prev {
|
||||
errors.push(format!("{label} broken tail backlink"));
|
||||
}
|
||||
// Every flagged member cell must be reachable from the head.
|
||||
for (idx, cell) in self.cells.iter().enumerate().skip(OFFSET) {
|
||||
if cell.in_list && !visited.contains(&idx) {
|
||||
errors.push(format!(
|
||||
"{label} node {} flagged but unreachable",
|
||||
idx - OFFSET
|
||||
));
|
||||
}
|
||||
}
|
||||
if visited.len() != self.len {
|
||||
errors.push(format!(
|
||||
"{label} list={} != len={}",
|
||||
visited.len(),
|
||||
self.len
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Eviction priority keys.
|
||||
|
||||
/// Eviction-priority key, ordered lexicographically; lower evicts first.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
|
||||
pub struct PriorityKey(pub i64, pub i64);
|
||||
|
||||
/// Ranks nodes for eviction; lower priority evicts first.
|
||||
pub trait EvictionStrategy<K: ChildKeyType> {
|
||||
/// The node's eviction priority.
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey;
|
||||
}
|
||||
|
||||
/// Least-recently-used.
|
||||
pub struct LruStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for LruStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(node.last_access_counter, 0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Least-frequently-used; LRU within a hit count.
|
||||
pub struct LfuStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for LfuStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(node.hit_count, node.last_access_counter)
|
||||
}
|
||||
}
|
||||
|
||||
/// First-in-first-out over creation order.
|
||||
pub struct FifoStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for FifoStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(node.creation_counter, 0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Most-recently-used first.
|
||||
pub struct MruStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for MruStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(-node.last_access_counter, 0)
|
||||
}
|
||||
}
|
||||
|
||||
/// First-in-last-out over creation order.
|
||||
pub struct FiloStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for FiloStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(-node.creation_counter, 0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Priority-aware: lower node priority evicts first, LRU within a priority.
|
||||
pub struct PriorityStrategy;
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for PriorityStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(node.priority, node.last_access_counter)
|
||||
}
|
||||
}
|
||||
|
||||
/// Segmented LRU: probationary nodes (hits below the threshold) evict before
|
||||
/// protected ones, LRU within a segment.
|
||||
pub struct SlruStrategy {
|
||||
pub protected_threshold: i64,
|
||||
}
|
||||
|
||||
impl<K: ChildKeyType> EvictionStrategy<K> for SlruStrategy {
|
||||
fn get_priority(&self, node: &Node<K>) -> PriorityKey {
|
||||
PriorityKey(
|
||||
(node.hit_count >= self.protected_threshold) as i64,
|
||||
node.last_access_counter,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// The strategy for an eviction-policy name.
|
||||
pub fn get_eviction_strategy<K: ChildKeyType>(policy: &str) -> Box<dyn EvictionStrategy<K> + Send> {
|
||||
match policy.to_lowercase().as_str() {
|
||||
"lru" => Box::new(LruStrategy),
|
||||
"lfu" => Box::new(LfuStrategy),
|
||||
"fifo" => Box::new(FifoStrategy),
|
||||
"mru" => Box::new(MruStrategy),
|
||||
"filo" => Box::new(FiloStrategy),
|
||||
"priority" => Box::new(PriorityStrategy),
|
||||
"slru" => Box::new(SlruStrategy {
|
||||
protected_threshold: 2,
|
||||
}),
|
||||
other => panic!(
|
||||
"Unknown eviction policy: {other}. Supported policies: \
|
||||
'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'."
|
||||
),
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
#[path = "tests/unified_lru_list.rs"]
|
||||
mod tests;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,15 @@
|
||||
#pragma once
|
||||
|
||||
#include <stdexcept>
|
||||
#include <torch/version.h>
|
||||
|
||||
#if TORCH_VERSION_MAJOR > 2 || \
|
||||
(TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 13)
|
||||
// Keep tch 0.24's removed alignment wrappers as explicit runtime errors.
|
||||
#define align_as(...) \
|
||||
alias(); \
|
||||
throw std::runtime_error("align_as is unavailable in PyTorch 2.13+")
|
||||
#define align_tensors(...) \
|
||||
autograd::variable_list{}; \
|
||||
throw std::runtime_error("align_tensors is unavailable in PyTorch 2.13+")
|
||||
#endif
|
||||
Reference in New Issue
Block a user