diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,5 @@
+# Revision history for CudaHashMap
+
+## 0.1.0.0 -- 2026-11-10
+
+* First version. Released on an unsuspecting world.
diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,29 @@
+Copyright (c) 2026, Susumu Katayama
+
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+    * Redistributions of source code must retain the above copyright
+      notice, this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above
+      copyright notice, this list of conditions and the following
+      disclaimer in the documentation and/or other materials provided
+      with the distribution.
+
+    * Neither the name of the copyright holder nor the names of its
+      contributors may be used to endorse or promote products derived
+      from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/README.md b/README.md
new file mode 100644
--- /dev/null
+++ b/README.md
@@ -0,0 +1,190 @@
+# hasktorch-CudaHashMap
+
+`hasktorch-CudaHashMap` provides a [HaskTorch](https://github.com/hasktorch/hasktorch) interface to the GPU-parallel hash map implementations (`cuco::dynamic_map` and `cuco::static_map`) in NVIDIA [cuCollections (`cuco`)](https://github.com/NVIDIA/cuCollections).
+
+It supports high-throughput batch insertions and lookups directly on GPU-resident `Tensor`s without round-tripping through host (CPU) memory.
+
+## Features
+
+- **Two-Stage Lifecycle (`dynamic_map` $\rightarrow$ `static_map`)**:
+  - Starts as a dynamically resizing GPU hash map (`cuco::dynamic_map`) for workloads where the number of unique keys is not known in advance (e.g., training or incremental vocabulary/graph construction).
+  - Supports one-time compaction into a single flat `cuco::static_map` via `finalizeCudaHashMap` for maximum lookup throughput during inference or read-heavy phases.
+- **Typed Key-Value Batch Operations (`insertBatch` / `lookupBatch`)**:
+  - Supports `Int32`, `Int64`, `Float`, and `Double` `Tensor`s for both keys and values.
+  - Uses overwrite (`insert_or_assign`) semantics when duplicate keys are inserted across or within batches.
+- **Key-to-Index Mapping for `indexPut` / `indexSelect` (`insertKeys` / `lookupIndices`)**:
+  - Automatically assigns 0-based sequential `Int64` indices to unique keys on the GPU.
+  - Produces index `Tensor`s that can be passed directly to HaskTorch's `indexSelect` and `indexPut` (or via the convenience wrappers `hashMapIndexPut` and `hashMapIndexSelect`).
+- **CUDA Device & Stream Awareness**:
+  - Automatically moves input tensors to the target CUDA device and executes kernels on LibTorch's current CUDA stream.
+
+---
+
+## System & Toolchain Requirements
+
+Building the C++/CUDA wrapper (`cbits/cuco_wrapper.cu`) and linking it with GHC and LibTorch is sensitive to the versions of **`cmake`**, **`nvcc`**, **`gcc`**, and **`g++`**.
+
+> **Important OS & Compiler Note:**
+> **Ubuntu 26.04 LTS** (or newer) is recommended. **Ubuntu 24.04 LTS is not sufficient** out of the box because its standard package repositories only provide `gcc`/`g++` up to version 14, which is not enough for building this project's toolchain and dependencies. Ubuntu 26.04 provides **GCC/G++ 15** alongside newer CMake and CUDA packages.
+
+### Verified Environment (Ubuntu 26.04.1 LTS)
+
+| Component | Verified Version | Notes |
+| :--- | :--- | :--- |
+| **OS** | Ubuntu 26.04.1 LTS | Ubuntu 24.04 (max GCC 14) is insufficient. |
+| **`gcc` / `g++`** | **15.2.0** (`>= 15`) | Used both for host C++ compilation and as GHC's linker driver (`-pgml g++`). |
+| **`cmake`** | **4.4.5** | Minimum `3.23.1` required by `CMakeLists.txt`; a recent CMake is needed for CUDA architecture handling and CPM (`cuCollections` / CCCL). |
+| **`nvcc` (CUDA)** | **12.4** (or newer, e.g. `12.8`) | Must support C++17 (`-std=c++17`) and extended device lambdas (`--expt-extended-lambda`). CUDA libraries (`libcuda`, `libcudart`, `libnvrtc`) are required at link time. |
+| **GPU Architecture** | **Volta (`sm_70`) or newer** | `cuCollections` requires `sm_70+` (`70;75;80;86;89;90`) for scoped atomic operations. |
+| **GHC / Cabal** | GHC 9.10.3, Cabal `>= 3.0` | Configured via `Setup.hs` (`build-type: Custom`). |
+| **LibTorch** | **2.9.1** (CUDA build) | Installed via `libtorch-ffi` with the `+cuda` flag enabled (configured in `cabal.project`). |
+
+---
+
+## Environment Setup & Building
+
+### 1. LibTorch & `LD_LIBRARY_PATH` Configuration
+
+`cabal.project` enables the `+cuda` flag on `libtorch-ffi`:
+
+```cabal
+packages: .
+
+package libtorch-ffi
+  flags: +cuda
+```
+
+Both `cbits/CMakeLists.txt` and `setenv` resolve the LibTorch path using the following environment variables (shown with their defaults):
+
+- `LIBTORCH_HOME`: defaults to `$XDG_CACHE_HOME/libtorch` or `$HOME/.cache/libtorch`
+- `LIBTORCH_VERSION`: defaults to `2.9.1`
+- `LIBTORCH_CUDA_VERSION`: defaults to `cu128`
+
+Before running executables, tests, or benchmarks, source `setenv` to add LibTorch's `lib` directory to `LD_LIBRARY_PATH`:
+
+```bash
+source ./setenv
+```
+
+### 2. Building with Cabal
+
+The custom `Setup.hs` hook automatically runs CMake during `cabal configure` / `cabal build` to fetch `NVIDIA/cuCollections` (via CPM) and compile `libcuco_wrapper.a` into `dist-newstyle/cbits-build`:
+
+```bash
+cabal build
+```
+
+*(Optional)* If you prefer to configure and build the C++/CUDA wrapper manually (e.g., to pass custom CMake variables), create `cbits/build` before running `cabal build`. When `cbits/build/CMakeCache.txt` exists, `Setup.hs` uses `cbits/build` instead of `dist-newstyle/cbits-build`:
+
+```bash
+cmake -S cbits -B cbits/build
+cmake --build cbits/build -j4
+cabal build
+```
+
+Note that `cabal clean` removes both `dist-newstyle` and `cbits/build`.
+
+### 3. Running the Demo, Tests, and Benchmarks
+
+```bash
+# Run the example executable (app/Main.hs)
+cabal run CudaHashMap
+
+# Run the test suite (test/Main.hs)
+cabal test
+
+# Run the benchmark comparing Torch.CudaHashMap vs Data.HashMap.Strict (bench/Main.hs)
+cabal bench
+```
+
+---
+
+## Usage Examples
+
+### Pattern A: Key-Value Batch Insert & Lookup
+
+```haskell
+import Data.Int (Int64)
+import Torch.Tensor (asTensor)
+import Torch.CudaHashMap
+
+main :: IO ()
+main = do
+  -- 1. Create a dynamic GPU hash map
+  cudaMap <- newCudaHashMap
+
+  let keysTensor      = asTensor ([10, 20, 30] :: [Int64])
+      valuesTensor    = asTensor ([1.5, 2.5, 3.5] :: [Float])
+      queryKeysTensor = asTensor ([20, 10, 99] :: [Int64])
+
+  -- 2. Insert key-value batches (overwrites existing keys with new values)
+  insertBatch cudaMap keysTensor valuesTensor
+
+  -- 3. (Optional) Compact into a static_map once insertions are finished
+  finalizeCudaHashMap cudaMap
+
+  -- 4. Batch lookup (unregistered keys yield -1 for Int32/Int64, or NaN for Float/Double)
+  resultTensor <- lookupBatch cudaMap queryKeysTensor
+  print resultTensor
+  -- Output: Tensor Float [3] [ 2.5000,  1.5000,  NaN]
+```
+
+### Pattern B: Key-to-Index Mapping (`indexPut` / `indexSelect`)
+
+When storing multi-dimensional feature rows or embedding vectors in a separate `Tensor` table, `CudaHashMap` can maintain a mapping from arbitrary keys to 0-based sequential `Int64` row indices:
+
+```haskell
+import Data.Int (Int64)
+import Torch.Tensor (asTensor)
+import Torch.TensorFactories (zeros')
+import Torch.CudaHashMap
+
+main :: IO ()
+main = do
+  idxMap <- newCudaHashMap
+  let table0     = zeros' [1000]
+      keysTensor = asTensor ([500, 600, 500, 700] :: [Int64])
+      valsTensor = asTensor ([5.0, 6.0, 5.5, 7.0] :: [Float])
+
+  -- Assigns 0-based Int64 indices to unique keys and writes valsTensor into table0
+  table1 <- hashMapIndexPut idxMap False keysTensor valsTensor table0
+
+  finalizeCudaHashMap idxMap
+
+  -- Look up indices for registered keys and select from table1 along dimension 0
+  selected <- hashMapIndexSelect idxMap 0 (asTensor ([700, 500, 600] :: [Int64])) table1
+  print selected
+```
+
+You can also obtain the raw `Int64` index tensors directly via `insertKeys` and `lookupIndices` (which returns `-1` for missing keys).
+
+---
+
+## API Summary (`Torch.CudaHashMap`)
+
+| Function | Type | Description |
+| :--- | :--- | :--- |
+| `newCudaHashMap` | `IO CudaHashMap` | Create a new GPU hash map with default initial capacity (`100,000`). Device is inferred from the first tensor (or `cuda:0` for CPU tensors). |
+| `newCudaHashMapWithCapacity` | `Int -> IO CudaHashMap` | Create a new GPU hash map with a specified minimum capacity. |
+| `newCudaHashMapOnDevice` | `Device -> Int -> IO CudaHashMap` | Create a new GPU hash map pinned to a specific CUDA `Device` (e.g. `Device CUDA 1`). |
+| `hashMapDevice` | `CudaHashMap -> IO Device` | Query the CUDA `Device` associated with the hash map. |
+| `sizeCudaHashMap` | `CudaHashMap -> IO Int` | Return the number of unique keys currently stored in the map. |
+| `finalizeCudaHashMap` | `CudaHashMap -> IO ()` | Convert the internal `cuco::dynamic_map` into a single `cuco::static_map` for faster lookups. |
+| `insertBatch` | `CudaHashMap -> Tensor -> Tensor -> IO ()` | Insert or overwrite (`insert_or_assign`) key-value pairs (`Int32`, `Int64`, `Float`, `Double`). |
+| `lookupBatch` | `CudaHashMap -> Tensor -> IO Tensor` | Look up values for a tensor of keys. Missing keys yield `-1` (`Int32`/`Int64`) or `NaN` (`Float`/`Double`). |
+| `insertKeys` | `CudaHashMap -> Tensor -> IO Tensor` | Assign 0-based sequential `Int64` indices to unseen keys and return the `Int64` index tensor for all input keys. |
+| `lookupIndices` | `CudaHashMap -> Tensor -> IO Tensor` | Look up 0-based `Int64` indices for input keys without inserting (`-1` for missing keys). |
+| `hashMapIndexPut` | `CudaHashMap -> Bool -> Tensor -> Tensor -> Tensor -> IO Tensor` | Map keys via `insertKeys` and update a table tensor using `indexPut`. |
+| `hashMapIndexSelect` | `CudaHashMap -> Int -> Tensor -> Tensor -> IO Tensor` | Map keys via `lookupIndices` and gather from a table tensor using `indexSelect`. |
+
+### Sentinel Value Conventions
+
+`cuCollections` reserves sentinel values in open-addressing hash tables:
+- **Empty key sentinel**: `std::numeric_limits<K>::min()` for `Int32`/`Int64`, and quiet `NaN` for `Float`/`Double`. Avoid inserting these reserved values as keys.
+- **Missing value sentinel (on lookup)**: `-1` for `Int32`/`Int64` values/indices, and quiet `NaN` for `Float`/`Double` values.
+
+---
+
+## License
+
+BSD-3-Clause (see [LICENSE](LICENSE)).
diff --git a/Setup.hs b/Setup.hs
new file mode 100644
--- /dev/null
+++ b/Setup.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE CPP #-}
+
+import Control.Monad (filterM, unless, when)
+import Distribution.Simple
+import Distribution.Simple.Setup (ConfigFlags (..))
+import System.Directory
+    ( doesDirectoryExist
+    , doesFileExist
+    , makeAbsolute
+    , removePathForcibly
+    )
+import System.Process (callProcess)
+
+#if MIN_VERSION_Cabal(3,14,0)
+import Distribution.Utils.Path (makeSymbolicPath)
+#else
+makeSymbolicPath :: a -> a
+makeSymbolicPath = id
+#endif
+
+main :: IO ()
+main = defaultMainWithHooks simpleUserHooks
+    { confHook = \(gpd, hbi) flags -> do
+        -- Use cbits/build if it was created manually;
+        -- otherwise use dist-newstyle/cbits-build, which is removed automatically by 'cabal clean'.
+        manualBuildExists <- doesFileExist "cbits/build/CMakeCache.txt"
+        let buildDir = if manualBuildExists
+                         then "cbits/build"
+                         else "dist-newstyle/cbits-build"
+
+        cmakeCacheExists <- doesFileExist (buildDir ++ "/CMakeCache.txt")
+        unless cmakeCacheExists $
+            callProcess "cmake" ["-S", "cbits", "-B", buildDir]
+        callProcess "cmake" ["--build", buildDir, "-j4"]
+
+        absBuildDir <- makeAbsolute buildDir
+        cudaDirs <- filterM doesDirectoryExist
+            [ "/usr/local/cuda/lib64"
+            , "/usr/lib/x86_64-linux-gnu"
+            ]
+        let updatedFlags = flags
+                { configExtraLibDirs =
+                    map makeSymbolicPath (absBuildDir : cudaDirs)
+                    ++ configExtraLibDirs flags
+                }
+        confHook simpleUserHooks (gpd, hbi) updatedFlags
+
+    , cleanHook = \pd z uh cf -> do
+        exists <- doesDirectoryExist "cbits/build"
+        when exists $ removePathForcibly "cbits/build"
+        cleanHook simpleUserHooks pd z uh cf
+    }
diff --git a/app/Main.hs b/app/Main.hs
new file mode 100644
--- /dev/null
+++ b/app/Main.hs
@@ -0,0 +1,53 @@
+module Main (main) where
+
+import Control.Monad (forM_)
+import Data.Int (Int64)
+import Torch.Functional (indexPut)
+import Torch.Tensor (Tensor, asTensor, device, indexSelect, toDevice)
+import Torch.TensorFactories (zeros')
+import Torch.CudaHashMap
+
+main :: IO ()
+main = do
+  -- === Pattern A: Key-Value Hash Map (insertBatch / lookupBatch) ===
+  -- 1. Initialize the hash map (internally a dynamic_map)
+  --    To use a specific GPU, you can also call: newCudaHashMapOnDevice (Device CUDA 1) 100000
+  cudaMap <- newCudaHashMap
+
+  -- 2. Training loop (frequent insertions when total size is unpredictable)
+  forM_ ([1..100] :: [Int]) $ \_epoch -> do
+    -- Overwrites with the new value if the key already exists (insert_or_assign)
+    insertBatch cudaMap keysTensor valuesTensor
+
+  -- 3. Once insertions have stabilized, run the one-time optimization
+  finalizeCudaHashMap cudaMap  -- Converts the internal representation into an ultra-fast single static_map
+
+  -- 4. Inference and subsequent fixed lookups (runs at very high speed)
+  resultTensor <- lookupBatch cudaMap queryKeysTensor
+  print resultTensor
+
+  -- === Pattern B: Obtain sequential indices from keys for indexPut / indexSelect ===
+  idxMap <- newCudaHashMap
+  let table0 = zeros' [1000]
+  -- Automatically assign and retrieve 0-based Int64 indices for the keys
+  idx <- insertKeys idxMap keysTensor
+  -- Move valuesTensor and table0 to the same CUDA device as idx and write via indexPut
+  -- (alternatively, you can use: hashMapIndexPut idxMap False keysTensor valuesTensor table0)
+  let dev    = device idx
+      table1 = indexPut False [idx] (toDevice dev valuesTensor) (toDevice dev table0)
+  finalizeCudaHashMap idxMap
+
+  -- During inference, look up indices with lookupIndices (unregistered key 99 becomes -1)
+  qIdx <- lookupIndices idxMap queryKeysTensor
+  print qIdx
+
+  -- indexSelect for registered keys (or: hashMapIndexSelect idxMap 0 validQueryKeysTensor table1)
+  validQIdx <- lookupIndices idxMap validQueryKeysTensor
+  let selected = indexSelect 0 validQIdx table1
+  print selected
+
+keysTensor, valuesTensor, queryKeysTensor, validQueryKeysTensor :: Tensor
+keysTensor           = asTensor ([10, 20, 30] :: [Int64])
+valuesTensor         = asTensor ([1.5, 2.5, 3.5] :: [Float])
+queryKeysTensor      = asTensor ([20, 10, 99] :: [Int64])
+validQueryKeysTensor = asTensor ([20, 10, 30] :: [Int64])
diff --git a/bench/Main.hs b/bench/Main.hs
new file mode 100644
--- /dev/null
+++ b/bench/Main.hs
@@ -0,0 +1,199 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Main (main) where
+
+import Control.DeepSeq (force)
+import Control.Exception (evaluate)
+import Control.Monad (forM_, replicateM_)
+import Data.HashMap.Strict (HashMap)
+import qualified Data.HashMap.Strict as HM
+import Data.Int (Int64)
+import Data.List (foldl')
+import Foreign.C.Types (CInt(..))
+import System.Clock (Clock(Monotonic), getTime, toNanoSecs)
+import Text.Printf (printf)
+
+import Torch.Device (Device(..), DeviceType(..))
+import Torch.Tensor (Tensor, asTensor, toDevice)
+import Torch.CudaHashMap
+  ( finalizeCudaHashMap
+  , insertBatch
+  , insertKeys
+  , lookupBatch
+  , lookupIndices
+  , newCudaHashMapWithCapacity
+  )
+
+foreign import ccall "cudaDeviceSynchronize"
+  c_cudaDeviceSynchronize :: IO CInt
+
+syncCuda :: IO ()
+syncCuda = do
+  _ <- c_cudaDeviceSynchronize
+  return ()
+
+-- | Measure wall-clock execution time in milliseconds (averaged over @iters@ runs).
+timeMs :: forall a. Int -> IO a -> IO (Double, a)
+timeMs iters action = do
+  syncCuda
+  t0 <- getTime Monotonic
+  let loop :: Int -> IO a
+      loop k
+        | k <= 1    = action
+        | otherwise = action >> loop (k - 1)
+  !res <- loop (max 1 iters)
+  syncCuda
+  t1 <- getTime Monotonic
+  let ns = fromIntegral (toNanoSecs t1 - toNanoSecs t0) :: Double
+      ms = (ns / 1.0e6) / fromIntegral (max 1 iters)
+  return (ms, res)
+
+-- | Deterministic pseudo-random Int64 key generator.
+genKeys :: Int -> Int64 -> [Int64]
+genKeys n seed = take n $ drop 1 $ iterate step seed
+  where
+    step :: Int64 -> Int64
+    step x = (x * 6364136223846793005 + 1442695040888963407) `mod` 1000000007
+
+-- | CPU reference implementation for 'insertKeys' (assigns 0-based sequential Int64 indices).
+cpuInsertKeys :: [Int64] -> (HashMap Int64 Int64, [Int64])
+cpuInsertKeys ks =
+  let step :: (HashMap Int64 Int64, Int64, [Int64]) -> Int64 -> (HashMap Int64 Int64, Int64, [Int64])
+      step (!hm, !nextIdx, acc) k =
+        case HM.lookup k hm of
+          Just idx -> (hm, nextIdx, idx : acc)
+          Nothing  ->
+            let !hm'      = HM.insert k nextIdx hm
+                !nextIdx' = nextIdx + 1
+            in (hm', nextIdx', nextIdx : acc)
+      (!finalHm, _, revIndices) = foldl' step (HM.empty, 0, []) ks
+  in (finalHm, reverse revIndices)
+
+printRow :: String -> Int -> Double -> Double -> IO ()
+printRow label n baseMs ms = do
+  let mops    = (fromIntegral n / 1.0e6) / (ms / 1.0e3)
+      speedup = baseMs / ms
+  putStrLn $ printf "  %-42s %10.3f ms  (%8.2f Mops/s, %6.2fx)" label ms mops speedup
+
+runBenchmarkForSize :: Int -> IO ()
+runBenchmarkForSize n = do
+  putStrLn $ replicate 86 '='
+  putStrLn $ printf "Benchmark size: N = %d elements (Key: Int64, Value: Float / Index: Int64)" n
+  putStrLn $ replicate 86 '='
+
+  let iters   = if n <= 100000 then 5 else 3
+      cudaDev = Device CUDA 0
+      cpuDev  = Device CPU 0
+
+  -- Prepare input data on CPU and pre-materialize tensors
+  !keysList  <- evaluate $ force $ genKeys n 42
+  !valsList  <- evaluate $ force [fromIntegral (i `mod` 1000) * 0.5 :: Float | i <- [1 .. n]]
+  -- Query keys: 50% hit, 50% miss
+  let half = n `div` 2
+  !queryList <- evaluate $ force (take half keysList ++ genKeys (n - half) 99999)
+  !kvPairs   <- evaluate $ force (zip keysList valsList)
+
+  !cpuKeysTensor  <- evaluate (asTensor keysList :: Tensor)
+  !cpuValsTensor  <- evaluate (asTensor valsList :: Tensor)
+  !cpuQueryTensor <- evaluate (asTensor queryList :: Tensor)
+
+  !gpuKeysTensor  <- evaluate (toDevice cudaDev cpuKeysTensor)
+  !gpuValsTensor  <- evaluate (toDevice cudaDev cpuValsTensor)
+  !gpuQueryTensor <- evaluate (toDevice cudaDev cpuQueryTensor)
+  syncCuda
+
+  ----------------------------------------------------------------------------
+  -- 1. Key-Value Batch Insert
+  ----------------------------------------------------------------------------
+  putStrLn "\n[1] Batch Insert (Int64 -> Float)"
+
+  (cpuInsertMs, cpuMap) <- timeMs iters $
+    evaluate $ force $ HM.fromList kvPairs
+  printRow "Data.HashMap.Strict (fromList)" n cpuInsertMs cpuInsertMs
+
+  -- Benchmark CudaHashMap insert (GPU-resident tensors)
+  (gpuInsertMs, gpuMap) <- timeMs iters $ do
+    m <- newCudaHashMapWithCapacity (n * 2)
+    insertBatch m gpuKeysTensor gpuValsTensor
+    return m
+  printRow "CudaHashMap insertBatch (GPU tensor)" n cpuInsertMs gpuInsertMs
+
+  -- Benchmark CudaHashMap insert (including CPU -> GPU transfer)
+  (e2eInsertMs, _) <- timeMs iters $ do
+    m <- newCudaHashMapWithCapacity (n * 2)
+    insertBatch m cpuKeysTensor cpuValsTensor
+    return m
+  printRow "CudaHashMap insertBatch (CPU -> GPU)" n cpuInsertMs e2eInsertMs
+
+  ----------------------------------------------------------------------------
+  -- 2. Key-Value Batch Lookup (dynamic_map vs static_map vs Data.HashMap)
+  ----------------------------------------------------------------------------
+  putStrLn "\n[2] Batch Lookup (Int64 -> Float, 50% hit / 50% miss)"
+
+  (cpuLookupMs, _) <- timeMs iters $
+    evaluate $ force $ map (\k -> HM.lookupDefault (-1.0) k cpuMap) queryList
+  printRow "Data.HashMap.Strict (lookupDefault)" n cpuLookupMs cpuLookupMs
+
+  -- Warmup GPU lookup
+  replicateM_ 2 $ lookupBatch gpuMap gpuQueryTensor
+  syncCuda
+
+  (gpuDynLookupMs, _) <- timeMs iters $
+    lookupBatch gpuMap gpuQueryTensor
+  printRow "CudaHashMap lookupBatch (dynamic_map, GPU)" n cpuLookupMs gpuDynLookupMs
+
+  -- Finalize (dynamic_map -> static_map)
+  (finalizeMs, _) <- timeMs 1 $
+    finalizeCudaHashMap gpuMap
+  putStrLn $ printf "  [finalizeCudaHashMap: %.3f ms]" finalizeMs
+
+  (gpuStaticLookupMs, _) <- timeMs iters $
+    lookupBatch gpuMap gpuQueryTensor
+  printRow "CudaHashMap lookupBatch (static_map, GPU)" n cpuLookupMs gpuStaticLookupMs
+
+  (e2eStaticLookupMs, _) <- timeMs iters $ do
+    resGpu <- lookupBatch gpuMap cpuQueryTensor
+    evaluate (toDevice cpuDev resGpu)
+  printRow "CudaHashMap lookupBatch (static_map, CPU<->GPU)" n cpuLookupMs e2eStaticLookupMs
+
+  ----------------------------------------------------------------------------
+  -- 3. Key-to-Index Assignment & Lookup (insertKeys / lookupIndices)
+  ----------------------------------------------------------------------------
+  putStrLn "\n[3] Key-to-Index Assignment (insertKeys) & Lookup (lookupIndices)"
+
+  (cpuIdxInsertMs, (cpuIdxMap, _)) <- timeMs iters $
+    evaluate $ force $ cpuInsertKeys keysList
+  printRow "Data.HashMap.Strict (key -> 0..K-1 insert)" n cpuIdxInsertMs cpuIdxInsertMs
+
+  (gpuIdxInsertMs, gpuIdxMap) <- timeMs iters $ do
+    m <- newCudaHashMapWithCapacity (n * 2)
+    _ <- insertKeys m gpuKeysTensor
+    return m
+  printRow "CudaHashMap insertKeys (GPU tensor)" n cpuIdxInsertMs gpuIdxInsertMs
+
+  finalizeCudaHashMap gpuIdxMap
+  syncCuda
+
+  (cpuIdxLookupMs, _) <- timeMs iters $
+    evaluate $ force $ map (\k -> HM.lookupDefault (-1) k cpuIdxMap) queryList
+  printRow "Data.HashMap.Strict (index lookup)" n cpuIdxLookupMs cpuIdxLookupMs
+
+  (gpuIdxLookupMs, _) <- timeMs iters $
+    lookupIndices gpuIdxMap gpuQueryTensor
+  printRow "CudaHashMap lookupIndices (static_map, GPU)" n cpuIdxLookupMs gpuIdxLookupMs
+  putStrLn ""
+
+main :: IO ()
+main = do
+  putStrLn "Benchmarking Torch.CudaHashMap vs Data.HashMap.Strict (unordered-containers)"
+  -- Warm up CUDA context and allocator
+  warmupMap <- newCudaHashMapWithCapacity 100000
+  let wKeys = asTensor ([1, 2, 3] :: [Int64])
+      wVals = asTensor ([1.0, 2.0, 3.0] :: [Float])
+  insertBatch warmupMap wKeys wVals
+  _ <- lookupBatch warmupMap wKeys
+  syncCuda
+
+  forM_ [10000, 100000, 1000000] runBenchmarkForSize
diff --git a/cbits/CMakeLists.txt b/cbits/CMakeLists.txt
new file mode 100644
--- /dev/null
+++ b/cbits/CMakeLists.txt
@@ -0,0 +1,91 @@
+cmake_minimum_required(VERSION 3.23.1)
+
+# cuCollections requires Volta (sm_70) or newer for scoped atomic operations
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "" OR CMAKE_CUDA_ARCHITECTURES STREQUAL "52")
+    set(CMAKE_CUDA_ARCHITECTURES "70;75;80;86;89;90" CACHE STRING "CUDA architectures" FORCE)
+endif()
+
+project(cuco_wrapper LANGUAGES CXX CUDA)
+
+# nvcc 12.x Extended device lambdas support
+set(CMAKE_CUDA_STANDARD 17)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_EXTENSIONS ON)
+set(CMAKE_CUDA_STANDARD_COMPUSE ON)
+
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda")
+
+include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/CPM.cmake")
+
+# Fetch cuCollections via CPM
+CPMAddPackage(
+    NAME cucollections
+    GITHUB_REPOSITORY NVIDIA/cuCollections
+    GIT_TAG dev
+)
+
+# Configure LibTorch include paths (matching the LibTorch version used under the hood by HaskTorch).
+#    Resolved in the order: CMake variables (-DLIBTORCH_HOME=..., etc.) -> environment variables -> default values.
+
+if(NOT DEFINED LIBTORCH_HOME OR LIBTORCH_HOME STREQUAL "")
+    if(DEFINED ENV{LIBTORCH_HOME} AND NOT "$ENV{LIBTORCH_HOME}" STREQUAL "")
+        set(LIBTORCH_HOME "$ENV{LIBTORCH_HOME}")
+    elseif(DEFINED ENV{XDG_CACHE_HOME} AND NOT "$ENV{XDG_CACHE_HOME}" STREQUAL "")
+        set(LIBTORCH_HOME "$ENV{XDG_CACHE_HOME}/libtorch")
+    elseif(DEFINED ENV{HOME} AND NOT "$ENV{HOME}" STREQUAL "")
+        set(LIBTORCH_HOME "$ENV{HOME}/.cache/libtorch")
+    endif()
+endif()
+
+if(NOT DEFINED LIBTORCH_VERSION OR LIBTORCH_VERSION STREQUAL "")
+    if(DEFINED ENV{LIBTORCH_VERSION} AND NOT "$ENV{LIBTORCH_VERSION}" STREQUAL "")
+        set(LIBTORCH_VERSION "$ENV{LIBTORCH_VERSION}")
+    else()
+        set(LIBTORCH_VERSION "2.9.1")
+    endif()
+endif()
+
+if(NOT DEFINED LIBTORCH_CUDA_VERSION OR LIBTORCH_CUDA_VERSION STREQUAL "")
+    if(DEFINED ENV{LIBTORCH_CUDA_VERSION} AND NOT "$ENV{LIBTORCH_CUDA_VERSION}" STREQUAL "")
+        set(LIBTORCH_CUDA_VERSION "$ENV{LIBTORCH_CUDA_VERSION}")
+    else()
+        set(LIBTORCH_CUDA_VERSION "cu128")
+    endif()
+endif()
+
+if(NOT DEFINED OS OR OS STREQUAL "")
+    if(DEFINED ENV{OS} AND NOT "$ENV{OS}" STREQUAL "")
+        set(OS "$ENV{OS}")
+    else()
+        execute_process(COMMAND uname -s OUTPUT_VARIABLE OS OUTPUT_STRIP_TRAILING_WHITESPACE)
+    endif()
+endif()
+string(TOLOWER "${OS}" OS)
+
+execute_process(COMMAND uname -m OUTPUT_VARIABLE LIBTORCH_ARCH OUTPUT_STRIP_TRAILING_WHITESPACE)
+
+set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/${LIBTORCH_CUDA_VERSION}")
+if(NOT EXISTS "${LIBTORCH_TARGET_DIR}/include" AND EXISTS "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu/include")
+    set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu")
+endif()
+
+include_directories("${LIBTORCH_HOME}/include")
+include_directories("${LIBTORCH_HOME}/include/torch/csrc/api/include")
+include_directories("${LIBTORCH_TARGET_DIR}/include")
+include_directories("${LIBTORCH_TARGET_DIR}/include/torch/csrc/api/include")
+
+# Define the wrapper library
+add_library(cuco_wrapper STATIC cuco_wrapper.cu)
+
+# Automatically configure header paths for cuCollections
+target_link_libraries(cuco_wrapper PRIVATE cuco::cuco)
+
+# Place output directly in the build directory and also remove _deps and CMakeCache.txt on 'make clean'
+set_target_properties(cuco_wrapper PROPERTIES
+    POSITION_INDEPENDENT_CODE ON
+    CUDA_SEPARABLE_COMPILATION OFF
+    ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+    ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/_deps;${CMAKE_BINARY_DIR}/CMakeCache.txt"
+)
diff --git a/cbits/cmake/CPM.cmake b/cbits/cmake/CPM.cmake
new file mode 100644
--- /dev/null
+++ b/cbits/cmake/CPM.cmake
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: MIT
+#
+# SPDX-FileCopyrightText: Copyright (c) 2019-2023 Lars Melchior and contributors
+
+set(CPM_DOWNLOAD_VERSION 0.43.2)
+set(CPM_HASH_SUM "49a3bef91ceb65bb66d57255e12d1ffd22abc2f6408fa9fe4534c544a2f232aa")
+
+if(CPM_PATH)
+  set(CPM_DOWNLOAD_LOCATION "${CPM_PATH}/CPM.cmake")
+elseif(DEFINED ENV{CPM_PATH})
+  file(TO_CMAKE_PATH "$ENV{CPM_PATH}/CPM.cmake" CPM_DOWNLOAD_LOCATION)
+elseif(CPM_SOURCE_CACHE)
+  set(CPM_DOWNLOAD_LOCATION "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+elseif(DEFINED ENV{CPM_SOURCE_CACHE})
+  set(CPM_DOWNLOAD_LOCATION "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+else()
+  set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
+endif()
+
+# Expand relative path. This is important if the provided path contains a tilde (~)
+get_filename_component(CPM_DOWNLOAD_LOCATION ${CPM_DOWNLOAD_LOCATION} ABSOLUTE)
+
+file(DOWNLOAD
+     https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake
+     ${CPM_DOWNLOAD_LOCATION} EXPECTED_HASH SHA256=${CPM_HASH_SUM}
+)
+
+include(${CPM_DOWNLOAD_LOCATION})
diff --git a/cbits/cuco_wrapper.cu b/cbits/cuco_wrapper.cu
new file mode 100644
--- /dev/null
+++ b/cbits/cuco_wrapper.cu
@@ -0,0 +1,520 @@
+#include "cuco_wrapper.h"
+
+#include <cuco/dynamic_map.cuh>
+#include <cuco/static_map.cuh>
+#include <cuda_runtime.h>
+#include <dlfcn.h>
+
+#include <thrust/copy.h>
+#include <thrust/device_vector.h>
+#include <thrust/execution_policy.h>
+#include <thrust/fill.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/sort.h>
+#include <thrust/tuple.h>
+#include <thrust/unique.h>
+
+#include <cuda/std/functional>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <memory>
+#include <type_traits>
+#include <variant>
+
+namespace {
+
+// c10::cuda::CUDAStream (16-byte trivially copyable value wrapping c10::Stream)
+struct RawTorchCUDAStream {
+    std::uint64_t lo;
+    std::uint64_t hi;
+};
+
+// Dynamically resolve LibTorch's (libc10_cuda.so) stream accessor functions at runtime.
+// This avoids "DSO missing from command line" errors in environments where libtorch-ffi
+// does not pass -lc10_cuda directly on the linker command line, as well as undefined
+// symbol errors when linking against CPU-only LibTorch builds.
+inline cudaStream_t get_current_torch_cuda_stream(int device_index = -1) {
+    using GetCurrentCUDAStreamFn = RawTorchCUDAStream (*)(std::int8_t);
+    using CUDAStreamStreamFn     = cudaStream_t (*)(const RawTorchCUDAStream*);
+
+    static const auto get_current_fn = reinterpret_cast<GetCurrentCUDAStreamFn>(
+        dlsym(RTLD_DEFAULT, "_ZN3c104cuda20getCurrentCUDAStreamEa"));
+    static const auto stream_fn = reinterpret_cast<CUDAStreamStreamFn>(
+        dlsym(RTLD_DEFAULT, "_ZNK3c104cuda10CUDAStream6streamEv"));
+
+    if (get_current_fn != nullptr && stream_fn != nullptr) {
+        RawTorchCUDAStream raw = get_current_fn(static_cast<std::int8_t>(device_index));
+        return stream_fn(&raw);
+    }
+    return nullptr;
+}
+
+// RAII guard that switches to the target GPU device within the current scope (equivalent to c10::cuda::CUDAGuard)
+struct ScopedCudaDevice {
+    int prev_device = -1;
+    bool switched = false;
+
+    explicit ScopedCudaDevice(int target_device) {
+        if (target_device >= 0) {
+            if (cudaGetDevice(&prev_device) == cudaSuccess && prev_device != target_device) {
+                cudaSetDevice(target_device);
+                switched = true;
+            }
+        }
+    }
+
+    ~ScopedCudaDevice() {
+        if (switched && prev_device >= 0) {
+            cudaSetDevice(prev_device);
+        }
+    }
+
+    ScopedCudaDevice(const ScopedCudaDevice&) = delete;
+    ScopedCudaDevice& operator=(const ScopedCudaDevice&) = delete;
+};
+
+template <typename To, typename From>
+inline To bit_cast_val(From src) noexcept {
+    static_assert(sizeof(To) == sizeof(From), "bit_cast_val requires matching sizes");
+    To dst;
+    std::memcpy(&dst, &src, sizeof(To));
+    return dst;
+}
+
+template <typename K>
+K make_empty_key_sentinel(int key_dtype) noexcept {
+    if constexpr (sizeof(K) == 4) {
+        if (key_dtype == CUCO_DTYPE_FLOAT) {
+            return bit_cast_val<K>(std::numeric_limits<float>::quiet_NaN());
+        }
+        return std::numeric_limits<K>::min();
+    } else {
+        if (key_dtype == CUCO_DTYPE_DOUBLE) {
+            return bit_cast_val<K>(std::numeric_limits<double>::quiet_NaN());
+        }
+        return std::numeric_limits<K>::min();
+    }
+}
+
+template <typename V>
+V make_empty_val_sentinel(int val_dtype) noexcept {
+    if constexpr (sizeof(V) == 4) {
+        if (val_dtype == CUCO_DTYPE_FLOAT) {
+            return bit_cast_val<V>(std::numeric_limits<float>::quiet_NaN());
+        }
+        return static_cast<V>(-1);
+    } else {
+        if (val_dtype == CUCO_DTYPE_DOUBLE) {
+            return bit_cast_val<V>(std::numeric_limits<double>::quiet_NaN());
+        }
+        return static_cast<V>(-1);
+    }
+}
+
+template <typename K, typename V>
+struct TupleToCucoPair {
+    template <typename TupleLike>
+    __host__ __device__ cuco::pair<K, V> operator()(TupleLike const& t) const noexcept {
+        return cuco::pair<K, V>{
+            static_cast<K>(thrust::get<0>(t)),
+            static_cast<V>(thrust::get<1>(t))
+        };
+    }
+};
+
+template <typename K, typename V, typename KeyIt, typename ValIt>
+auto make_cuco_pair_iterator(KeyIt k_it, ValIt v_it) {
+    return thrust::make_transform_iterator(
+        thrust::make_zip_iterator(thrust::make_tuple(k_it, v_it)),
+        TupleToCucoPair<K, V>{}
+    );
+}
+
+struct MapBase {
+    virtual ~MapBase() = default;
+    virtual int key_dtype() const noexcept = 0;
+    virtual int val_dtype() const noexcept = 0;
+    virtual size_t size(cudaStream_t stream) const = 0;
+    virtual void insert_batch(int num_elements, const void* d_keys, const void* d_values, cudaStream_t stream) = 0;
+    virtual void lookup_batch(int num_elements, const void* d_keys, void* d_values, cudaStream_t stream) const = 0;
+    virtual void insert_keys(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) = 0;
+    virtual void lookup_indices(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) const = 0;
+    virtual void finalize(cudaStream_t stream) = 0;
+};
+
+template <typename K, typename V>
+class TypedCudaHashMap final : public MapBase {
+    using dynamic_map_t = cuco::dynamic_map<K, V>;
+    using static_map_t  = cuco::static_map<K, V>;
+
+    int key_dtype_;
+    int val_dtype_;
+    K empty_key_;
+    V empty_val_;
+    std::variant<dynamic_map_t, static_map_t> map_storage_;
+
+public:
+    TypedCudaHashMap(int k_dtype, int v_dtype, size_t initial_capacity, cudaStream_t stream)
+        : key_dtype_(k_dtype),
+          val_dtype_(v_dtype),
+          empty_key_(make_empty_key_sentinel<K>(k_dtype)),
+          empty_val_(make_empty_val_sentinel<V>(v_dtype)),
+          map_storage_(
+              std::in_place_type<dynamic_map_t>,
+              // cuCollections' dynamic_map requires an initial capacity of at least 100,000 due to its internal min_insert_size (1e4)
+              std::max<size_t>(initial_capacity, 100000),
+              cuco::empty_key<K>{empty_key_},
+              cuco::empty_value<V>{empty_val_},
+              cuda::std::equal_to<K>{},
+              cuco::linear_probing<4, cuco::default_hash_function<K>>{},
+              cuco::cuda_thread_scope<cuda::thread_scope_device>{},
+              cuco::storage<1>{},
+              cuco::cuda_allocator<cuco::pair<K, V>>{},
+              cuda::stream_ref{stream})
+    {}
+
+    int key_dtype() const noexcept override { return key_dtype_; }
+    int val_dtype() const noexcept override { return val_dtype_; }
+
+    size_t size(cudaStream_t stream) const override {
+        if (std::holds_alternative<dynamic_map_t>(map_storage_)) {
+            return std::get<dynamic_map_t>(map_storage_).size();
+        } else {
+            return std::get<static_map_t>(map_storage_).size(cuda::stream_ref{stream});
+        }
+    }
+
+    void insert_batch(int num_elements, const void* d_keys, const void* d_values, cudaStream_t stream) override {
+        if (num_elements <= 0) return;
+
+        const K* keys = static_cast<const K*>(d_keys);
+        const V* values = static_cast<const V*>(d_values);
+        auto pairs_begin = make_cuco_pair_iterator<K, V>(keys, values);
+        auto pairs_end = pairs_begin + num_elements;
+
+        // Overwrite with the new value if the key already exists (insert_or_assign)
+        std::visit([&](auto& m) {
+            m.insert_or_assign(pairs_begin, pairs_end, cuda::stream_ref{stream});
+        }, map_storage_);
+    }
+
+    void lookup_batch(int num_elements, const void* d_keys, void* d_values, cudaStream_t stream) const override {
+        if (num_elements <= 0) return;
+
+        const K* keys = static_cast<const K*>(d_keys);
+        V* values = static_cast<V*>(d_values);
+
+        std::visit([&](auto const& m) {
+            m.find_async(keys, keys + num_elements, values, cuda::stream_ref{stream});
+        }, map_storage_);
+    }
+
+    void insert_keys(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) override {
+        if (num_elements <= 0) return;
+
+        if constexpr (std::is_same_v<V, int64_t>) {
+            const K* keys = static_cast<const K*>(d_keys);
+            auto exec = thrust::cuda::par.on(stream);
+
+            // 1. Deduplicate keys within the input batch using sort + unique
+            thrust::device_vector<K> unique_keys(num_elements);
+            thrust::copy(exec, keys, keys + num_elements, unique_keys.begin());
+            thrust::sort(exec, unique_keys.begin(), unique_keys.end());
+            auto unique_end = thrust::unique(exec, unique_keys.begin(), unique_keys.end());
+            size_t num_unique = static_cast<size_t>(unique_end - unique_keys.begin());
+
+            // 2. Check in batch which unique keys are already present in the map
+            thrust::device_vector<bool> is_contained(num_unique);
+            std::visit([&](auto const& m) {
+                m.contains_async(unique_keys.begin(), unique_end, is_contained.begin(), cuda::stream_ref{stream});
+            }, map_storage_);
+
+            // 3. Extract only the unregistered keys
+            thrust::device_vector<K> new_keys(num_unique);
+            auto new_keys_end = thrust::copy_if(
+                exec,
+                unique_keys.begin(),
+                unique_end,
+                is_contained.begin(),
+                new_keys.begin(),
+                cuda::std::logical_not<bool>{}
+            );
+            size_t num_new = static_cast<size_t>(new_keys_end - new_keys.begin());
+
+            // 4. Assign sequential Int64 indices in [base_idx, base_idx + num_new) to the new keys and insert them
+            if (num_new > 0) {
+                int64_t base_idx = static_cast<int64_t>(size(stream));
+                auto idx_begin = thrust::make_counting_iterator<int64_t>(base_idx);
+                auto pairs_begin = make_cuco_pair_iterator<K, int64_t>(new_keys.begin(), idx_begin);
+                auto pairs_end = pairs_begin + num_new;
+
+                std::visit([&](auto& m) {
+                    m.insert(pairs_begin, pairs_end, cuda::stream_ref{stream});
+                }, map_storage_);
+            }
+
+            // 5. Write the 0-based indices corresponding to all original input keys into d_indices
+            std::visit([&](auto const& m) {
+                m.find_async(keys, keys + num_elements, d_indices, cuda::stream_ref{stream});
+            }, map_storage_);
+        }
+    }
+
+    void lookup_indices(int num_elements, const void* d_keys, int64_t* d_indices, cudaStream_t stream) const override {
+        if (num_elements <= 0) return;
+
+        if constexpr (std::is_same_v<V, int64_t>) {
+            const K* keys = static_cast<const K*>(d_keys);
+            std::visit([&](auto const& m) {
+                m.find_async(keys, keys + num_elements, d_indices, cuda::stream_ref{stream});
+            }, map_storage_);
+        }
+    }
+
+    void finalize(cudaStream_t stream) override {
+        if (!std::holds_alternative<dynamic_map_t>(map_storage_)) {
+            return;
+        }
+
+        auto& d_map = std::get<dynamic_map_t>(map_storage_);
+        size_t current_size = d_map.size();
+
+        thrust::device_vector<K> all_keys(current_size);
+        thrust::device_vector<V> all_vals(current_size);
+        size_t actual_size = 0;
+
+        if (current_size > 0) {
+            auto [keys_end, vals_end] = d_map.retrieve_all(
+                all_keys.begin(), all_vals.begin(), cuda::stream_ref{stream});
+            actual_size = static_cast<size_t>(keys_end - all_keys.begin());
+        }
+
+        size_t optimal_capacity = std::max<size_t>(actual_size * 2, 1024);
+
+        // Construct static_map in-place while destroying the old dynamic_map
+        auto& s_map = map_storage_.template emplace<static_map_t>(
+            optimal_capacity,
+            cuco::empty_key<K>{empty_key_},
+            cuco::empty_value<V>{empty_val_},
+            cuda::std::equal_to<K>{},
+            cuco::linear_probing<4, cuco::default_hash_function<K>>{},
+            cuco::cuda_thread_scope<cuda::thread_scope_device>{},
+            cuco::storage<1>{},
+            cuco::cuda_allocator<cuco::pair<K, V>>{},
+            cuda::stream_ref{stream}
+        );
+
+        if (actual_size > 0) {
+            auto pairs_begin = make_cuco_pair_iterator<K, V>(all_keys.begin(), all_vals.begin());
+            auto pairs_end = pairs_begin + actual_size;
+            s_map.insert(pairs_begin, pairs_end, cuda::stream_ref{stream});
+        }
+    }
+};
+
+std::unique_ptr<MapBase> create_typed_map(int key_dtype, int val_dtype, size_t initial_capacity, cudaStream_t stream) {
+    bool key_is_64 = (key_dtype == CUCO_DTYPE_INT64 || key_dtype == CUCO_DTYPE_DOUBLE);
+    bool val_is_64 = (val_dtype == CUCO_DTYPE_INT64 || val_dtype == CUCO_DTYPE_DOUBLE);
+
+    if (!key_is_64 && !val_is_64) {
+        return std::make_unique<TypedCudaHashMap<int32_t, int32_t>>(key_dtype, val_dtype, initial_capacity, stream);
+    } else if (!key_is_64 && val_is_64) {
+        return std::make_unique<TypedCudaHashMap<int32_t, int64_t>>(key_dtype, val_dtype, initial_capacity, stream);
+    } else if (key_is_64 && !val_is_64) {
+        return std::make_unique<TypedCudaHashMap<int64_t, int32_t>>(key_dtype, val_dtype, initial_capacity, stream);
+    } else {
+        return std::make_unique<TypedCudaHashMap<int64_t, int64_t>>(key_dtype, val_dtype, initial_capacity, stream);
+    }
+}
+
+template <typename V>
+void fill_uninitialized_lookup(int val_dtype, int num_elements, void* d_values, cudaStream_t stream) {
+    V sentinel = make_empty_val_sentinel<V>(val_dtype);
+    V* out = static_cast<V*>(d_values);
+    thrust::fill(thrust::cuda::par.on(stream), out, out + num_elements, sentinel);
+}
+
+} // namespace
+
+struct CudaHashMapOpaque {
+    int device_index;
+    size_t initial_capacity;
+    std::unique_ptr<MapBase> impl;
+
+    explicit CudaHashMapOpaque(size_t cap = 100000, int dev = -1)
+        : device_index(dev), initial_capacity(cap), impl(nullptr) {}
+
+    MapBase* ensure_initialized(int key_dtype, int val_dtype, cudaStream_t stream) {
+        if (!impl) {
+            if (device_index < 0) {
+                int cur_dev = 0;
+                if (cudaGetDevice(&cur_dev) == cudaSuccess) {
+                    device_index = cur_dev;
+                } else {
+                    device_index = 0;
+                }
+            }
+            impl = create_typed_map(key_dtype, val_dtype, initial_capacity, stream);
+        }
+        return impl.get();
+    }
+};
+
+extern "C" {
+
+CudaHashMapOpaque* custom_cuco_create(void) {
+    return new CudaHashMapOpaque(100000, -1);
+}
+
+CudaHashMapOpaque* custom_cuco_create_with_capacity(size_t initial_capacity) {
+    return new CudaHashMapOpaque(initial_capacity, -1);
+}
+
+CudaHashMapOpaque* custom_cuco_create_on_device(int device_index, size_t initial_capacity) {
+    return new CudaHashMapOpaque(initial_capacity, device_index);
+}
+
+void custom_cuco_destroy(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr) return;
+    if (map_ptr->impl && map_ptr->device_index >= 0) {
+        ScopedCudaDevice guard(map_ptr->device_index);
+        map_ptr->impl.reset();
+    }
+    delete map_ptr;
+}
+
+int custom_cuco_device_index(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr) return -1;
+    return map_ptr->device_index;
+}
+
+void custom_cuco_set_device_index(CudaHashMapOpaque* map_ptr, int device_index) {
+    if (!map_ptr) return;
+    if (!map_ptr->impl) {
+        map_ptr->device_index = device_index;
+    }
+}
+
+size_t custom_cuco_size(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr || !map_ptr->impl) return 0;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    return map_ptr->impl->size(stream);
+}
+
+int custom_cuco_key_dtype(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr || !map_ptr->impl) return -1;
+    return map_ptr->impl->key_dtype();
+}
+
+int custom_cuco_val_dtype(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr || !map_ptr->impl) return -1;
+    return map_ptr->impl->val_dtype();
+}
+
+void custom_cuco_insert_keys(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int num_elements,
+    const void* d_keys,
+    int64_t* d_indices)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    MapBase* m = map_ptr->ensure_initialized(key_dtype, CUCO_DTYPE_INT64, stream);
+    m->insert_keys(num_elements, d_keys, d_indices, stream);
+}
+
+void custom_cuco_lookup_indices(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int num_elements,
+    const void* d_keys,
+    int64_t* d_indices)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    if (!map_ptr->impl) {
+        fill_uninitialized_lookup<int64_t>(CUCO_DTYPE_INT64, num_elements, d_indices, stream);
+        return;
+    }
+    map_ptr->impl->lookup_indices(num_elements, d_keys, d_indices, stream);
+}
+
+void custom_cuco_insert_batch_typed(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int val_dtype,
+    int num_elements,
+    const void* d_keys,
+    const void* d_values)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    MapBase* m = map_ptr->ensure_initialized(key_dtype, val_dtype, stream);
+    m->insert_batch(num_elements, d_keys, d_values, stream);
+}
+
+void custom_cuco_lookup_batch_typed(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int val_dtype,
+    int num_elements,
+    const void* d_keys,
+    void* d_values)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    if (!map_ptr->impl) {
+        if (val_dtype == CUCO_DTYPE_INT64 || val_dtype == CUCO_DTYPE_DOUBLE) {
+            fill_uninitialized_lookup<int64_t>(val_dtype, num_elements, d_values, stream);
+        } else {
+            fill_uninitialized_lookup<int32_t>(val_dtype, num_elements, d_values, stream);
+        }
+        return;
+    }
+    (void)key_dtype;
+    map_ptr->impl->lookup_batch(num_elements, d_keys, d_values, stream);
+}
+
+void custom_cuco_insert_batch(
+    CudaHashMapOpaque* map_ptr,
+    int num_elements,
+    void* d_keys,
+    void* d_values)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    int k_dtype = map_ptr->impl ? map_ptr->impl->key_dtype() : CUCO_DTYPE_INT32;
+    int v_dtype = map_ptr->impl ? map_ptr->impl->val_dtype() : CUCO_DTYPE_INT32;
+    custom_cuco_insert_batch_typed(map_ptr, k_dtype, v_dtype, num_elements, d_keys, d_values);
+}
+
+void custom_cuco_lookup_batch(
+    CudaHashMapOpaque* map_ptr,
+    int num_elements,
+    void* d_keys,
+    void* d_values)
+{
+    if (!map_ptr || num_elements <= 0) return;
+    int k_dtype = map_ptr->impl ? map_ptr->impl->key_dtype() : CUCO_DTYPE_INT32;
+    int v_dtype = map_ptr->impl ? map_ptr->impl->val_dtype() : CUCO_DTYPE_INT32;
+    custom_cuco_lookup_batch_typed(map_ptr, k_dtype, v_dtype, num_elements, d_keys, d_values);
+}
+
+void custom_cuco_finalize(CudaHashMapOpaque* map_ptr) {
+    if (!map_ptr || !map_ptr->impl) return;
+    ScopedCudaDevice guard(map_ptr->device_index);
+    cudaStream_t stream = get_current_torch_cuda_stream(map_ptr->device_index);
+    map_ptr->impl->finalize(stream);
+}
+
+} // extern "C"
diff --git a/cbits/cuco_wrapper.h b/cbits/cuco_wrapper.h
new file mode 100644
--- /dev/null
+++ b/cbits/cuco_wrapper.h
@@ -0,0 +1,92 @@
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Opaque struct representing a C++ hash map instance
+typedef struct CudaHashMapOpaque CudaHashMapOpaque;
+
+// Supported tensor data types (DType)
+typedef enum {
+    CUCO_DTYPE_INT32  = 0,
+    CUCO_DTYPE_INT64  = 1,
+    CUCO_DTYPE_FLOAT  = 2,
+    CUCO_DTYPE_DOUBLE = 3
+} CucoDType;
+
+// Hash map creation and destruction
+// When device_index is -1, uses the device of the first tensor passed in or the current GPU device
+CudaHashMapOpaque* custom_cuco_create(void);
+CudaHashMapOpaque* custom_cuco_create_with_capacity(size_t initial_capacity);
+CudaHashMapOpaque* custom_cuco_create_on_device(int device_index, size_t initial_capacity);
+void custom_cuco_destroy(CudaHashMapOpaque* map_ptr);
+
+// Get or set the CUDA device index associated with the map (-1 when unset)
+int custom_cuco_device_index(CudaHashMapOpaque* map_ptr);
+void custom_cuco_set_device_index(CudaHashMapOpaque* map_ptr, int device_index);
+
+// Get the number of currently registered unique keys
+size_t custom_cuco_size(CudaHashMapOpaque* map_ptr);
+
+// Get the key and value DTypes configured for the map (-1 when uninitialized)
+int custom_cuco_key_dtype(CudaHashMapOpaque* map_ptr);
+int custom_cuco_val_dtype(CudaHashMapOpaque* map_ptr);
+
+// Automatically assign 0-based sequential indices (Int64) to keys and write an Int64
+// index tensor that can be used directly with indexSelect or indexPut
+void custom_cuco_insert_keys(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int num_elements,
+    const void* d_keys,
+    int64_t* d_indices);
+
+// Look up Int64 indices for keys without inserting new keys (unregistered keys get -1)
+void custom_cuco_lookup_indices(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int num_elements,
+    const void* d_keys,
+    int64_t* d_indices);
+
+// Typed batch insertion (existing keys are overwritten with new values: insert_or_assign)
+void custom_cuco_insert_batch_typed(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int val_dtype,
+    int num_elements,
+    const void* d_keys,
+    const void* d_values);
+
+// Typed batch lookup (unregistered keys get -1 for Int32/Int64, or NaN for Float/Double)
+void custom_cuco_lookup_batch_typed(
+    CudaHashMapOpaque* map_ptr,
+    int key_dtype,
+    int val_dtype,
+    int num_elements,
+    const void* d_keys,
+    void* d_values);
+
+// Compatibility interface (operates with default Int32 or the already configured types)
+void custom_cuco_insert_batch(
+    CudaHashMapOpaque* map_ptr,
+    int num_elements,
+    void* d_keys,
+    void* d_values);
+
+void custom_cuco_lookup_batch(
+    CudaHashMapOpaque* map_ptr,
+    int num_elements,
+    void* d_keys,
+    void* d_values);
+
+// Optimization function to call when entering a steady state (converts dynamic_map -> static_map)
+void custom_cuco_finalize(CudaHashMapOpaque* map_ptr);
+
+#ifdef __cplusplus
+}
+#endif
diff --git a/hasktorch-CudaHashMap.cabal b/hasktorch-CudaHashMap.cabal
new file mode 100644
--- /dev/null
+++ b/hasktorch-CudaHashMap.cabal
@@ -0,0 +1,158 @@
+cabal-version:      3.0
+-- The cabal-version field refers to the version of the .cabal specification,
+-- and can be different from the cabal-install (the tool) version and the
+-- Cabal (the library) version you are using. As such, the Cabal (the library)
+-- version used must be equal or greater than the version stated in this field.
+-- Starting from the specification version 2.2, the cabal-version field must be
+-- the first thing in the cabal file.
+
+-- Initial package description 'hasktorch-CudaHashMap' generated by
+-- 'cabal init'. For further documentation, see:
+--   http://haskell.org/cabal/users-guide/
+--
+-- The name of the package.
+name:               hasktorch-CudaHashMap
+
+-- The package version.
+-- See the Haskell package versioning policy (PVP) for standards
+-- guiding when and how versions should be incremented.
+-- https://pvp.haskell.org
+-- PVP summary:     +-+------- breaking API changes
+--                  | | +----- non-breaking API additions
+--                  | | | +--- code changes with no API change
+version:            0.1.0.0
+
+-- A short (one-line) description of the package.
+synopsis:           HaskTorch interface to the GPU-parallel hashmap in cuCollections
+
+-- A longer description of the package.
+-- description:
+
+license:            BSD-3-Clause
+license-file:       LICENSE
+author:             Susumu Katayama
+maintainer:         skata@cs.miyazaki-u.ac.jp
+
+-- A copyright notice.
+-- copyright:
+category:           Data
+build-type:         Custom
+
+-- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.
+extra-doc-files:    README.md, CHANGELOG.md
+
+-- Extra source files to be distributed with the package, such as examples, or a tutorial module.
+extra-source-files: cbits/cuco_wrapper.cu, cbits/cuco_wrapper.h, cbits/CMakeLists.txt, cbits/cmake/CPM.cmake
+
+custom-setup
+    setup-depends:
+        base >= 4.7 && < 5,
+        Cabal >= 3.0 && < 4,
+        directory >= 1.3 && < 2,
+        process >= 1.6 && < 2
+
+common warnings
+    ghc-options: -Wall
+
+library
+    -- Import common warning flags.
+    import:           warnings
+
+    exposed-modules:     Torch.CudaHashMap
+    hs-source-dirs:      src
+    default-language:    Haskell2010
+--    ghc-options:         -Wall
+
+    build-depends:       base >= 4.7 && < 5
+                     , hasktorch
+                     , libtorch-ffi
+
+    include-dirs:        cbits
+
+    -- Paths to the precompiled custom library (libcuco_wrapper.a) and CUDA libraries
+    -- are automatically resolved as absolute paths in the confHook of Setup.hs.
+
+    -- Libraries to link against
+    --    cuco_wrapper: custom C++ wrapper (libcuco_wrapper.a)
+    --    cuda, cudart: CUDA runtime
+    --    nvrtc: CUDA runtime compilation (required internally by cuco)
+    extra-libraries:     cuco_wrapper
+                       , cuda
+                       , cudart
+                       , nvrtc
+                       , stdc++
+
+    -- GHC options for C++ and CUDA integration
+    --    -pgml g++ : use g++ instead of gcc as the linker driver (required for C++ exception handling and standard library linking)
+    --    --copy-dt-needed-entries : allow resolving indirect DSO symbols such as libc10_cuda.so required by libtorch_cuda.so
+    ghc-options:         -pgml g++
+                         -optl-Wl,--copy-dt-needed-entries
+
+
+    -- Modules included in this library but not exported.
+    -- other-modules:
+
+    -- LANGUAGE extensions used by modules in this package.
+    -- other-extensions:
+
+
+executable CudaHashMap
+    import:           warnings
+    main-is:          Main.hs
+    build-depends:
+        base >= 4.7 && < 5,
+        hasktorch,
+        hasktorch-CudaHashMap
+
+    ghc-options:      -pgml g++
+                      -optl-no-pie
+                      -optl-Lcbits/build
+                      -optl-L/usr/local/cuda/lib64
+                      -optl-Wl,--copy-dt-needed-entries
+
+    hs-source-dirs:   app
+    default-language: Haskell2010
+
+test-suite CudaHashMap-test
+    import:           warnings
+    default-language: Haskell2010
+
+    -- The interface type and version of the test suite.
+    type:             exitcode-stdio-1.0
+
+    hs-source-dirs:   test
+    main-is:          Main.hs
+
+    build-depends:
+        base >= 4.7 && < 5,
+        hasktorch,
+        hasktorch-CudaHashMap
+
+    ghc-options:      -pgml g++
+                      -optl-no-pie
+                      -optl-Lcbits/build
+                      -optl-L/usr/local/cuda/lib64
+                      -optl-Wl,--copy-dt-needed-entries
+
+benchmark CudaHashMap-bench
+    import:           warnings
+    type:             exitcode-stdio-1.0
+    hs-source-dirs:   bench
+    main-is:          Main.hs
+    default-language: Haskell2010
+
+    build-depends:
+        base >= 4.7 && < 5,
+        clock >= 0.8,
+        deepseq >= 1.4,
+        hasktorch,
+        hasktorch-CudaHashMap,
+        unordered-containers >= 0.2
+
+    ghc-options:      -O2
+                      -pgml g++
+                      -optl-no-pie
+                      -optl-Lcbits/build
+                      -optl-L/usr/local/cuda/lib64
+                      -optl-Wl,--copy-dt-needed-entries
+
diff --git a/src/Torch/CudaHashMap.hs b/src/Torch/CudaHashMap.hs
new file mode 100644
--- /dev/null
+++ b/src/Torch/CudaHashMap.hs
@@ -0,0 +1,318 @@
+{-# LANGUAGE ForeignFunctionInterface #-}
+
+module Torch.CudaHashMap
+  ( -- * Data Types
+    CudaHashMap(..)
+  , CudaHashMapOpaque
+    -- * Creation, Optimization, and Metadata
+  , newCudaHashMap
+  , newCudaHashMapWithCapacity
+  , newCudaHashMapOnDevice
+  , hashMapDevice
+  , finalizeCudaHashMap
+  , sizeCudaHashMap
+    -- * Key-Value Batch Insert & Lookup (supports Int32, Int64, Float, Double)
+  , insertBatch
+  , lookupBatch
+    -- * Index Conversion API for indexSelect / indexPut
+  , insertKeys
+  , lookupIndices
+  , hashMapIndexPut
+  , hashMapIndexSelect
+    -- * Low-Level Helper
+  , withDevicePointer
+  ) where
+
+import Foreign.C.Types (CInt(..), CSize(..))
+import Foreign.ForeignPtr (ForeignPtr, newForeignPtr, withForeignPtr)
+import Foreign.Ptr (FunPtr, Ptr, castPtr)
+import Data.Int (Int64)
+
+import Torch.Device (Device(..), DeviceType(..))
+import Torch.DType (DType(..))
+import qualified Torch.DType as DType
+import Torch.Functional (indexPut)
+import Torch.Tensor
+  ( Tensor(..)
+  , contiguous
+  , device
+  , dtype
+  , indexSelect
+  , isContiguous
+  , numel
+  , reshape
+  , shape
+  , toDevice
+  )
+import Torch.TensorFactories (mkFactory)
+import Torch.TensorOptions (defaultOpts, withDevice, withDType)
+import qualified Torch.Internal.Managed.TensorFactories as LibTorch
+import qualified Torch.Internal.Managed.Type.Tensor.Tensor0 as ATen
+
+-- | Opaque type representing a pointer to a C++ hash map instance
+data CudaHashMapOpaque
+
+-- | GPU hash map type handled on the Haskell side
+newtype CudaHashMap = CudaHashMap (ForeignPtr CudaHashMapOpaque)
+
+-- FFI declarations for C++ creation, destruction, and metadata functions
+foreign import ccall "custom_cuco_create"
+  c_create :: IO (Ptr CudaHashMapOpaque)
+
+foreign import ccall "custom_cuco_create_with_capacity"
+  c_create_with_capacity :: CSize -> IO (Ptr CudaHashMapOpaque)
+
+foreign import ccall "custom_cuco_create_on_device"
+  c_create_on_device :: CInt -> CSize -> IO (Ptr CudaHashMapOpaque)
+
+foreign import ccall "&custom_cuco_destroy"
+  c_destroy :: FunPtr (Ptr CudaHashMapOpaque -> IO ())
+
+foreign import ccall "custom_cuco_device_index"
+  c_device_index :: Ptr CudaHashMapOpaque -> IO CInt
+
+foreign import ccall "custom_cuco_set_device_index"
+  c_set_device_index :: Ptr CudaHashMapOpaque -> CInt -> IO ()
+
+foreign import ccall "custom_cuco_size"
+  c_size :: Ptr CudaHashMapOpaque -> IO CSize
+
+foreign import ccall "custom_cuco_val_dtype"
+  c_val_dtype :: Ptr CudaHashMapOpaque -> IO CInt
+
+foreign import ccall "custom_cuco_finalize"
+  c_finalize :: Ptr CudaHashMapOpaque -> IO ()
+
+-- FFI declarations for C++ Key-to-Index functions
+foreign import ccall "custom_cuco_insert_keys"
+  c_insert_keys
+    :: Ptr CudaHashMapOpaque
+    -> CInt
+    -> CInt
+    -> Ptr ()
+    -> Ptr Int64
+    -> IO ()
+
+foreign import ccall "custom_cuco_lookup_indices"
+  c_lookup_indices
+    :: Ptr CudaHashMapOpaque
+    -> CInt
+    -> CInt
+    -> Ptr ()
+    -> Ptr Int64
+    -> IO ()
+
+-- FFI declarations for C++ typed batch Insert / Lookup functions
+foreign import ccall "custom_cuco_insert_batch_typed"
+  c_insert_batch_typed
+    :: Ptr CudaHashMapOpaque
+    -> CInt
+    -> CInt
+    -> CInt
+    -> Ptr ()
+    -> Ptr ()
+    -> IO ()
+
+foreign import ccall "custom_cuco_lookup_batch_typed"
+  c_lookup_batch_typed
+    :: Ptr CudaHashMapOpaque
+    -> CInt
+    -> CInt
+    -> CInt
+    -> Ptr ()
+    -> Ptr ()
+    -> IO ()
+
+-- | Convert a HaskTorch DType to the C++ wrapper's CucoDType code
+toCucoDType :: DType -> CInt
+toCucoDType DType.Int32  = 0
+toCucoDType DType.Int64  = 1
+toCucoDType DType.Float  = 2
+toCucoDType DType.Double = 3
+toCucoDType dt =
+  error $ "Torch.CudaHashMap: unsupported Tensor DType " ++ show dt
+       ++ " (supported: Int32, Int64, Float, Double)"
+
+fromCucoDType :: CInt -> Maybe DType
+fromCucoDType 0 = Just DType.Int32
+fromCucoDType 1 = Just DType.Int64
+fromCucoDType 2 = Just DType.Float
+fromCucoDType 3 = Just DType.Double
+fromCucoDType _ = Nothing
+
+-- | Determine or retrieve the CUDA device associated with the map, and ensure that the input tensor
+-- resides in contiguous memory on that CUDA device.
+-- If the map's device is not yet determined (-1), adopts the input tensor's device index if it is
+-- a CUDA tensor, or defaults to cuda:0 (Device CUDA 0) if it is a CPU tensor, and records it in the map.
+ensureCudaOnMap :: Ptr CudaHashMapOpaque -> Tensor -> IO Tensor
+ensureCudaOnMap mapPtr t = do
+  devIdx <- c_device_index mapPtr
+  targetDev <-
+    if devIdx >= 0
+      then return $ Device CUDA (fromIntegral devIdx)
+      else case device t of
+             d@(Device CUDA idx) -> do
+               c_set_device_index mapPtr (fromIntegral idx)
+               return d
+             _ -> do
+               c_set_device_index mapPtr 0
+               return $ Device CUDA 0
+  let tCuda = if device t == targetDev then t else toDevice targetDev t
+  return $ if isContiguous tCuda then tCuda else contiguous tCuda
+
+-- | Allocate an output tensor with the specified shape, DType, and Device
+allocTensor :: [Int] -> DType -> Device -> IO Tensor
+allocTensor dims dt dev =
+  let opts = withDType dt $ withDevice dev defaultOpts
+  in mkFactory LibTorch.empty_lo dims opts
+
+-- | Extract the raw device pointer (GPU pointer) from a Tensor and run an action with it
+withDevicePointer :: Tensor -> (Ptr () -> IO a) -> IO a
+withDevicePointer t action = do
+  let Unsafe fptr = if isContiguous t then t else contiguous t
+  withForeignPtr fptr $ \_ -> do
+    ptr <- ATen.tensor_data_ptr fptr
+    action ptr
+
+-- | Initialize a hash map (starts as a dynamically resizable dynamic_map).
+-- The target GPU device is automatically set to the device of the first CUDA tensor inserted or queried
+-- (or cuda:0 if the first tensor is on the CPU).
+newCudaHashMap :: IO CudaHashMap
+newCudaHashMap = do
+  ptr <- c_create
+  fptr <- newForeignPtr c_destroy ptr
+  return $ CudaHashMap fptr
+
+-- | Initialize a hash map with a specified initial capacity
+newCudaHashMapWithCapacity :: Int -> IO CudaHashMap
+newCudaHashMapWithCapacity cap = do
+  ptr <- c_create_with_capacity (fromIntegral (max 100000 cap))
+  fptr <- newForeignPtr c_destroy ptr
+  return $ CudaHashMap fptr
+
+-- | Initialize a hash map with an explicit CUDA device (e.g., @Device CUDA 1@) and initial capacity
+newCudaHashMapOnDevice :: Device -> Int -> IO CudaHashMap
+newCudaHashMapOnDevice (Device CUDA devIdx) cap = do
+  ptr <- c_create_on_device (fromIntegral devIdx) (fromIntegral (max 100000 cap))
+  fptr <- newForeignPtr c_destroy ptr
+  return $ CudaHashMap fptr
+newCudaHashMapOnDevice dev _ =
+  error $ "Torch.CudaHashMap.newCudaHashMapOnDevice: expected CUDA device, got " ++ show dev
+
+-- | Get the CUDA device associated with the hash map (returns @Device CUDA 0@ if not yet determined)
+hashMapDevice :: CudaHashMap -> IO Device
+hashMapDevice (CudaHashMap fptr) =
+  withForeignPtr fptr $ \mapPtr -> do
+    devIdx <- c_device_index mapPtr
+    return $ Device CUDA (if devIdx >= 0 then fromIntegral devIdx else 0)
+
+-- | Get the number of unique keys currently registered in the hash map
+sizeCudaHashMap :: CudaHashMap -> IO Int
+sizeCudaHashMap (CudaHashMap fptr) =
+  withForeignPtr fptr $ \mapPtr ->
+    fromIntegral <$> c_size mapPtr
+
+-- | Convert and freeze the internal dynamic_map into a single high-speed static_map once insertions have stabilized
+finalizeCudaHashMap :: CudaHashMap -> IO ()
+finalizeCudaHashMap (CudaHashMap fptr) =
+  withForeignPtr fptr c_finalize
+
+-- | Batch-insert key and value tensors (overwrites with the new value if a key already exists: insert_or_assign).
+-- Both keys and values support Int32, Int64, Float, and Double.
+insertBatch :: CudaHashMap -> Tensor -> Tensor -> IO ()
+insertBatch (CudaHashMap fptr) keysTensor valuesTensor =
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaKeys   <- ensureCudaOnMap mapPtr keysTensor
+    cudaValues <- ensureCudaOnMap mapPtr valuesTensor
+    let kDType  = toCucoDType (dtype cudaKeys)
+        vDType  = toCucoDType (dtype cudaValues)
+        numKeys = numel cudaKeys
+        numVals = numel cudaValues
+    if numKeys /= numVals
+      then error $ "Torch.CudaHashMap.insertBatch: key and value element counts do not match ("
+                ++ show numKeys ++ " vs " ++ show numVals ++ ")"
+      else do
+        let numElements = fromIntegral numKeys
+        withDevicePointer cudaKeys $ \keyPtr ->
+          withDevicePointer cudaValues $ \valPtr ->
+            c_insert_batch_typed mapPtr kDType vDType numElements keyPtr valPtr
+
+-- | Batch-lookup a tensor of keys and return the corresponding tensor of values.
+-- Unregistered keys receive -1 for Int32 / Int64, or NaN for Float / Double.
+lookupBatch :: CudaHashMap -> Tensor -> IO Tensor
+lookupBatch (CudaHashMap fptr) keysTensor =
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor
+    let kDType      = toCucoDType (dtype cudaKeys)
+        numElements = fromIntegral (numel cudaKeys)
+    storedValCode <- c_val_dtype mapPtr
+    let outDType = case fromCucoDType storedValCode of
+                     Just dt -> dt
+                     Nothing -> dtype cudaKeys
+        vDType   = toCucoDType outDType
+    valuesTensor <- allocTensor (shape cudaKeys) outDType (device cudaKeys)
+    withDevicePointer cudaKeys $ \keyPtr ->
+      withDevicePointer valuesTensor $ \valPtr ->
+        c_lookup_batch_typed mapPtr kDType vDType numElements keyPtr valPtr
+    return valuesTensor
+
+-- | Assign and return 0-based sequential indices (DType.Int64) for a tensor of keys.
+-- Returns the existing index for already registered keys, and assigns new sequential indices to unregistered keys.
+-- The returned tensor can be used directly as the index tensor for HaskTorch's 'indexSelect' or 'indexPut'.
+insertKeys :: CudaHashMap -> Tensor -> IO Tensor
+insertKeys (CudaHashMap fptr) keysTensor =
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor
+    let kDType      = toCucoDType (dtype cudaKeys)
+        numElements = fromIntegral (numel cudaKeys)
+    indicesTensor <- allocTensor (shape cudaKeys) DType.Int64 (device cudaKeys)
+    withDevicePointer cudaKeys $ \keyPtr ->
+      withDevicePointer indicesTensor $ \idxPtr ->
+        c_insert_keys mapPtr kDType numElements keyPtr (castPtr idxPtr)
+    return indicesTensor
+
+-- | Look up 0-based Int64 indices corresponding to a tensor of keys without registering new keys.
+-- Unregistered keys receive -1.
+lookupIndices :: CudaHashMap -> Tensor -> IO Tensor
+lookupIndices (CudaHashMap fptr) keysTensor =
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaKeys <- ensureCudaOnMap mapPtr keysTensor
+    let kDType      = toCucoDType (dtype cudaKeys)
+        numElements = fromIntegral (numel cudaKeys)
+    indicesTensor <- allocTensor (shape cudaKeys) DType.Int64 (device cudaKeys)
+    withDevicePointer cudaKeys $ \keyPtr ->
+      withDevicePointer indicesTensor $ \idxPtr ->
+        c_lookup_indices mapPtr kDType numElements keyPtr (castPtr idxPtr)
+    return indicesTensor
+
+-- | Obtain sequential indices from a tensor of keys (registering any unseen keys)
+-- and write valuesTensor into the corresponding indices of tableTensor via indexPut.
+hashMapIndexPut
+  :: CudaHashMap
+  -> Bool    -- ^ accumulate (True to add, False to overwrite)
+  -> Tensor  -- ^ keysTensor
+  -> Tensor  -- ^ valuesTensor
+  -> Tensor  -- ^ tableTensor (target tensor to write into)
+  -> IO Tensor
+hashMapIndexPut hmap@(CudaHashMap fptr) accumulate keysTensor valuesTensor tableTensor = do
+  indices <- insertKeys hmap keysTensor
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaValues <- ensureCudaOnMap mapPtr valuesTensor
+    cudaTable  <- ensureCudaOnMap mapPtr tableTensor
+    let flatIndices = reshape [-1] indices
+    return $ indexPut accumulate [flatIndices] cudaValues cudaTable
+
+-- | Look up registered indices from a tensor of keys and extract elements
+-- from the specified dimension of tableTensor via indexSelect.
+hashMapIndexSelect
+  :: CudaHashMap
+  -> Int     -- ^ dimension (typically 0)
+  -> Tensor  -- ^ keysTensor
+  -> Tensor  -- ^ tableTensor
+  -> IO Tensor
+hashMapIndexSelect hmap@(CudaHashMap fptr) selectDim keysTensor tableTensor = do
+  indices <- lookupIndices hmap keysTensor
+  withForeignPtr fptr $ \mapPtr -> do
+    cudaTable <- ensureCudaOnMap mapPtr tableTensor
+    let flatIndices = reshape [-1] indices
+    return $ indexSelect selectDim flatIndices cudaTable
diff --git a/test/Main.hs b/test/Main.hs
new file mode 100644
--- /dev/null
+++ b/test/Main.hs
@@ -0,0 +1,42 @@
+module Main (main) where
+
+import Data.Int (Int32, Int64)
+import Torch.Functional (indexPut)
+import Torch.Tensor (asTensor, device, indexSelect, toDevice)
+import Torch.TensorFactories (zeros')
+import Torch.CudaHashMap
+
+main :: IO ()
+main = do
+  putStrLn "=== Test 1: Int32 Key -> Int32 Value ==="
+  m1 <- newCudaHashMap
+  let k1 = asTensor ([1, 2, 3] :: [Int32])
+      v1 = asTensor ([10, 20, 30] :: [Int32])
+      v1_overwrite = asTensor ([100, 200, 300] :: [Int32])
+  insertBatch m1 k1 v1
+  insertBatch m1 k1 v1_overwrite
+  finalizeCudaHashMap m1
+  res1 <- lookupBatch m1 (asTensor ([3, 1, 99] :: [Int32]))
+  print res1
+
+  putStrLn "=== Test 2: Int64 Key -> Float Value ==="
+  m2 <- newCudaHashMap
+  let k2 = asTensor ([100, 200, 300] :: [Int64])
+      v2 = asTensor ([1.25, 2.5, 3.75] :: [Float])
+  insertBatch m2 k2 v2
+  finalizeCudaHashMap m2
+  res2 <- lookupBatch m2 (asTensor ([200, 100, 999] :: [Int64]))
+  print res2
+
+  putStrLn "=== Test 3: Key -> Index (Int64) for indexPut / indexSelect ==="
+  m3 <- newCudaHashMap
+  let table0 = zeros' [10]
+      batchKeys1 = asTensor ([500, 600, 500, 700] :: [Int64])
+      batchVals1 = asTensor ([5.0, 6.0, 5.5, 7.0] :: [Float])
+  idx1 <- insertKeys m3 batchKeys1
+  let dev    = device idx1
+      table1 = indexPut False [idx1] (toDevice dev batchVals1) (toDevice dev table0)
+  finalizeCudaHashMap m3
+  qIdx <- lookupIndices m3 (asTensor ([700, 500, 600] :: [Int64]))
+  let out = indexSelect 0 qIdx table1
+  print out