Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Adding a test for FD. (#2516) 

* Adding a test for FD.

* Fixing flashdecoding (empty batch doesn't work).

* Fixing the invalid popping.

* Fixing radix with block_size > 1

* Last reference.

* Use an actual hash.

* Update hash for slice.len() == 1

* Update the locks.

* Increasing docker timeout.
This commit is contained in:
Nicolas Patry 2024-09-16 17:00:54 +02:00 committed by GitHub
parent 7774655297
commit 38fcafcf96
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 3003 additions and 191 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

312
Cargo.lock generated
View File

@ -4,9 +4,9 @@ version = 3
[[package]] [[package]]
name = "addr2line" name = "addr2line"
version = "0.22.0" version = "0.24.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678" checksum = "f5fb1d8e4442bd405fdfd1dacb42792696b0cf9cb15882e5d097b742a676d375"
dependencies = [ dependencies = [
"gimli", "gimli",
] ]
@ -109,9 +109,9 @@ dependencies = [
[[package]] [[package]]
name = "anyhow" name = "anyhow"
version = "1.0.86" version = "1.0.88"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" checksum = "4e1496f8fb1fbf272686b8d37f523dab3e4a7443300055e74cdaa449f3114356"
[[package]] [[package]]
name = "arbitrary" name = "arbitrary"
@ -133,7 +133,7 @@ checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -172,18 +172,18 @@ checksum = "16e62a023e7c117e27523144c5d2459f4397fcc3cab0085af8e2224f643a0193"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
name = "async-trait" name = "async-trait"
version = "0.1.81" version = "0.1.82"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e0c28dcc82d7c8ead5cb13beb15405b57b8546e93215673ff8ca0349a028107" checksum = "a27b8a3a6e1a44fa4c8baf1f653e4172e81486d4941f2237e20dc2d0cf4ddff1"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -257,9 +257,9 @@ dependencies = [
[[package]] [[package]]
name = "aws-lc-rs" name = "aws-lc-rs"
version = "1.8.1" version = "1.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ae74d9bd0a7530e8afd1770739ad34b36838829d6ad61818f9230f683f5ad77" checksum = "2f95446d919226d587817a7d21379e6eb099b97b45110a7f272a444ca5c54070"
dependencies = [ dependencies = [
"aws-lc-sys", "aws-lc-sys",
"mirai-annotations", "mirai-annotations",
@ -269,9 +269,9 @@ dependencies = [
[[package]] [[package]]
name = "aws-lc-sys" name = "aws-lc-sys"
version = "0.20.1" version = "0.21.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f0e249228c6ad2d240c2dc94b714d711629d52bad946075d8e9b2f5391f0703" checksum = "234314bd569802ec87011d653d6815c6d7b9ffb969e9fee5b8b20ef860e8dce9"
dependencies = [ dependencies = [
"bindgen", "bindgen",
"cc", "cc",
@ -406,17 +406,17 @@ dependencies = [
[[package]] [[package]]
name = "backtrace" name = "backtrace"
version = "0.3.73" version = "0.3.74"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5cc23269a4f8976d0a4d2e7109211a419fe30e8d88d677cd60b6bc79c5732e0a" checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a"
dependencies = [ dependencies = [
"addr2line", "addr2line",
"cc",
"cfg-if", "cfg-if",
"libc", "libc",
"miniz_oxide 0.7.4", "miniz_oxide 0.8.0",
"object", "object",
"rustc-demangle", "rustc-demangle",
"windows-targets 0.52.6",
] ]
[[package]] [[package]]
@ -456,7 +456,7 @@ dependencies = [
"regex", "regex",
"rustc-hash", "rustc-hash",
"shlex", "shlex",
"syn 2.0.76", "syn 2.0.77",
"which", "which",
] ]
@ -528,9 +528,9 @@ checksum = "5ce89b21cab1437276d2650d57e971f9d548a2d9037cc231abdc0562b97498ce"
[[package]] [[package]]
name = "bytemuck" name = "bytemuck"
version = "1.17.1" version = "1.18.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773d90827bc3feecfb67fab12e24de0749aad83c74b9504ecde46237b5cd24e2" checksum = "94bbb0ad554ad961ddc5da507a12a29b14e4ae5bda06b19f575a3e6079d2e2ae"
[[package]] [[package]]
name = "byteorder" name = "byteorder"
@ -605,9 +605,9 @@ dependencies = [
[[package]] [[package]]
name = "cc" name = "cc"
version = "1.1.15" version = "1.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57b6a275aa2903740dc87da01c62040406b8812552e97129a63ea8850a17c6e6" checksum = "b62ac837cdb5cb22e10a256099b4fc502b1dfe560cb282963a974d7abd80e476"
dependencies = [ dependencies = [
"jobserver", "jobserver",
"libc", "libc",
@ -675,9 +675,9 @@ dependencies = [
[[package]] [[package]]
name = "clap" name = "clap"
version = "4.5.16" version = "4.5.17"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed6719fffa43d0d87e5fd8caeab59be1554fb028cd30edc88fc4369b17971019" checksum = "3e5a21b8495e732f1b3c364c9949b201ca7bae518c502c80256c96ad79eaf6ac"
dependencies = [ dependencies = [
"clap_builder", "clap_builder",
"clap_derive", "clap_derive",
@ -685,9 +685,9 @@ dependencies = [
[[package]] [[package]]
name = "clap_builder" name = "clap_builder"
version = "4.5.15" version = "4.5.17"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "216aec2b177652e3846684cbfe25c9964d18ec45234f0f5da5157b207ed1aab6" checksum = "8cf2dd12af7a047ad9d6da2b6b249759a22a7abc0f474c1dae1777afa4b21a73"
dependencies = [ dependencies = [
"anstream", "anstream",
"anstyle", "anstyle",
@ -704,7 +704,7 @@ dependencies = [
"heck 0.5.0", "heck 0.5.0",
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -789,9 +789,9 @@ checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
[[package]] [[package]]
name = "cpufeatures" name = "cpufeatures"
version = "0.2.13" version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51e852e6dc9a5bed1fae92dd2375037bf2b768725bf3be87811edee3249d09ad" checksum = "608697df725056feaccfa42cffdaeeec3fccc4ffc38358ecd19b243e716a78e0"
dependencies = [ dependencies = [
"libc", "libc",
] ]
@ -949,9 +949,9 @@ dependencies = [
[[package]] [[package]]
name = "cxx" name = "cxx"
version = "1.0.126" version = "1.0.128"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3c4eae4b7fc8dcb0032eb3b1beee46b38d371cdeaf2d0c64b9944f6f69ad7755" checksum = "54ccead7d199d584d139148b04b4a368d1ec7556a1d9ea2548febb1b9d49f9a4"
dependencies = [ dependencies = [
"cc", "cc",
"cxxbridge-flags", "cxxbridge-flags",
@ -961,9 +961,9 @@ dependencies = [
[[package]] [[package]]
name = "cxx-build" name = "cxx-build"
version = "1.0.126" version = "1.0.128"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c822bf7fb755d97328d6c337120b6f843678178751cba33c9da25cf522272e0" checksum = "c77953e99f01508f89f55c494bfa867171ef3a6c8cea03d26975368f2121a5c1"
dependencies = [ dependencies = [
"cc", "cc",
"codespan-reporting", "codespan-reporting",
@ -971,24 +971,24 @@ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"scratch", "scratch",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
name = "cxxbridge-flags" name = "cxxbridge-flags"
version = "1.0.126" version = "1.0.128"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "719d6197dc016c88744aff3c0d0340a01ecce12e8939fc282e7c8f583ee64bc6" checksum = "65777e06cc48f0cb0152024c77d6cf9e4bdb4408e7b48bea993d42fa0f5b02b6"
[[package]] [[package]]
name = "cxxbridge-macro" name = "cxxbridge-macro"
version = "1.0.126" version = "1.0.128"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "35de3b547387863c8f82013c4f79f1c2162edee956383e4089e1d04c18c4f16c" checksum = "98532a60dedaebc4848cb2cba5023337cc9ea3af16a5b062633fabfd9f18fb60"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1012,7 +1012,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"strsim", "strsim",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1023,7 +1023,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806"
dependencies = [ dependencies = [
"darling_core", "darling_core",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1037,33 +1037,33 @@ dependencies = [
[[package]] [[package]]
name = "derive_builder" name = "derive_builder"
version = "0.20.0" version = "0.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0350b5cb0331628a5916d6c5c0b72e97393b8b6b03b47a9284f4e7f5a405ffd7" checksum = "cd33f37ee6a119146a1781d3356a7c26028f83d779b2e04ecd45fdc75c76877b"
dependencies = [ dependencies = [
"derive_builder_macro", "derive_builder_macro",
] ]
[[package]] [[package]]
name = "derive_builder_core" name = "derive_builder_core"
version = "0.20.0" version = "0.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d48cda787f839151732d396ac69e3473923d54312c070ee21e9effcaa8ca0b1d" checksum = "7431fa049613920234f22c47fdc33e6cf3ee83067091ea4277a3f8c4587aae38"
dependencies = [ dependencies = [
"darling", "darling",
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
name = "derive_builder_macro" name = "derive_builder_macro"
version = "0.20.0" version = "0.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "206868b8242f27cecce124c19fd88157fbd0dd334df2587f36417bafbc85097b" checksum = "4abae7035bf79b9877b779505d8cf3749285b80c43941eda66604841889451dc"
dependencies = [ dependencies = [
"derive_builder_core", "derive_builder_core",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1338,7 +1338,7 @@ checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1415,9 +1415,9 @@ dependencies = [
[[package]] [[package]]
name = "gimli" name = "gimli"
version = "0.29.0" version = "0.31.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd" checksum = "32085ea23f3234fc7846555e85283ba4de91e21016dc0455a16286d87a292d64"
[[package]] [[package]]
name = "glob" name = "glob"
@ -1447,7 +1447,7 @@ dependencies = [
"futures-sink", "futures-sink",
"futures-util", "futures-util",
"http 0.2.12", "http 0.2.12",
"indexmap 2.4.0", "indexmap 2.5.0",
"slab", "slab",
"tokio", "tokio",
"tokio-util", "tokio-util",
@ -1466,7 +1466,7 @@ dependencies = [
"futures-core", "futures-core",
"futures-sink", "futures-sink",
"http 1.1.0", "http 1.1.0",
"indexmap 2.4.0", "indexmap 2.5.0",
"slab", "slab",
"tokio", "tokio",
"tokio-util", "tokio-util",
@ -1688,16 +1688,16 @@ dependencies = [
[[package]] [[package]]
name = "hyper-rustls" name = "hyper-rustls"
version = "0.27.2" version = "0.27.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5ee4be2c948921a1a5320b629c4193916ed787a7f7f293fd3f7f5a6c9de74155" checksum = "08afdbb5c31130e3034af566421053ab03787c640246a446327f550d11bcb333"
dependencies = [ dependencies = [
"futures-util", "futures-util",
"http 1.1.0", "http 1.1.0",
"hyper 1.4.1", "hyper 1.4.1",
"hyper-util", "hyper-util",
"log", "log",
"rustls 0.23.12", "rustls 0.23.13",
"rustls-native-certs", "rustls-native-certs",
"rustls-pki-types", "rustls-pki-types",
"tokio", "tokio",
@ -1732,9 +1732,9 @@ dependencies = [
[[package]] [[package]]
name = "hyper-util" name = "hyper-util"
version = "0.1.7" version = "0.1.8"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cde7055719c54e36e95e8719f95883f22072a48ede39db7fc17a4e1d5281e9b9" checksum = "da62f120a8a37763efb0cf8fdf264b884c7b8b9ac8660b900c8661030c00e6ba"
dependencies = [ dependencies = [
"bytes", "bytes",
"futures-channel", "futures-channel",
@ -1817,9 +1817,9 @@ dependencies = [
[[package]] [[package]]
name = "indexmap" name = "indexmap"
version = "2.4.0" version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93ead53efc7ea8ed3cfb0c79fc8023fbb782a5432b52830b6518941cebe6505c" checksum = "68b900aa2f7301e21c36462b170ee99994de34dff39a4a6a528e80e7376d07e5"
dependencies = [ dependencies = [
"equivalent", "equivalent",
"hashbrown 0.14.5", "hashbrown 0.14.5",
@ -1839,6 +1839,12 @@ dependencies = [
"unicode-width", "unicode-width",
] ]
[[package]]
name = "indoc"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"
[[package]] [[package]]
name = "init-tracing-opentelemetry" name = "init-tracing-opentelemetry"
version = "0.14.1" version = "0.14.1"
@ -1859,7 +1865,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b23a0c8dfe501baac4adf6ebbfa6eddf8f0c07f56b058cc1288017e32397846c" checksum = "b23a0c8dfe501baac4adf6ebbfa6eddf8f0c07f56b058cc1288017e32397846c"
dependencies = [ dependencies = [
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -1879,14 +1885,14 @@ checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
name = "ipnet" name = "ipnet"
version = "2.9.0" version = "2.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3" checksum = "187674a687eed5fe42285b40c6291f9a01517d415fad1c3cbc6a9f778af7fcd4"
[[package]] [[package]]
name = "is_terminal_polyfill" name = "is_terminal_polyfill"
@ -1979,7 +1985,7 @@ dependencies = [
"anyhow", "anyhow",
"base64 0.21.7", "base64 0.21.7",
"bytecount", "bytecount",
"clap 4.5.16", "clap 4.5.17",
"fancy-regex", "fancy-regex",
"fraction", "fraction",
"getrandom", "getrandom",
@ -2191,7 +2197,7 @@ dependencies = [
"hyper 1.4.1", "hyper 1.4.1",
"hyper-rustls", "hyper-rustls",
"hyper-util", "hyper-util",
"indexmap 2.4.0", "indexmap 2.5.0",
"ipnet", "ipnet",
"metrics", "metrics",
"metrics-util", "metrics-util",
@ -2314,7 +2320,7 @@ checksum = "a7ce64b975ed4f123575d11afd9491f2e37bbd5813fbfbc0f09ae1fbddea74e0"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -2514,7 +2520,7 @@ checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -2585,18 +2591,18 @@ checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
[[package]] [[package]]
name = "object" name = "object"
version = "0.36.3" version = "0.36.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "27b64972346851a39438c60b341ebc01bba47464ae329e55cf343eb93964efd9" checksum = "084f1a5821ac4c651660a94a7153d27ac9d8a53736203f58b31945ded098070a"
dependencies = [ dependencies = [
"memchr", "memchr",
] ]
[[package]] [[package]]
name = "once_cell" name = "once_cell"
version = "1.19.0" version = "1.20.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" checksum = "33ea5043e58958ee56f3e15a90aee535795cd7dfd319846288d93c5b57d85cbe"
[[package]] [[package]]
name = "onig" name = "onig"
@ -2649,7 +2655,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -2688,7 +2694,7 @@ checksum = "1e32339a5dc40459130b3bd269e9892439f55b33e772d2a9d402a789baaf4e8a"
dependencies = [ dependencies = [
"futures-core", "futures-core",
"futures-sink", "futures-sink",
"indexmap 2.4.0", "indexmap 2.5.0",
"js-sys", "js-sys",
"once_cell", "once_cell",
"pin-project-lite", "pin-project-lite",
@ -2912,7 +2918,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db"
dependencies = [ dependencies = [
"fixedbitset", "fixedbitset",
"indexmap 2.4.0", "indexmap 2.5.0",
] ]
[[package]] [[package]]
@ -2932,7 +2938,7 @@ checksum = "2f38a4412a78282e09a2cf38d195ea5420d15ba0602cb375210efbc877243965"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -2955,9 +2961,9 @@ checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec"
[[package]] [[package]]
name = "plotters" name = "plotters"
version = "0.3.6" version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a15b6eccb8484002195a3e44fe65a4ce8e93a625797a063735536fd59cb01cf3" checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
dependencies = [ dependencies = [
"num-traits", "num-traits",
"plotters-backend", "plotters-backend",
@ -2968,15 +2974,15 @@ dependencies = [
[[package]] [[package]]
name = "plotters-backend" name = "plotters-backend"
version = "0.3.6" version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "414cec62c6634ae900ea1c56128dfe87cf63e7caece0852ec76aba307cebadb7" checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
[[package]] [[package]]
name = "plotters-svg" name = "plotters-svg"
version = "0.3.6" version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "81b30686a7d9c3e010b84284bdd26a29f2138574f52f5eb6f794fc0ad924e705" checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
dependencies = [ dependencies = [
"plotters-backend", "plotters-backend",
] ]
@ -3022,7 +3028,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "479cf940fbbb3426c32c5d5176f62ad57549a0bb84773423ba8be9d089f5faba" checksum = "479cf940fbbb3426c32c5d5176f62ad57549a0bb84773423ba8be9d089f5faba"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -3074,7 +3080,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8021cf59c8ec9c432cfc2526ac6b8aa508ecaf29cd415f271b8406c1b851c3fd" checksum = "8021cf59c8ec9c432cfc2526ac6b8aa508ecaf29cd415f271b8406c1b851c3fd"
dependencies = [ dependencies = [
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -3114,7 +3120,7 @@ dependencies = [
"prost 0.12.6", "prost 0.12.6",
"prost-types", "prost-types",
"regex", "regex",
"syn 2.0.76", "syn 2.0.77",
"tempfile", "tempfile",
] ]
@ -3141,7 +3147,7 @@ dependencies = [
"itertools 0.12.1", "itertools 0.12.1",
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -3200,7 +3206,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"pyo3-macros-backend", "pyo3-macros-backend",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -3213,7 +3219,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"pyo3-build-config", "pyo3-build-config",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -3397,9 +3403,9 @@ dependencies = [
[[package]] [[package]]
name = "redox_syscall" name = "redox_syscall"
version = "0.5.3" version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a908a6e00f1fdd0dfd9c0eb08ce85126f6d8bbda50017e74bc4a4b7d4a926a4" checksum = "0884ad60e090bf1345b93da0a5de8923c93884cd03f40dfcfddd3b4bee661853"
dependencies = [ dependencies = [
"bitflags 2.6.0", "bitflags 2.6.0",
] ]
@ -3501,9 +3507,9 @@ dependencies = [
[[package]] [[package]]
name = "rgb" name = "rgb"
version = "0.8.48" version = "0.8.50"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f86ae463694029097b846d8f99fd5536740602ae00022c0c50c5600720b2f71" checksum = "57397d16646700483b67d2dd6511d79318f9d057fdbd21a4066aeac8b41d310a"
dependencies = [ dependencies = [
"bytemuck", "bytemuck",
] ]
@ -3558,7 +3564,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"rust-embed-utils", "rust-embed-utils",
"syn 2.0.76", "syn 2.0.77",
"walkdir", "walkdir",
] ]
@ -3586,18 +3592,18 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
[[package]] [[package]]
name = "rustc_version" name = "rustc_version"
version = "0.4.0" version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366" checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
dependencies = [ dependencies = [
"semver", "semver",
] ]
[[package]] [[package]]
name = "rustix" name = "rustix"
version = "0.38.34" version = "0.38.37"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "70dc5ec042f7a43c4a73241207cecc9873a06d45debb38b329f8541d85c2730f" checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811"
dependencies = [ dependencies = [
"bitflags 2.6.0", "bitflags 2.6.0",
"errno", "errno",
@ -3634,9 +3640,9 @@ dependencies = [
[[package]] [[package]]
name = "rustls" name = "rustls"
version = "0.23.12" version = "0.23.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c58f8c84392efc0a126acce10fa59ff7b3d2ac06ab451a33f2741989b806b044" checksum = "f2dabaac7466917e566adb06783a81ca48944c6898a1b08b9374106dd671f4c8"
dependencies = [ dependencies = [
"aws-lc-rs", "aws-lc-rs",
"log", "log",
@ -3649,9 +3655,9 @@ dependencies = [
[[package]] [[package]]
name = "rustls-native-certs" name = "rustls-native-certs"
version = "0.7.2" version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "04182dffc9091a404e0fc069ea5cd60e5b866c3adf881eff99a32d048242dffa" checksum = "fcaf18a4f2be7326cd874a5fa579fae794320a0f388d365dca7e480e55f83f8a"
dependencies = [ dependencies = [
"openssl-probe", "openssl-probe",
"rustls-pemfile 2.1.3", "rustls-pemfile 2.1.3",
@ -3687,9 +3693,9 @@ checksum = "fc0a2ce646f8655401bb81e7927b812614bd5d91dbc968696be50603510fcaf0"
[[package]] [[package]]
name = "rustls-webpki" name = "rustls-webpki"
version = "0.102.7" version = "0.102.8"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "84678086bd54edf2b415183ed7a94d0efb049f1b646a33e22a36f3794be6ae56" checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9"
dependencies = [ dependencies = [
"aws-lc-rs", "aws-lc-rs",
"ring 0.17.8", "ring 0.17.8",
@ -3720,11 +3726,11 @@ dependencies = [
[[package]] [[package]]
name = "schannel" name = "schannel"
version = "0.1.23" version = "0.1.24"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fbc91545643bcf3a0bbb6569265615222618bdf33ce4ffbbd13c4bbd4c093534" checksum = "e9aaafd5a2b6e3d657ff009d82fbd630b6bd54dd4eb06f21693925cdf80f9b8b"
dependencies = [ dependencies = [
"windows-sys 0.52.0", "windows-sys 0.59.0",
] ]
[[package]] [[package]]
@ -3783,9 +3789,9 @@ dependencies = [
[[package]] [[package]]
name = "serde" name = "serde"
version = "1.0.209" version = "1.0.210"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "99fce0ffe7310761ca6bf9faf5115afbc19688edd00171d81b1bb1b116c63e09" checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a"
dependencies = [ dependencies = [
"serde_derive", "serde_derive",
] ]
@ -3802,20 +3808,20 @@ dependencies = [
[[package]] [[package]]
name = "serde_derive" name = "serde_derive"
version = "1.0.209" version = "1.0.210"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5831b979fd7b5439637af1752d535ff49f4860c0f341d1baeb6faf0f4242170" checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
name = "serde_json" name = "serde_json"
version = "1.0.127" version = "1.0.128"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8043c06d9f82bd7271361ed64f415fe5e12a77fdb52e573e7f06a516dea329ad" checksum = "6ff5456707a1de34e7e37f2a6fd3d3f808c318259cbd01ab6377795054b483d8"
dependencies = [ dependencies = [
"itoa", "itoa",
"memchr", "memchr",
@ -4023,7 +4029,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"rustversion", "rustversion",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -4045,9 +4051,9 @@ dependencies = [
[[package]] [[package]]
name = "syn" name = "syn"
version = "2.0.76" version = "2.0.77"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "578e081a14e0cefc3279b0472138c513f37b41a08d5a3cca9b6e4e8ceb6cd525" checksum = "9f35bcdf61fd8e7be6caf75f429fdca8beb3ed76584befb503b1569faee373ed"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
@ -4173,7 +4179,7 @@ version = "2.2.1-dev0"
dependencies = [ dependencies = [
"async-stream", "async-stream",
"async-trait", "async-trait",
"clap 4.5.16", "clap 4.5.17",
"cmake", "cmake",
"cxx", "cxx",
"cxx-build", "cxx-build",
@ -4195,7 +4201,7 @@ name = "text-generation-benchmark"
version = "2.2.1-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"average", "average",
"clap 4.5.16", "clap 4.5.17",
"crossterm", "crossterm",
"float-ord", "float-ord",
"hf-hub", "hf-hub",
@ -4233,7 +4239,7 @@ dependencies = [
name = "text-generation-launcher" name = "text-generation-launcher"
version = "2.2.1-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"clap 4.5.16", "clap 4.5.17",
"ctrlc", "ctrlc",
"float_eq", "float_eq",
"hf-hub", "hf-hub",
@ -4257,7 +4263,7 @@ dependencies = [
"axum 0.7.5", "axum 0.7.5",
"axum-tracing-opentelemetry", "axum-tracing-opentelemetry",
"base64 0.22.1", "base64 0.22.1",
"clap 4.5.16", "clap 4.5.17",
"csv", "csv",
"futures", "futures",
"futures-util", "futures-util",
@ -4306,7 +4312,7 @@ dependencies = [
"axum 0.7.5", "axum 0.7.5",
"axum-tracing-opentelemetry", "axum-tracing-opentelemetry",
"base64 0.22.1", "base64 0.22.1",
"clap 4.5.16", "clap 4.5.17",
"criterion", "criterion",
"futures", "futures",
"futures-util", "futures-util",
@ -4374,7 +4380,7 @@ checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -4524,9 +4530,9 @@ dependencies = [
[[package]] [[package]]
name = "tokio" name = "tokio"
version = "1.39.3" version = "1.40.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9babc99b9923bfa4804bd74722ff02c0381021eafa4db9949217e3be8e84fff5" checksum = "e2b070231665d27ad9ec9b8df639893f46727666c6767db40317fbe920a5d998"
dependencies = [ dependencies = [
"backtrace", "backtrace",
"bytes", "bytes",
@ -4558,7 +4564,7 @@ checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -4588,16 +4594,16 @@ version = "0.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4" checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4"
dependencies = [ dependencies = [
"rustls 0.23.12", "rustls 0.23.13",
"rustls-pki-types", "rustls-pki-types",
"tokio", "tokio",
] ]
[[package]] [[package]]
name = "tokio-stream" name = "tokio-stream"
version = "0.1.15" version = "0.1.16"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "267ac89e0bec6e691e5813911606935d77c476ff49024f98abcea3e7b15e37af" checksum = "4f4e6ce100d0eb49a2734f8c0812bcd324cf357d21810932c5df6b96ef2b86f1"
dependencies = [ dependencies = [
"futures-core", "futures-core",
"pin-project-lite", "pin-project-lite",
@ -4606,9 +4612,9 @@ dependencies = [
[[package]] [[package]]
name = "tokio-util" name = "tokio-util"
version = "0.7.11" version = "0.7.12"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cf6b47b3771c49ac75ad09a6162f53ad4b8088b76ac60e8ec1455b31a189fe1" checksum = "61e7c3654c13bcd040d4a03abee2c75b1d14a37b423cf5a813ceae1cc903ec6a"
dependencies = [ dependencies = [
"bytes", "bytes",
"futures-core", "futures-core",
@ -4645,7 +4651,7 @@ version = "0.22.20"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "583c44c02ad26b0c3f3066fe629275e50627026c51ac2e595cca4c230ce1ce1d" checksum = "583c44c02ad26b0c3f3066fe629275e50627026c51ac2e595cca4c230ce1ce1d"
dependencies = [ dependencies = [
"indexmap 2.4.0", "indexmap 2.5.0",
"serde", "serde",
"serde_spanned", "serde_spanned",
"toml_datetime", "toml_datetime",
@ -4717,7 +4723,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"prost-build", "prost-build",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -4788,7 +4794,7 @@ checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -4947,9 +4953,9 @@ checksum = "08f95100a766bf4f8f28f90d77e0a5461bbdb219042e7679bebe79004fed8d75"
[[package]] [[package]]
name = "unicode-ident" name = "unicode-ident"
version = "1.0.12" version = "1.0.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe"
[[package]] [[package]]
name = "unicode-normalization" name = "unicode-normalization"
@ -4971,9 +4977,9 @@ dependencies = [
[[package]] [[package]]
name = "unicode-segmentation" name = "unicode-segmentation"
version = "1.11.0" version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
[[package]] [[package]]
name = "unicode-truncate" name = "unicode-truncate"
@ -5065,7 +5071,7 @@ version = "4.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c5afb1a60e207dca502682537fefcfd9921e71d0b83e9576060f09abc6efab23" checksum = "c5afb1a60e207dca502682537fefcfd9921e71d0b83e9576060f09abc6efab23"
dependencies = [ dependencies = [
"indexmap 2.4.0", "indexmap 2.5.0",
"serde", "serde",
"serde_json", "serde_json",
"utoipa-gen", "utoipa-gen",
@ -5081,7 +5087,7 @@ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"regex", "regex",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -5119,7 +5125,7 @@ checksum = "ee1cd046f83ea2c4e920d6ee9f7c3537ef928d75dce5d84a87c2c5d6b3999a3a"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -5220,7 +5226,7 @@ dependencies = [
"once_cell", "once_cell",
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
"wasm-bindgen-shared", "wasm-bindgen-shared",
] ]
@ -5254,7 +5260,7 @@ checksum = "afc340c74d9005395cf9dd098506f7f44e38f2b4a21c6aaacf9a105ea5e1e836"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
"wasm-bindgen-backend", "wasm-bindgen-backend",
"wasm-bindgen-shared", "wasm-bindgen-shared",
] ]
@ -5307,9 +5313,9 @@ dependencies = [
[[package]] [[package]]
name = "webpki-roots" name = "webpki-roots"
version = "0.26.3" version = "0.26.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd7c23921eeb1713a4e851530e9b9756e4fb0e89978582942612524cf09f01cd" checksum = "0bd24728e5af82c6c4ec1b66ac4844bdf8156257fccda846ec58b42cd0cdbe6a"
dependencies = [ dependencies = [
"rustls-pki-types", "rustls-pki-types",
] ]
@ -5633,7 +5639,7 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"syn 2.0.76", "syn 2.0.77",
] ]
[[package]] [[package]]
@ -5641,20 +5647,6 @@ name = "zeroize"
version = "1.8.1" version = "1.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde" checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde"
dependencies = [
"zeroize_derive",
]
[[package]]
name = "zeroize_derive"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ce36e65b0d2999d2aafac989fb249189a141aee1f53c612c1f37d72631959f69"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.76",
]
[[package]] [[package]]
name = "zip" name = "zip"

8
backends/v3/src/queue.rs
View File

@ -364,7 +364,7 @@ impl State {
// Add it back to the front // Add it back to the front
tracing::debug!("Over budget: not enough free blocks"); tracing::debug!("Over budget: not enough free blocks");
self.entries.push_front((id, entry)); self.entries.push_front((id, entry));
break; continue;
} }
Some(block_allocation) => { Some(block_allocation) => {
tracing::debug!("Allocation: {block_allocation:?}"); tracing::debug!("Allocation: {block_allocation:?}");
@ -436,6 +436,12 @@ impl State {
batch_entries.insert(id, entry); batch_entries.insert(id, entry);
} }
// Empty batch
if batch_requests.is_empty() {
tracing::debug!("Filterered out all entries");
return None;
}
// Final batch size // Final batch size
let size = batch_requests.len() as u32; let size = batch_requests.len() as u32;
next_batch_span.record("batch_size", size); next_batch_span.record("batch_size", size);

76
backends/v3/src/radix.rs
View File

@ -1,10 +1,22 @@
use crate::block_allocator::{Allocator, BlockAllocation}; use crate::block_allocator::{Allocator, BlockAllocation};
use slotmap::{DefaultKey, SlotMap}; use slotmap::{DefaultKey, SlotMap};
use std::hash::{Hash, Hasher};
use std::{ use std::{
collections::{BTreeSet, HashMap}, collections::{BTreeSet, HashMap},
sync::Arc, sync::Arc,
}; };
fn hash(slice: &[u32]) -> u64 {
assert!(!slice.is_empty());
if slice.len() == 1 {
slice[0] as u64
} else {
let mut s = std::hash::DefaultHasher::new();
slice.hash(&mut s);
s.finish()
}
}
pub struct RadixAllocator { pub struct RadixAllocator {
allocation_id: u64, allocation_id: u64,
@ -44,6 +56,10 @@ impl RadixAllocator {
// the free list if we cannot allocate enough blocks. This is only // the free list if we cannot allocate enough blocks. This is only
// temporary, the trie needs to be able to report whether it can // temporary, the trie needs to be able to report whether it can
// allocate the requested amount. Just not implemented yet. // allocate the requested amount. Just not implemented yet.
tracing::debug!(
"Free blocks {} need {n_blocks_needed}",
self.free_blocks.len()
);
self.free_blocks.extend( self.free_blocks.extend(
self.cache_blocks self.cache_blocks
.evict(n_blocks_needed - self.free_blocks.len()), .evict(n_blocks_needed - self.free_blocks.len()),
@ -94,6 +110,9 @@ impl Allocator for RadixAllocator {
match self.alloc_or_reclaim(suffix_blocks as usize) { match self.alloc_or_reclaim(suffix_blocks as usize) {
Some(suffix_blocks) => blocks.extend(suffix_blocks), Some(suffix_blocks) => blocks.extend(suffix_blocks),
None => { None => {
tracing::debug!("Cannot allocate {:?}", self.cache_blocks);
tracing::debug!("Found {prefix_len} prefix tokens need {suffix_blocks} suffix blocks for {tokens} tokens");
tracing::debug!("Block size {}", self.block_size);
self.cache_blocks self.cache_blocks
.decref(prefix_node) .decref(prefix_node)
.expect("Failed to decrement refcount"); .expect("Failed to decrement refcount");
@ -211,7 +230,6 @@ struct RadixAllocation {
pub enum TrieError { pub enum TrieError {
InvalidNodeId, InvalidNodeId,
RefCountUnderflow, RefCountUnderflow,
BlockTokenCountMismatch,
} }
pub type NodeId = DefaultKey; pub type NodeId = DefaultKey;
@ -268,16 +286,19 @@ impl RadixTrie {
fn find_(&mut self, mut node_id: NodeId, key: &[u32], blocks: &mut Vec<u32>) -> NodeId { fn find_(&mut self, mut node_id: NodeId, key: &[u32], blocks: &mut Vec<u32>) -> NodeId {
let node = &self.nodes[node_id]; let node = &self.nodes[node_id];
if let Some(&child_id) = node.children.get(&key[0]) { if key.len() >= self.block_size {
self.update_access_time(child_id); let node_key = hash(&key[..self.block_size]);
let child = self.nodes.get(child_id).expect("Invalid child identifier"); if let Some(&child_id) = node.children.get(&node_key) {
let shared_prefix_len = shared_prefix(&child.key, key, self.block_size); self.update_access_time(child_id);
assert_eq!(shared_prefix_len % self.block_size, 0); let child = self.nodes.get(child_id).expect("Invalid child identifier");
blocks.extend(&child.blocks[..shared_prefix_len / self.block_size]); let shared_prefix_len = shared_prefix(&child.key, key, self.block_size);
assert_eq!(shared_prefix_len % self.block_size, 0);
blocks.extend(&child.blocks[..shared_prefix_len / self.block_size]);
let key = &key[shared_prefix_len..]; let key = &key[shared_prefix_len..];
if !key.is_empty() { if !key.is_empty() {
node_id = self.find_(child_id, key, blocks); node_id = self.find_(child_id, key, blocks);
}
} }
} }
@ -344,9 +365,11 @@ impl RadixTrie {
// evict n_blocks and return `None` if we can't. We are now needlessly // evict n_blocks and return `None` if we can't. We are now needlessly
// evicting prefixes from the cache in such a case. // evicting prefixes from the cache in such a case.
let mut evicted = Vec::new(); let mut evicted = Vec::new();
tracing::debug!("Evicting in search of {n_blocks}");
while let Some((last_access, node_id)) = self.leaves.pop_first() { while let Some((last_access, node_id)) = self.leaves.pop_first() {
let blocks_needed = n_blocks - evicted.len(); let blocks_needed = n_blocks.saturating_sub(evicted.len());
tracing::debug!("Evicting node {node_id:?} ");
let node = self.nodes.get(node_id).expect("Leave does not exist"); let node = self.nodes.get(node_id).expect("Leave does not exist");
assert_eq!( assert_eq!(
@ -368,8 +391,11 @@ impl RadixTrie {
// the required number of blocks and leave the remaining blocks // the required number of blocks and leave the remaining blocks
// untouched. // untouched.
let node = self.nodes.get_mut(node_id).expect("Leave does not exist"); let node = self.nodes.get_mut(node_id).expect("Leave does not exist");
node.key.truncate(node.blocks.len() - blocks_needed);
evicted.extend(node.blocks.split_off(node.blocks.len() - blocks_needed)); let truncate_blocks = node.blocks.len() - blocks_needed;
let truncate_tokens = truncate_blocks * self.block_size;
node.key.truncate(truncate_tokens);
evicted.extend(node.blocks.split_off(truncate_blocks));
self.leaves.insert((last_access, node_id)); self.leaves.insert((last_access, node_id));
break; break;
} }
@ -400,11 +426,10 @@ impl RadixTrie {
// the part of the prefix that is already in the trie to detect // the part of the prefix that is already in the trie to detect
// mismatches. // mismatches.
if tokens.len() != blocks.len() * self.block_size { assert_eq!(tokens.len(), blocks.len() * self.block_size);
return Err(TrieError::BlockTokenCountMismatch);
}
if let Some(&child_id) = self.nodes[node_id].children.get(&tokens[0]) { let node_key = hash(&tokens[..self.block_size]);
if let Some(&child_id) = self.nodes[node_id].children.get(&node_key) {
self.update_access_time(child_id); self.update_access_time(child_id);
let child = self let child = self
.nodes .nodes
@ -452,14 +477,15 @@ impl RadixTrie {
.get_mut(node_id) .get_mut(node_id)
.expect("Node to-be split does not exist"); .expect("Node to-be split does not exist");
let mut parent_key = node.key.split_off(prefix_len); let mut parent_key = node.key.split_off(prefix_len);
let mut parent_blocks = node.blocks.split_off(prefix_len); let prefix_blocks = prefix_len / self.block_size;
let mut parent_blocks = node.blocks.split_off(prefix_blocks);
// Move first part of the prefix to the parent. We swap to avoid // Move first part of the prefix to the parent. We swap to avoid
// an allocation + copy for both splits of the key/blocks. // an allocation + copy for both splits of the key/blocks.
std::mem::swap(&mut node.key, &mut parent_key); std::mem::swap(&mut node.key, &mut parent_key);
std::mem::swap(&mut node.blocks, &mut parent_blocks); std::mem::swap(&mut node.blocks, &mut parent_blocks);
let node_key = node.key[0]; let node_key = hash(&node.key[..self.block_size]);
let grandparent_id = node.parent.expect("Node does not have a parent"); let grandparent_id = node.parent.expect("Node does not have a parent");
let parent_id = self.add_node(grandparent_id, parent_key, parent_blocks); let parent_id = self.add_node(grandparent_id, parent_key, parent_blocks);
@ -484,7 +510,7 @@ impl RadixTrie {
) -> NodeId { ) -> NodeId {
let key = key.into(); let key = key.into();
let blocks = blocks.into(); let blocks = blocks.into();
let first = key[0]; let first = hash(&key[..self.block_size]);
let child = TrieNode::new(key, blocks, self.time, Some(parent_id)); let child = TrieNode::new(key, blocks, self.time, Some(parent_id));
let child_id = self.nodes.insert(child); let child_id = self.nodes.insert(child);
@ -496,10 +522,10 @@ impl RadixTrie {
} }
/// Add a node to the parent. /// Add a node to the parent.
fn add_node_to_parent(&mut self, parent_id: NodeId, first: u32, child_id: NodeId) { fn add_node_to_parent(&mut self, parent_id: NodeId, hash: u64, child_id: NodeId) {
// Unwrap here, passing in an unknown id is a programming error. // Unwrap here, passing in an unknown id is a programming error.
let parent = self.nodes.get_mut(parent_id).expect("Unknown parent node"); let parent = self.nodes.get_mut(parent_id).expect("Unknown parent node");
if parent.children.insert(first, child_id).is_none() { if parent.children.insert(hash, child_id).is_none() {
// Only increase reference count if child does not replace another child. // Only increase reference count if child does not replace another child.
self.incref(parent_id) self.incref(parent_id)
.expect("Failed to increase parent refcount"); .expect("Failed to increase parent refcount");
@ -517,7 +543,9 @@ impl RadixTrie {
); );
let parent_id = node.parent.expect("Attempted to remove root node"); let parent_id = node.parent.expect("Attempted to remove root node");
let parent = self.nodes.get_mut(parent_id).expect("Unknown parent node"); let parent = self.nodes.get_mut(parent_id).expect("Unknown parent node");
parent.children.remove(&node.key[0]);
let node_key = hash(&node.key[..self.block_size]);
parent.children.remove(&node_key);
self.decref(parent_id) self.decref(parent_id)
.expect("Failed to decrease parent refcount"); .expect("Failed to decrease parent refcount");
node node
@ -571,7 +599,7 @@ impl RadixTrie {
#[derive(Debug)] #[derive(Debug)]
struct TrieNode { struct TrieNode {
blocks: Vec<u32>, blocks: Vec<u32>,
children: HashMap<u32, NodeId>, children: HashMap<u64, NodeId>,
key: Vec<u32>, key: Vec<u32>,
last_accessed: u64, last_accessed: u64,
parent: Option<NodeId>, parent: Option<NodeId>,

12
flake.lock
View File

@ -853,11 +853,11 @@
] ]
}, },
"locked": { "locked": {
"lastModified": 1726021481, "lastModified": 1726280639,
"narHash": "sha256-4J4E+Fh+77XIYnq2RVtg+ENWXpu6t74P0jKN/f2RQmI=", "narHash": "sha256-YfLRPlFZWrT2oRLNAoqf7G3+NnUTDdlIJk6tmBU7kXM=",
"owner": "oxalica", "owner": "oxalica",
"repo": "rust-overlay", "repo": "rust-overlay",
"rev": "1c2c120246c51a644c20ba2a36a33d3bd4860d70", "rev": "e9f8641c92f26fd1e076e705edb12147c384171d",
"type": "github" "type": "github"
}, },
"original": { "original": {
@ -978,11 +978,11 @@
"nixpkgs": "nixpkgs_6" "nixpkgs": "nixpkgs_6"
}, },
"locked": { "locked": {
"lastModified": 1725950569, "lastModified": 1726229792,
"narHash": "sha256-nJHA1SvIQbXySpL2ueNbzQOhnkQASa5tOLz/kdW0PWA=", "narHash": "sha256-9xsLmjc9nr7a4PTddKv2DOi82ompTtJNyjO6R67y5tE=",
"owner": "danieldk", "owner": "danieldk",
"repo": "tgi-nix", "repo": "tgi-nix",
"rev": "d40f3c22e9bcc5e16c94d4605cf6a7d74dd07f46", "rev": "1a902f4818e94c3f8d95f6000db17bc3fadd0ce7",
"type": "github" "type": "github"
}, },
"original": { "original": {

7
integration-tests/conftest.py
View File

@ -342,6 +342,7 @@ def launcher(event_loop):
max_total_tokens: Optional[int] = None, max_total_tokens: Optional[int] = None,
lora_adapters: Optional[List[str]] = None, lora_adapters: Optional[List[str]] = None,
cuda_graphs: Optional[List[int]] = None, cuda_graphs: Optional[List[int]] = None,
attention: Optional[str] = None,
): ):
port = random.randint(8000, 10_000) port = random.randint(8000, 10_000)
master_port = random.randint(10_000, 20_000) master_port = random.randint(10_000, 20_000)
@ -401,6 +402,8 @@ def launcher(event_loop):
if not use_flash_attention: if not use_flash_attention:
env["USE_FLASH_ATTENTION"] = "false" env["USE_FLASH_ATTENTION"] = "false"
if attention is not None:
env["ATTENTION"] = attention
with tempfile.TemporaryFile("w+") as tmp: with tempfile.TemporaryFile("w+") as tmp:
# We'll output stdout/stderr to a temporary file. Using a pipe # We'll output stdout/stderr to a temporary file. Using a pipe
@ -437,6 +440,7 @@ def launcher(event_loop):
max_total_tokens: Optional[int] = None, max_total_tokens: Optional[int] = None,
lora_adapters: Optional[List[str]] = None, lora_adapters: Optional[List[str]] = None,
cuda_graphs: Optional[List[int]] = None, cuda_graphs: Optional[List[int]] = None,
attention: Optional[str] = None,
): ):
port = random.randint(8000, 10_000) port = random.randint(8000, 10_000)
@ -491,6 +495,8 @@ def launcher(event_loop):
} }
if not use_flash_attention: if not use_flash_attention:
env["USE_FLASH_ATTENTION"] = "false" env["USE_FLASH_ATTENTION"] = "false"
if attention is not None:
env["ATTENTION"] = attention
if HF_TOKEN is not None: if HF_TOKEN is not None:
env["HF_TOKEN"] = HF_TOKEN env["HF_TOKEN"] = HF_TOKEN
@ -522,6 +528,7 @@ def launcher(event_loop):
devices=devices, devices=devices,
volumes=volumes, volumes=volumes,
ports={"80/tcp": port}, ports={"80/tcp": port},
healthcheck={"timeout": int(10 * 1e9)},
shm_size="1G", shm_size="1G",
) )

2550
integration-tests/models/__snapshots__/test_flash_llama_prefix_flashdecoding/test_flash_llama_flashdecoding.json Normal file
View File

File diff suppressed because it is too large Load Diff

229
integration-tests/models/test_flash_llama_prefix_flashdecoding.py Normal file
View File

File diff suppressed because one or more lines are too long