diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 50099695120..3917bbba720 100644 --- a/megatron/core/inference/inference_request.py +++ b/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 diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 9fc9ac9aac3..c2deb6d74c7 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/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()