fix: Replace polynomial rolling hash with SHA-256 for prefix caching

megatron/core/inference/inference_request.py

@@ -1,11 +1,11 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
+import hashlib
 import time
 import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum, auto
-from itertools import accumulate
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
@@ -87,53 +87,44 @@ class Status(Enum):
 # Hash computation for prefix caching
 # =========================================================================
 
-# Constants for hash computation
-# Using 2^61 - 1 (Mersenne prime) for ~10^18 hash space, reducing collision probability
-# from ~10^-9 to ~10^-18 compared to the previous prime (1000000007).
-HASH_PRIME = 2305843009213693951
-HASH_BASE = 31
-
-_hash_powers: Optional[torch.Tensor] = None
-
 
 def compute_block_hashes_batched(prompt_tokens: torch.Tensor, block_size: int) -> List[int]:
-    """Compute hashes for all complete blocks in a prompt in one batched operation.
+    """Compute SHA-256 based hashes for all complete blocks in a prompt.
 
-    Reshapes prompt tokens into [num_blocks, block_size], computes all per-block
-    token hashes via a single GPU matmul, transfers results with one .tolist() call,
-    and chains parent hashes on CPU.
+    Each block hash is computed as SHA-256(parent_digest || block_bytes), where
+    parent_digest chains from the previous block (starting from a zero digest).
+    This provides cryptographic collision resistance with no exploitable algebraic
+    structure.
 
     Args:
         prompt_tokens: All prompt token IDs, shape [seq_len].
         block_size: Number of tokens per block.
 
     Returns:
-        List of positive integer hash values (1 to HASH_PRIME), one per complete block.
+        List of positive integer hash values in [1, 2^63-1], one per complete block.
     """
     num_complete_blocks = len(prompt_tokens) // block_size
     if num_complete_blocks == 0:
         return []
 
-    global _hash_powers
-    if _hash_powers is None or _hash_powers.shape[0] != block_size:
-        positions = torch.arange(block_size, device=prompt_tokens.device, dtype=torch.int64)
-        _hash_powers = torch.pow(HASH_BASE, positions).to(torch.int64) % HASH_PRIME
+    # Single GPU->CPU transfer, get contiguous bytes
+    tokens_cpu = prompt_tokens[: num_complete_blocks * block_size].to(torch.int64).cpu()
+    tokens_bytes = tokens_cpu.numpy().tobytes()
+    block_byte_size = block_size * tokens_cpu.element_size()  # 8 bytes per int64
 
-    # Reshape to [num_blocks, block_size] (zero-copy view) and compute all token hashes
-    blocks = prompt_tokens[: num_complete_blocks * block_size].view(num_complete_blocks, block_size)
-    token_hashes = (blocks.to(torch.int64) * _hash_powers).sum(dim=1) % HASH_PRIME
+    hashes = []
+    parent_digest = b'\x00' * 32  # SHA-256 digest size
 
-    # Single GPU→CPU transfer
-    token_hashes_list = token_hashes.tolist()
+    for i in range(num_complete_blocks):
+        block_bytes = tokens_bytes[i * block_byte_size : (i + 1) * block_byte_size]
+        digest = hashlib.sha256(parent_digest + block_bytes).digest()
 
-    # Chain parent hashes on CPU (C-level accumulate, no Python loop)
-    hashes = list(
-        accumulate(
-            token_hashes_list,
-            lambda parent, th: (parent * HASH_BASE + th) % HASH_PRIME + 1,
-            initial=0,
-        )
-    )[1:]
+        # Map to positive int64 range [1, 2^63-1], avoiding sentinels -1 and 0
+        raw = int.from_bytes(digest[:8], byteorder='little', signed=False)
+        hash_val = (raw % (2**63 - 1)) + 1
+
+        hashes.append(hash_val)
+        parent_digest = digest  # Full 32-byte digest chains into next block
 
     return hashes
 

tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py

@@ -11,7 +11,6 @@
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
 from megatron.core.inference.inference_request import (
-    HASH_PRIME,
     DynamicInferenceRequest,
     DynamicInferenceRequestRecord,
     Status,
@@ -142,7 +141,7 @@ def test_hash_computation(self):
         tokens = self._prompt(32)
         h1 = compute_block_hashes_batched(tokens, 32)
         h2 = compute_block_hashes_batched(tokens, 32)
-        assert h1 == h2 and len(h1) == 1 and 1 <= h1[0] <= HASH_PRIME
+        assert h1 == h2 and len(h1) == 1 and h1[0] >= 1
         assert compute_block_hashes_batched(self._prompt(32, offset=1), 32)[0] != h1[0]
 
         # parent chaining: 4 blocks of all-zero tokens produce distinct hashes
@@ -161,6 +160,47 @@ def test_hash_computation(self):
         )
         assert len(long_h) == 120 and all(v > 0 for v in long_h)
 
+    @pytest.mark.internal
+    def test_hash_collision_resistance(self):
+        """Regression tests: old polynomial collision attacks must fail with SHA-256."""
+        bs = 32
+
+        # V2 regression: algebraic attack (token[j] += 31, token[j+1] -= 1)
+        # This was a zero-delta exploit against the old polynomial hash.
+        tokens = self._prompt(bs)
+        collision = tokens.clone()
+        collision[0] += 31
+        collision[1] -= 1
+        h_orig = compute_block_hashes_batched(tokens, bs)
+        h_coll = compute_block_hashes_batched(collision, bs)
+        assert h_orig != h_coll, "V2 algebraic collision: token[j]+=31, token[j+1]-=1"
+
+        # V2 at different positions within the block
+        for j in range(bs - 1):
+            c = tokens.clone()
+            c[j] += 31
+            c[j + 1] -= 1
+            assert compute_block_hashes_batched(c, bs) != h_orig, f"V2 at position {j}"
+
+        # V2 across multiple blocks: modify one block, verify all downstream hashes change
+        tokens_multi = self._prompt(bs * 4)
+        h_multi = compute_block_hashes_batched(tokens_multi, bs)
+        modified = tokens_multi.clone()
+        modified[0] += 31
+        modified[1] -= 1
+        h_mod = compute_block_hashes_batched(modified, bs)
+        assert h_mod[0] != h_multi[0], "modified block hash must differ"
+        # Parent chaining: all subsequent blocks must also differ
+        for i in range(1, 4):
+            assert h_mod[i] != h_multi[i], f"parent chain: block {i} must differ"
+
+        # V2 generalized: arbitrary linear combinations (token[j] += k*31, token[j+1] -= k)
+        for k in [1, 2, 5, 100]:
+            c = tokens.clone()
+            c[0] += k * 31
+            c[1] -= k
+            assert compute_block_hashes_batched(c, bs) != h_orig, f"V2 generalized k={k}"
+
     @pytest.mark.internal
     def test_registration_and_discovery(self):
         ctx = self._ctx()