Skip to main content

dynamo_kv_hashing/
compute.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Block-hash computation for a [`crate::Request`].
5
6use dynamo_tokens::{
7    BlockHash, PositionalLineageHash, SaltHash, SequenceHash, TokenBlockSequence, Tokens,
8};
9
10use crate::block::UniversalBlock;
11use crate::error::KvHashingError;
12use crate::request::Request;
13use crate::salt::compute_salt_hash;
14
15impl Request {
16    /// Returns the canonical [`SaltHash`] for this request.
17    pub fn salt_hash(&self) -> Result<SaltHash, KvHashingError> {
18        compute_salt_hash(self.salt(), self.lora_name())
19    }
20
21    /// Returns the rich per-block result.
22    ///
23    /// One [`UniversalBlock`] per *complete* `block_size`-sized window in the request's
24    /// token stream (placeholder slots count toward `block_size`). A trailing partial
25    /// block — fewer than `block_size` slots — is not hashed and not returned.
26    pub fn into_blocks(&self, block_size: u32) -> Result<Vec<UniversalBlock>, KvHashingError> {
27        let salt_hash = self.salt_hash()?;
28        let token_mm = self.token_mm_info();
29        let seq = TokenBlockSequence::new_with_mm(
30            Tokens::from(self.tokens.clone()),
31            &token_mm,
32            block_size,
33            Some(salt_hash),
34        )?;
35        Ok(seq.blocks().iter().map(UniversalBlock::from).collect())
36    }
37
38    fn into_blocks_consuming(self, block_size: u32) -> Result<Vec<UniversalBlock>, KvHashingError> {
39        let salt_hash = compute_salt_hash(self.salt(), self.lora_name())?;
40        let token_mm = self.mm_info.into_iter().map(Into::into).collect::<Vec<_>>();
41        let seq = TokenBlockSequence::new_with_mm(
42            Tokens::from(self.tokens),
43            &token_mm,
44            block_size,
45            Some(salt_hash),
46        )?;
47        Ok(seq.blocks().iter().map(UniversalBlock::from).collect())
48    }
49
50    /// Projection: per-block [`BlockHash`].
51    pub fn block_hashes(&self, block_size: u32) -> Result<Vec<BlockHash>, KvHashingError> {
52        Ok(self
53            .into_blocks(block_size)?
54            .into_iter()
55            .map(|b| b.block_hash)
56            .collect())
57    }
58
59    /// Projection: per-block [`SequenceHash`] (parent-chained, derived from PLH).
60    pub fn sequence_hashes(&self, block_size: u32) -> Result<Vec<SequenceHash>, KvHashingError> {
61        Ok(self
62            .into_blocks(block_size)?
63            .into_iter()
64            .map(|b| b.sequence_hash())
65            .collect())
66    }
67
68    /// Consuming projection: per-block [`SequenceHash`].
69    ///
70    /// This preserves the borrowed [`Self::sequence_hashes`] API for callers that need to
71    /// keep the request, while allowing one-shot producers to move the token vector into
72    /// block construction and avoid an extra full-prompt clone.
73    pub fn into_sequence_hashes(
74        self,
75        block_size: u32,
76    ) -> Result<Vec<SequenceHash>, KvHashingError> {
77        Ok(self
78            .into_blocks_consuming(block_size)?
79            .into_iter()
80            .map(|b| b.sequence_hash())
81            .collect())
82    }
83
84    /// Projection: per-block [`PositionalLineageHash`] (the universal identifier).
85    pub fn positional_lineage_hashes(
86        &self,
87        block_size: u32,
88    ) -> Result<Vec<PositionalLineageHash>, KvHashingError> {
89        Ok(self
90            .into_blocks(block_size)?
91            .into_iter()
92            .map(|b| b.plh)
93            .collect())
94    }
95}