feat: positional encoded sequence hashes (#4000)

Signed-off-by: Ryan Olson <rolson@nvidia.com>

Cargo.lock

@@ -1788,6 +1788,20 @@ dependencies = [
  "parking_lot_core",
 ]
 
+[[package]]
+name = "dashmap"
+version = "6.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf"
+dependencies = [
+ "cfg-if 1.0.3",
+ "crossbeam-utils",
+ "hashbrown 0.14.5",
+ "lock_api",
+ "once_cell",
+ "parking_lot_core",
+]
+
 [[package]]
 name = "data-encoding"
 version = "2.9.0"
@@ -2166,7 +2180,7 @@ dependencies = [
  "chrono",
  "criterion 0.3.6",
  "cudarc 0.17.3",
- "dashmap",
+ "dashmap 5.5.3",
  "derive-getters",
  "derive_builder",
  "dialoguer",
@@ -2351,8 +2365,10 @@ name = "dynamo-tokens"
 version = "0.6.1"
 dependencies = [
  "bytemuck",
+ "dashmap 6.1.0",
  "derive-getters",
- "rayon",
+ "serde",
+ "thiserror 2.0.16",
  "xxhash-rust",
 ]
 
@@ -10322,7 +10338,7 @@ dependencies = [
  "asynchronous-codec",
  "bytes",
  "crossbeam-queue",
- "dashmap",
+ "dashmap 5.5.3",
  "futures-channel",
  "futures-io",
  "futures-task",

Cargo.toml

@@ -64,6 +64,7 @@ chrono = { version = "0.4", default-features = false, features = [
     "serde",
 ] }
 cudarc = { version = "0.17.1", features = ["cuda-12020"] }
+dashmap = { version = "6.1" }
 derive_builder = { version = "0.20" }
 derive-getters = { version = "0.5" }
 either = { version = "1.13", features = ["serde"] }
@@ -120,7 +121,7 @@ insta.opt-level = 3
 
 [profile.dev]
 # release level optimizations otherwise everything feels slow
-opt-level = 3
+# opt-level = 3
 
 [profile.release]
 # These make the build much slower but shrink the binary, and could help performance

lib/tokens/Cargo.toml

@@ -1,32 +1,21 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+
 
 [package]
 name = "dynamo-tokens"
+description = "Token management tools"
 version.workspace = true
 edition.workspace = true
-description.workspace = true
 authors.workspace = true
 license.workspace = true
-homepage.workspace = true
 repository.workspace = true
-keywords.workspace = true
 
 [dependencies]
+dashmap = { workspace = true }
 derive-getters = { workspace = true }
+serde = { workspace = true }
+thiserror = { workspace = true }
 xxhash-rust = { workspace = true }
 
 bytemuck = "1.22"
-rayon = "1"

lib/tokens/src/radix.rs

@@ -0,0 +1,59 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use dashmap::DashMap;
+
+use crate::PositionalSequenceHash;
+
+/// Positionally sparse radix tree for efficient indexing of [PositionalSequenceHashes][`crate::PositionalSequenceHash`].
+#[derive(Clone)]
+pub struct PositionalRadixTree<T> {
+    map: DashMap<u64, DashMap<PositionalSequenceHash, T>>,
+}
+
+impl<T> PositionalRadixTree<T> {
+    /// Creates a new empty [`PositionalRadixTree`].
+    pub fn new() -> Self {
+        Self {
+            map: DashMap::new(),
+        }
+    }
+
+    /// Provides the entry for the [`PositionalSequenceHash`] at the given position.
+    pub fn prefix(
+        &self,
+        seq_hash: &PositionalSequenceHash,
+    ) -> dashmap::mapref::one::RefMut<'_, u64, DashMap<PositionalSequenceHash, T>> {
+        let position = seq_hash.position();
+        self.map.entry(position).or_default()
+    }
+
+    /// Provides the sub-map for all [`PositionalSequenceHash`] entries at the given position.
+    pub fn position(
+        &self,
+        position: u64,
+    ) -> Option<dashmap::mapref::one::RefMut<'_, u64, DashMap<PositionalSequenceHash, T>>> {
+        self.map.get_mut(&position)
+    }
+
+    /// Returns the number of entries [`PositionalSequenceHashes`][`crate::PositionalSequenceHash`] in the [`PositionalRadixTree`].
+    pub fn len(&self) -> usize {
+        if self.map.is_empty() {
+            return 0;
+        }
+        self.map.iter().map(|level| level.len()).sum()
+    }
+
+    /// Returns true if the [`PositionalRadixTree`] is empty of [`PositionalSequenceHashes`][`crate::PositionalSequenceHash`]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
+impl<T> Default for PositionalRadixTree<T> {
+    fn default() -> Self {
+        Self {
+            map: DashMap::new(),
+        }
+    }
+}