# Copyright 2025 The HuggingFace Team Inc. # # 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 clone 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. import copy import functools import gc import itertools import math import os import threading import time import unittest from types import SimpleNamespace from typing import Any from unittest.mock import patch import torch from parameterized import parameterized import transformers.modeling_flash_attention_utils as fa_utils from transformers import ( AutoConfig, AutoModelForCausalLM, AutoTokenizer, CompileConfig, ContinuousBatchingConfig, GenerationConfig, GenerationMixin, StaticCache, ) from transformers.generation.continuous_batching.cache import ( PagedAttentionCache, PagedAttentionMemoryHandler, group_layers_by_attn_type, ) from transformers.generation.continuous_batching.cache_allocators import ( CacheAllocator, CachePool, FullAttentionCacheAllocator, SlidingAttentionCacheAllocator, ) from transformers.generation.continuous_batching.continuous_api import ( BackgroundThreadStatus, ContinuousBatchingManager, OutputRouter, ) from transformers.generation.continuous_batching.distributed import DistributedHelper from transformers.generation.continuous_batching.input_outputs import build_attention_mask from transformers.generation.continuous_batching.offloading_manager import OffloadingManager from transformers.generation.continuous_batching.requests import ( GenerationOutput, RequestState, RequestStatus, get_device_and_memory_breakdown, ) from transformers.generation.continuous_batching.utils import DEVICE_TYPE_TO_GRAPH_NAME from transformers.integrations.eager_paged import eager_paged_attention_forward from transformers.testing_utils import ( Expectations, backend_empty_cache, backend_memory_allocated, require_deterministic_for_xpu, require_flash_attn, require_flash_attn_3, require_torch_accelerator, require_torch_multi_accelerator, slow, torch_device, ) from transformers.utils import ( is_flash_attn_2_available, is_flash_attn_3_available, is_torch_cuda_available, is_torch_mps_available, is_torch_xpu_available, ) from transformers.utils.generic import is_flash_attention_requested from ..test_tensor_parallel_mixin import _init_distributed # Constants for tests _DEFAULT_USER_MESSAGES = [ "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?", "Josh decides to try flipping a house. He buys a house for $80,000 and then puts in $50,000 in repairs. This increased the value of the house by 150%. How much profit did he make?", "A basket contains 25 oranges among which 1 is bad, 20% are unripe, 2 are sour and the rest are good. How many oranges are good?", ] # fmt: skip # Helper functions def flush_memory(flush_compile: bool = True) -> None: """Flushes the memory of the current device and, if the flush_compile flag is True, all data related to torch.compile.""" gc.collect() # If needed, flush everything related to torch.compile if flush_compile: # Dynamo resets torch._dynamo.reset() torch._dynamo.reset_code_caches() if hasattr(torch._inductor, "codecache"): # Clear FX graph cache if hasattr(torch._inductor.codecache, "FxGraphCache"): torch._inductor.codecache.FxGraphCache.clear() # Clear PyCodeCache if hasattr(torch._inductor.codecache, "PyCodeCache"): torch._inductor.codecache.PyCodeCache.cache_clear() # Clear TritonFuture cache (for async compilation) if hasattr(torch._inductor.codecache, "TritonFuture"): if hasattr(torch._inductor.codecache.TritonFuture, "_compile_cache"): torch._inductor.codecache.TritonFuture._compile_cache.clear() # Clear CUDA cache if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() elif torch.xpu.is_available(): torch.xpu.empty_cache() torch.xpu.synchronize() gc.collect() def get_tokenizer_and_model( model_id: str, attn_implementation: str, device: str, dtype: str | torch.dtype = "auto", upcast_lm_head: bool = False, ) -> tuple[AutoTokenizer, GenerationMixin]: """Returns a tokenizer and a model for the given model ID. Attributes to setup the models (attn_implementation, dtype and device) are needed as arguments.""" # Get tokenizer, with a padding token if not present tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") if not hasattr(tokenizer, "pad_token") and hasattr(tokenizer, "eos_token"): tokenizer.pad_token = tokenizer.eos_token # Load model model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=attn_implementation, torch_dtype=dtype) model = model.to(device).eval() # If needed, upcast the lm_head to fp32 for added precision. This helps break ties in bf16 that would lead to token # divergence between CB and generate, while not affecting the rest of the model: if there is a real divergence in # the model, it will accumulate over layers, and having a more precise LM head will not close the gap. if upcast_lm_head: model.lm_head = copy.deepcopy(model.lm_head).to(torch.float32) # copy in case embedding are tied model.lm_head.register_forward_pre_hook(lambda m, args: (args[0].float(), *args[1:])) return tokenizer, model def with_flush_memory(func): """Decorator that ensures flush_memory is called after the test, even if it fails.""" @functools.wraps(func) def wrapper(*args, **kwargs): # Determine flush_compile value from continuous_batching_config or generation_config cb_config = kwargs.get("continuous_batching_config") generation_config = kwargs.get("generation_config") if isinstance(cb_config, ContinuousBatchingConfig): flush_compile = ( cb_config.default_compile_level > 0 or cb_config.varlen_compile_config is not None or cb_config.decode_compile_config is not None ) elif isinstance(generation_config, GenerationConfig): flush_compile = generation_config.compile_config is not None else: flush_compile = False # Run the test and always flush memory try: return func(*args, **kwargs) finally: flush_memory(flush_compile=flush_compile) return wrapper def flexible_flash_skip(test_case: unittest.TestCase, attn_implementation: str) -> None: """Skip the test if Flash Attention 2 or 3 is required but not available.""" is_fa2 = is_flash_attention_requested(requested_attention_implementation=attn_implementation, version=2) if is_fa2 and not is_flash_attn_2_available(kernels_fallback_ok=True): test_case.skipTest("Flash Attention 2 is not available, as a package or through `kernels`. Skipping test.") is_fa3 = is_flash_attention_requested(requested_attention_implementation=attn_implementation, version=3) if is_fa3 and not is_flash_attn_3_available(kernels_fallback_ok=True): test_case.skipTest("Flash Attention 3 is not available, as a package or through `kernels`. Skipping test.") def get_generation_inputs( user_messages: list[str], tokenizer: AutoTokenizer, for_continuous_batching: bool = False ) -> Any: """Returns the tokenized inputs for batched or non-batched generation.""" chats = [[{"role": "user", "content": user_message}] for user_message in user_messages] if for_continuous_batching: tokenized = [tokenizer.apply_chat_template(chat, add_generation_prompt=True) for chat in chats] input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized] return input_ids else: inputs = tokenizer.apply_chat_template( chats, add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True, return_attention_mask=True, ) return inputs def regular_generate( model: GenerationMixin, tokenizer: AutoTokenizer, user_messages: list[str], **generate_kwargs, ) -> tuple[list[list[int]], list[list[float]]]: # Run generation inputs = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=False) generate_outputs = model.generate(**inputs.to(model.device), return_dict_in_generate=True, **generate_kwargs) # Keep only generated tokens all_generated_tokens = [] num_input_tokens = inputs.input_ids.shape[1] for i in range(len(user_messages)): # Remove left-side input and padding tokens generated_toks = generate_outputs.sequences[i, num_input_tokens:].tolist() # Remove right-side padding tokens while generated_toks[-1] == model.generation_config.pad_token_id: generated_toks.pop() all_generated_tokens.append(generated_toks) # Retrieve logprobs if the scores were requested per_prompt_logprobs = [] if generate_kwargs.get("output_scores", False): # Loop over prompts for i in range(len(user_messages)): logprobs = [] tokens_for_prompt = generate_outputs.sequences[i, num_input_tokens:].tolist() for score, token in zip(generate_outputs.scores, tokens_for_prompt): # Scores already have logits processors applied (including temperature) probs = torch.nn.functional.softmax(score[i], dim=-1) logprobs.append(probs[token].log().item()) per_prompt_logprobs.append(logprobs) # Otherwise, return an empty list else: per_prompt_logprobs = [] return all_generated_tokens, per_prompt_logprobs def _setup_cache_pool(allocators: list[CacheAllocator], num_sectors: int) -> CachePool: """Computes the sector geometry the same way as PagedAttentionCache, then builds the cache tensor and the pool and registers both on the given allocators.""" bytes_per_sector = math.lcm(*(allocator.bytes_per_block for allocator in allocators), 128) non_trash_bytes = num_sectors * bytes_per_sector cache_tensor = torch.zeros(non_trash_bytes + 2 * bytes_per_sector, dtype=torch.uint8) pool = CachePool(num_sectors=num_sectors, num_allocators=len(allocators)) for allocator in allocators: allocator.register_cache_tensor(bytes_per_sector, non_trash_bytes, cache_tensor, pool) return pool def _make_allocator( cls: type[FullAttentionCacheAllocator | SlidingAttentionCacheAllocator], head_dim: int, num_kv_heads: int, page_size: int, layer_indices: list[int] | None = None, index: int = 0, sliding_window: int | None = None, allow_block_sharing: bool = True, ) -> FullAttentionCacheAllocator | SlidingAttentionCacheAllocator: """Builds an allocator on a minimal namespace config: allocators only read head_dim, num_key_value_heads and sliding_window from the model config, so tests do not need a full PreTrainedConfig.""" config = SimpleNamespace(head_dim=head_dim, num_key_value_heads=num_kv_heads, sliding_window=sliding_window) return cls( index=index, config=config, # type: ignore num_key_value_heads=num_kv_heads, cache_dtype=torch.float16, page_size=page_size, layer_indices=layer_indices if layer_indices is not None else [0], allow_block_sharing=allow_block_sharing, ) # Class for all continuous batching tests that do not require any accelerator. Usualy those test are faster to run. class ContinuousBatchingNoAcceleratorTest(unittest.TestCase): @parameterized.expand([("paged|eager", eager_paged_attention_forward)]) def test_paged_forward_without_cache_raises(self, attn_implementation, attention_forward): # A standard forward on a model switched to a paged implementation reaches these with no cache. They are # written for the packed inputs and the 4D mask continuous batching prepares, so they would silently attend # bidirectionally instead of causally: refuse the call rather than return a wrong result. module = torch.nn.Module() module.layer_idx = 0 module.num_key_value_groups = 1 q = k = v = torch.zeros(1, 1, 4, 8) with self.assertRaisesRegex(ValueError, attn_implementation.replace("|", r"\|")): attention_forward(module, q, k, v, None, scaling=1.0) def test_generation_outputs_are_snapshots(self): state = RequestState( request_id="r", initial_tokens=[10, 11], max_new_tokens=3, streaming=True, record_timestamps=True ) state._status = RequestStatus.DECODING self.assertFalse(state.update_and_check_completion(101, -1.01)) first_output = state.to_generation_output() self.assertEqual(first_output.generated_tokens, [101]) self.assertEqual(first_output.logprobs, [-1.01]) self.assertEqual(len(first_output.timestamps), 1) self.assertFalse(state.update_and_check_completion(102, -1.02)) second_output = state.to_generation_output() self.assertEqual(first_output.generated_tokens, [101]) self.assertEqual(first_output.logprobs, [-1.01]) self.assertEqual(len(first_output.timestamps), 1) self.assertEqual(second_output.generated_tokens, [101, 102]) self.assertEqual(second_output.logprobs, [-1.01, -1.02]) self.assertEqual(len(second_output.timestamps), 2) def test_streaming_output_after_soft_reset_does_not_shorten_generation(self): state = RequestState(request_id="r", initial_tokens=[10, 11], max_new_tokens=5, streaming=True) state._status = RequestStatus.DECODING for token_id in [101, 102]: self.assertFalse(state.update_and_check_completion(token_id, None)) reset_state = state.create_equivalent_initial_request() reset_state._status = RequestStatus.DECODING reset_state.streaming = True accepted_tokens = [] streamed_tokens = [] for token_id in [103, 104, 105, 106]: previous_tokens = reset_state.generated_tokens[:] is_finished = reset_state.update_and_check_completion(token_id, None) if reset_state.generated_tokens != previous_tokens: accepted_tokens.append(token_id) streamed_tokens.append(reset_state.to_generation_output().generated_tokens) if is_finished: break self.assertEqual(accepted_tokens, [103, 104, 105]) self.assertEqual(streamed_tokens, [[101, 102, 103], [101, 102, 103, 104], [101, 102, 103, 104, 105]]) self.assertEqual(reset_state.generated_tokens, [103, 104, 105]) self.assertEqual(reset_state.initial_tokens, [10, 11, 101, 102]) @parameterized.expand( [ ("f", None, {"f": [0]}), ("ffff", None, {"f": [0, 1, 2, 3]}), ("sssss", 4096, {"s": [0, 1, 2, 3, 4]}), ("fs", 4096, {"f": [0], "s": [1]}), ("ssfssf", 4096, {"f": [2, 5], "s": [0, 1, 3, 4]}), ("ssssf", 4096, {"f": [4], "s": [0, 1, 2, 3]}), ("fffsffs", 4096, {"f": [0, 1, 2, 4, 5], "s": [3, 6]}), # No layer_types attribute: falls back on the presence of a sliding window (None, None, {"f": [0, 1, 2, 3]}), (None, 4096, {"s": [0, 1, 2, 3]}), ] ) def test_group_layers( self, layer_types_str: str | None, sliding_window: int | None, expected_groups: dict[str, list[int]], ) -> None: """Test the layer grouping: one group per attention type, holding the indices of its layers in order.""" letter_to_type = {"f": "full_attention", "s": "sliding_attention"} # Take a config and change the layer_types attribute to the mix we want config = AutoConfig.from_pretrained("HuggingFaceTB/SmolLM-1.7B") config.sliding_window = sliding_window if layer_types_str is not None: config.layer_types = [letter_to_type[char] for char in layer_types_str] else: config.num_hidden_layers = 4 # Expand the expected groups expected_groups = {letter_to_type[key]: value for key, value in expected_groups.items()} layer_groups = group_layers_by_attn_type(config) self.assertEqual( layer_groups, expected_groups, f"Test failed for: {layer_types_str = }, {sliding_window = }, {expected_groups = }, {layer_groups = }", ) @parameterized.expand( [ ([0, 4], [0, 4], 1, ["1000", "1100", "1110", "1111"]), ([0, 4], [0, 4], 2, ["1000", "1100", "0110", "0011"]), ([0, 3], [0, 5], 1, ["11100", "11110", "11111"]), ([0, 3], [0, 5], 3, ["11100", "01110", "00111"]), ([0, 3, 6], [0, 3, 6], 1, ["100000", "110000", "111000", "000100", "000110", "000111"]), ([0, 3, 6], [0, 3, 6], 2, ["100000", "110000", "011000", "000100", "000110", "000011"]), ] ) def test_attention_mask( self, cumulative_seqlens_q: list[int], cumulative_seqlens_k: list[int], sliding_window: int, # the sliding window size, 1 means no sliding window str_expected_mask_lines: list[str], # the attention mask, broken down by line as a string of 0s and 1s ) -> None: """Tests the correctness of the attention mask used in the continuous batching API.""" # Build expected mask minus_inf = torch.finfo(torch.float32).min expected_mask = torch.empty((cumulative_seqlens_q[-1], cumulative_seqlens_k[-1]), dtype=torch.float32) for i, line in enumerate(str_expected_mask_lines): expected_mask[i, :] = torch.tensor([minus_inf if c == "0" else 0 for c in line]) # Build actual mask actual_mask = torch.full_like(expected_mask, minus_inf) # function modifies in place build_attention_mask(actual_mask, cumulative_seqlens_q, cumulative_seqlens_k, sliding_window) # Check that the actual mask matches the expected mask matches = (expected_mask == actual_mask).all() # If it doesn't match, print the masks in a readable form and fail the test if not matches: str_mask = [ "".join("1" if x == 0 else "0" for x in token_attn_vector) for token_attn_vector in actual_mask ] str_mask = "\n".join(str_mask) str_expected_mask = "\n".join(str_expected_mask_lines) self.fail( f"Test failed for: {cumulative_seqlens_q = }, {cumulative_seqlens_k = }, {sliding_window = }\n" f"Expected mask:\n{str_expected_mask}\n" f"Actual mask:\n{str_mask}" ) def test_continuous_batching_can_store_request_tokens(self) -> None: """Tests the allocation check on a hybrid full + sliding cache: the sliding allocator caps its need at the window, zero new tokens always fit, and a dry run leaves the pool untouched while a real allocation consumes it.""" def make_cache(num_sectors: int) -> PagedAttentionCache: common: dict[str, Any] = {"head_dim": 2, "num_kv_heads": 1, "page_size": 4, "allow_block_sharing": False} full = _make_allocator(FullAttentionCacheAllocator, **common) sliding = _make_allocator( SlidingAttentionCacheAllocator, layer_indices=[1], index=1, sliding_window=8, **common ) # bytes_per_block = 32 for both allocators, so each sector of lcm(32, 128) = 128 bytes holds 4 blocks pool = _setup_cache_pool([full, sliding], num_sectors) cache = PagedAttentionCache.__new__(PagedAttentionCache) cache.pool = pool cache.cache_allocators = {"full_attention": full, "sliding_attention": sliding} return cache state = RequestState(request_id="req_0", initial_tokens=[0]) # Each sector holds 4 pages of 4 tokens. 32 new tokens need 2 full-attention sectors but only 1 sliding # sector, since the sliding need is capped at the window (8 tokens = 2 pages): 3 sectors in total cache = make_cache(num_sectors=3) self.assertTrue(cache.can_store_request_tokens(state, 32, dry_run=True)) self.assertEqual(cache.pool.num_free_sectors, 3) # a dry run does not allocate self.assertFalse(make_cache(num_sectors=2).can_store_request_tokens(state, 32, dry_run=True)) # Zero new tokens always fit, even with no free sectors empty_cache = make_cache(num_sectors=3) for _ in range(3): empty_cache.pool.allocate_sector(0) self.assertEqual(empty_cache.pool.num_free_sectors, 0) self.assertTrue(empty_cache.can_store_request_tokens(state, 0, dry_run=True)) # A real allocation consumes the sectors cache = make_cache(num_sectors=3) self.assertTrue(cache.can_store_request_tokens(state, 32)) self.assertEqual(cache.pool.num_free_sectors, 0) # Growing the now window-full request only costs full-attention blocks (1 fresh sector the pool cannot serve) state.position_offset = 32 self.assertFalse(cache.can_store_request_tokens(state, 16, dry_run=True)) @parameterized.expand( [ # (block_size, block_table, past_length, query_length) # Contiguous blocks (32, [0, 1, 2], 0, 16), (32, [0, 1, 2], 0, 64), (32, [0, 1, 2], 16, 16), (32, [0, 1, 2], 31, 2), # Non-contiguous blocks (32, [0, 3, 6], 0, 64), (32, [2, 5, 8], 60, 10), # Different block sizes (16, [0, 1, 2, 3], 14, 4), (64, [0, 1], 60, 10), ] ) def test_full_attention_get_indices( self, block_size: int, block_table: list[int], past_length: int, query_length: int, ) -> None: """Test FullAttentionCacheAllocator.get_read_indices and get_write_indices return correct physical indices.""" allocator = _make_allocator(FullAttentionCacheAllocator, head_dim=8, num_kv_heads=2, page_size=block_size) allocator.block_table["req"] = block_table def reference_indices(start: int, end: int) -> list[int]: """Reference implementation: converts logical indices to physical row indices. The token slot t of block b lives at row b * block_physical_stride + t.""" return [ block_table[i // block_size] * allocator.block_physical_stride + i % block_size for i in range(start, end) ] # Test read indices (from 0 to past_length + query_length) expected_read = reference_indices(0, past_length + query_length) self.assertEqual(allocator.get_read_indices("req", past_length, query_length), expected_read) # Test write indices (from past_length to past_length + query_length) expected_write = reference_indices(past_length, past_length + query_length) self.assertEqual(allocator.get_write_indices("req", past_length, query_length), expected_write) @parameterized.expand( [ # (block_size, sliding_window, block_table, past_length, query_length) # Prefill from empty: no cache read, only sentinels in the read indices (4, 8, [0, 1], 0, 3), # Decode within the window (4, 8, [0, 1], 5, 1), # Chunked prefill in the middle of the window (4, 8, [0, 1], 2, 4), # Past length beyond the window: rolling-buffer wrap-around (4, 8, [0, 1], 10, 1), (4, 8, [0, 1], 14, 3), # Non-contiguous blocks (4, 8, [3, 5], 6, 2), # Query longer than the window: write indices get left-padded with the write trash index (4, 8, [0, 1], 0, 10), # Single-block window (block_size == sliding_window) (4, 4, [7], 3, 1), # Larger block size (16, 32, [0, 1], 20, 4), ] ) def test_sliding_attention_get_indices( self, block_size: int, sliding_window: int, block_table: list[int], past_length: int, query_length: int, ) -> None: """Test SlidingAttentionCacheAllocator.get_read_indices and get_write_indices place the cache, sentinel and write trash indices correctly, including for small block sizes and rolling-buffer wrap-around.""" allocator = _make_allocator( SlidingAttentionCacheAllocator, head_dim=8, num_kv_heads=2, page_size=block_size, sliding_window=sliding_window, ) allocator.block_table["req"] = block_table # The special indices live in the trash zone above the allocatable blocks (see PagedAttentionCache). We pick a # num_blocks larger than any block id used here so they never collide with real cache positions. num_blocks = 64 self.assertTrue(all(b < num_blocks for b in block_table)) allocator.read_trash_index = num_blocks * allocator.block_physical_stride allocator.sentinel_index = sentinel_index = num_blocks * allocator.block_physical_stride + 1 allocator.write_trash_index = write_trash_index = (num_blocks + 1) * allocator.block_physical_stride def to_physical(position: int) -> int: """Reference position-to-physical-row mapping inside the rolling buffer.""" slot = position % sliding_window return block_table[slot // block_size] * allocator.block_physical_stride + slot % block_size # Read indices: the last min(past_length, sliding_window - 1) positions, then one sentinel per query token read_cache_length = min(past_length, sliding_window - 1) read_start = past_length - read_cache_length expected_read = [to_physical(p) for p in range(read_start, read_start + read_cache_length)] expected_read += [sentinel_index] * query_length read = allocator.get_read_indices("req", past_length, query_length) # Main check self.assertEqual(read, expected_read) # No sentinel in the real indices self.assertNotIn(sentinel_index, read[:read_cache_length]) # Cache reads land in allocated blocks for idx in read[:read_cache_length]: self.assertIn(idx // allocator.block_physical_stride, block_table) # Write indices: the last min(query_length, sliding_window) positions are stored in their slots, left-padded # with the write trash index when the query overflows the window write_cache_length = min(query_length, sliding_window) padding_length = query_length - write_cache_length read_start = past_length + padding_length expected_write = [write_trash_index] * padding_length expected_write += [to_physical(p) for p in range(read_start, read_start + write_cache_length)] write = allocator.get_write_indices("req", past_length, query_length) # Main check self.assertEqual(write, expected_write) # Trash only at the front self.assertNotIn(write_trash_index, write[padding_length:]) # Cache writes land in allocated blocks for idx in write[padding_length:]: self.assertIn(idx // allocator.block_physical_stride, block_table) def test_fork_blocks_sharing_and_copy(self) -> None: """Tests forking at the allocator level: fully-written blocks are shared through a reference count, the others are backed by fresh blocks and copied, and a shared block returns to the pool only when its last owner frees.""" ca_kwargs: dict[str, Any] = {"head_dim": 8, "num_kv_heads": 2, "page_size": 4} full = _make_allocator(FullAttentionCacheAllocator, layer_indices=[0, 1], **ca_kwargs) sliding = _make_allocator( SlidingAttentionCacheAllocator, layer_indices=[2], index=1, sliding_window=8, **ca_kwargs ) num_sectors = 6 pool = _setup_cache_pool([full, sliding], num_sectors) # Allocate a source request with 10 cached tokens: 3 full-attention blocks, 2 sliding ones (window of 8) past_length = 10 for allocator in (full, sliding): for _ in range(allocator.needs_new_sectors("src", 0, past_length)): pool.allocate_sector(allocator.index) allocator.allocate_cache_to_request("src", 0, past_length) self.assertEqual(len(full.block_table["src"]), 3) self.assertEqual(len(sliding.block_table["src"]), 2) # Only the 2 fully-written full-attention blocks are shareable, and only when sharing is in use: the sliding # allocator never shares (use_block_sharing resolves to False despite allow_block_sharing=True) self.assertEqual(full.count_non_shareable_blocks("src", past_length), 1) self.assertEqual(sliding.count_non_shareable_blocks("src", past_length), 2) full.use_block_sharing = False self.assertEqual(full.count_non_shareable_blocks("src", past_length), 3) full.use_block_sharing = True # Stamp each source block with a distinct byte pattern, then fork a child request for allocator in (full, sliding): for block_id in allocator.block_table["src"]: allocator._copy_view[block_id] = (allocator.index * 100 + block_id) % 256 for allocator, num_copied in ((full, 1), (sliding, 2)): missing_blocks = num_copied - pool.count_free_blocks(allocator.index) for _ in range((missing_blocks + allocator.blocks_per_sector - 1) // allocator.blocks_per_sector): pool.allocate_sector(allocator.index) src_blocks, dst_blocks = allocator.fork_blocks("src", "child", past_length) self.assertEqual(len(src_blocks), num_copied) allocator.copy_blocks(src_blocks, dst_blocks) # The child's table starts with the shared blocks, and the copied content matches the source num_shared = len(allocator.block_table["src"]) - num_copied self.assertEqual(allocator.block_table["child"][:num_shared], allocator.block_table["src"][:num_shared]) for src_block, dst_block in zip(src_blocks, dst_blocks): self.assertTrue(torch.equal(allocator._copy_view[dst_block], allocator._copy_view[src_block])) # Shared blocks are ref-counted: freeing the source keeps them alive for the child shared_blocks = full.block_table["src"][:2] self.assertEqual(full.ledger.shared_ref_counts, dict.fromkeys(shared_blocks, 2)) self.assertEqual(sliding.ledger.shared_ref_counts, {}) # Sliding never shares: the child's table is fully disjoint from the source's self.assertTrue(set(sliding.block_table["child"]).isdisjoint(sliding.block_table["src"])) full.free_blocks("src") sliding.free_blocks("src") self.assertEqual(full.ledger.shared_ref_counts, {}) # counts of 1 are implicit self.assertEqual(full.block_table["child"][:2], shared_blocks) self.assertEqual(pool.count_free_blocks(full.index), 1) # only the copied-from block was freed self.assertEqual(pool.count_free_blocks(sliding.index), 2) # Freeing the child releases everything: after reclaim, the pool is whole again full.free_blocks("child") sliding.free_blocks("child") pool.try_to_free_sectors() self.assertEqual(pool.num_free_sectors, num_sectors) def test_prefix_match_and_dedup(self) -> None: """Tests block de-duplication at the allocator level: completed blocks are hashed, later requests match them as a prefix, unreferenced hashed blocks are kept cached instead of freed, and eviction returns them to the pool.""" allocator = _make_allocator(FullAttentionCacheAllocator, head_dim=8, num_kv_heads=2, page_size=4) num_sectors = 8 pool = _setup_cache_pool([allocator], num_sectors) ledger = allocator.ledger def allocate(request_id: str, num_tokens: int) -> None: for _ in range(allocator.needs_new_sectors(request_id, 0, num_tokens)): pool.allocate_sector(0) allocator.allocate_cache_to_request(request_id, 0, num_tokens) # "src" caches 10 tokens over 3 blocks, of which the first 2 are complete (page_size = 4) tokens = list(range(100, 110)) allocate("src", 10) src_table = allocator.block_table["src"][:] allocator.mark_complete_blocks("src", tokens, new_new_blocks=2) self.assertEqual(len(ledger.hash_to_block), 2) self.assertEqual(set(ledger.block_to_hash), set(src_table[:2])) # A request computing the same content is de-duplicated on the spot: it adopts the existing blocks, its # duplicates are freed immediately, and no new hash is registered allocate("dup", 10) allocator.mark_complete_blocks("dup", tokens, new_new_blocks=2) self.assertEqual(len(ledger.hash_to_block), 2) self.assertEqual(allocator.block_table["dup"][:2], src_table[:2]) self.assertEqual(pool.count_free_blocks(0), 2) # the two duplicate blocks were freed by the swap self.assertEqual(ledger.shared_ref_counts, dict.fromkeys(src_table[:2], 2)) allocator.free_blocks("dup") # releases the adopted blocks and frees its own incomplete one self.assertEqual(pool.count_free_blocks(0), 3) def match_and_acquire(request_id: str, prompt: list[int]) -> int: """Mirrors PagedAttentionCache.search_prefix_match at the single-allocator level.""" matched_blocks = allocator.match_prefix_blocks(prompt) prefix_len = len(matched_blocks) * allocator.tokens_per_page allocator.acquire_prefix_blocks(request_id, prefix_len, matched_blocks) return prefix_len # Prefix matching: same 8 first tokens match 2 blocks, 5 matching tokens only cover 1 complete block, and a # prompt that is an exact multiple of the page size never matches its own last block self.assertEqual(match_and_acquire("same", tokens[:8] + [1, 2, 3]), 8) self.assertEqual(allocator.block_table["same"], src_table[:2]) self.assertEqual(match_and_acquire("partial", tokens[:5] + [9, 9, 9]), 4) self.assertEqual(match_and_acquire("exact", tokens[:8]), 4) self.assertEqual(ledger.shared_ref_counts, {src_table[0]: 4, src_table[1]: 2}) # Freeing every request moves the hashed blocks to the cached set, not to the pool for request_id in ("src", "same", "partial", "exact"): allocator.free_blocks(request_id) self.assertEqual(set(ledger.cached_blocks), set(src_table[:2])) self.assertEqual(pool.count_free_blocks(0), 4) # only the incomplete block of "src" was freed # A new match claims the cached blocks back, and freeing it caches them again self.assertEqual(match_and_acquire("revive", tokens[:8] + [7]), 8) self.assertEqual(ledger.cached_blocks, {}) self.assertEqual(ledger.shared_ref_counts, {}) # single owner stays implicit allocator.free_blocks("revive") # Eviction drops the hashes and releases the blocks: the pool ends up whole again evicted_blocks = ledger.evict_cached_blocks() self.assertEqual(set(evicted_blocks), set(src_table[:2])) self.assertEqual(ledger.hash_to_block, {}) pool.free_blocks(0, evicted_blocks) pool.try_to_free_sectors() self.assertEqual(pool.num_free_sectors, num_sectors) def test_count_storable_requests(self) -> None: """Tests the cumulative fit simulation: when a request opens a fresh sector, the sector's unused blocks must be credited to the following requests, which a per-request or global-sum check would miss.""" allocator = _make_allocator( FullAttentionCacheAllocator, head_dim=2, num_kv_heads=1, page_size=4, allow_block_sharing=False ) # bytes_per_block = 32 and bytes_per_sector = lcm(32, 128) = 128, so each sector holds 4 blocks pool = _setup_cache_pool([allocator], num_sectors=2) self.assertEqual(allocator.blocks_per_sector, 4) # Build a bare cache exposing only what the simulation reads cache = PagedAttentionCache.__new__(PagedAttentionCache) cache.pool = pool cache.cache_allocators = {"full_attention": allocator} def blocks_needed(num_blocks: int) -> dict[str, int]: return {"full_attention": num_blocks} # 8 free blocks in 2 sectors: four 2-block requests fit thanks to the leftover blocks of opened sectors # (a per-request ceil would claim 1 sector each and only admit 2), and a fifth request does not fit requests = [blocks_needed(2) for _ in range(5)] self.assertEqual(cache.count_storable_requests(requests[:4]), 4) self.assertEqual(cache.count_storable_requests(requests), 4) # 3-block requests: r1 opens a sector (1 block left), r2 takes it and opens the last sector (2 blocks left), # r3 needs a third sector that does not exist requests = [blocks_needed(3) for _ in range(3)] self.assertEqual(cache.count_storable_requests(requests), 2) # Blocks already free in the allocator's bucket count towards the fit pool.allocate_sector(0) self.assertEqual(cache.count_storable_requests([blocks_needed(4), blocks_needed(4)]), 2) @parameterized.expand( [ # (block_size, sliding_window, block_table, query_lengths) # Prefill shorter than the window, then decode past the wrap-around (4, 8, [0, 1], [3] + [1] * 16), # Prefill exactly the size of the window, then decode (4, 8, [0, 1], [8] + [1] * 16), # Prefill longer than the window, then decode (4, 8, [3, 5], [11] + [1] * 16), # Chunked prefill, with chunks smaller and larger than the window (4, 8, [0, 1], [3, 5, 4, 10, 2] + [1] * 8), # Single-block window (4, 4, [7], [1] * 12), # Larger block size (16, 32, [0, 1], [5] + [1] * 40), ] ) def test_sliding_attention_read_write_round_trip( self, block_size: int, sliding_window: int, block_table: list[int], query_lengths: list[int], ) -> None: """Replays a sequence of prefill/decode steps through the rolling buffer, tracking which logical token each physical slot holds, and checks that every read returns the `sliding_window - 1` most recent tokens in chronological order. Reads and writes are only consistent with each other if both agree on where a given logical position lives, so this catches off-by-one errors that the two index computations would otherwise hide from each other.""" allocator = _make_allocator( SlidingAttentionCacheAllocator, head_dim=8, num_kv_heads=2, page_size=block_size, sliding_window=sliding_window, ) allocator.block_table["req"] = block_table # Special indices are normally set when registering the cache tensor, which we skip here. Physical indices are # always non-negative, so negative sentinels cannot collide with them. allocator.sentinel_index = -1 allocator.write_trash_index = -2 slot_to_token: dict[int, int] = {} # physical slot -> logical position currently stored there past_length = 0 for query_length in query_lengths: # Reads happen before writes, and only when there is cache to read from if past_length > 0: read = allocator.get_read_indices("req", past_length, query_length) self.assertEqual(read[-query_length:], [allocator.sentinel_index] * query_length) cached_tokens = [slot_to_token[idx] for idx in read[:-query_length]] expected_tokens = list(range(max(0, past_length - sliding_window + 1), past_length)) self.assertEqual( cached_tokens, expected_tokens, f"Wrong cache read at {past_length = } for {sliding_window = }, {block_size = }", ) for offset, idx in enumerate(allocator.get_write_indices("req", past_length, query_length)): if idx != allocator.write_trash_index: slot_to_token[idx] = past_length + offset past_length += query_length @slow def test_continuous_batching_no_accelerators(self) -> None: """Test continuous batching generation when no accelerator is available. It uses a simulated CPU-only PyTorch environment by mocking all acceleratoravailability checks to return False""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Accelerator availability are cached, so we clear the cache before mocking to ensure the CPU-only cached_availability_checks = [is_torch_cuda_available, is_torch_mps_available, is_torch_xpu_available] for check in cached_availability_checks: check.cache_clear() # Mock all accelerator availability checks to simulate CPU-only PyTorch try: with ( patch("torch.cuda.is_available", return_value=False), patch("torch.xpu.is_available", return_value=False), patch("torch.backends.mps.is_available", return_value=False), ): # Verify patches work self.assertFalse(torch.cuda.is_available()) self.assertFalse(is_torch_xpu_available()) self.assertFalse(torch.backends.mps.is_available()) tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", "cpu") user_messages = _DEFAULT_USER_MESSAGES[:1] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) model.generation_config.max_new_tokens = 10 model.generation_config.do_sample = False continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=False, use_async_batching=False, auto_switch_to_flash=False ) # This should not crash even with all accelerators unavailable outputs = model.generate_batch( inputs=input_ids, generation_config=model.generation_config, continuous_batching_config=continuous_batching_config, ) # Verify we got outputs self.assertEqual(len(outputs), len(input_ids)) for output in outputs.values(): self.assertIsNotNone(output.generated_tokens) self.assertGreater(len(output.generated_tokens), 0) finally: # Clear the caches again so the mocked availabilities do not leak into subsequent tests. for check in cached_availability_checks: check.cache_clear() def test_output_router_deliver_to_queue(self): """Test that OutputRouter.deliver places outputs on the queue when no handler is registered.""" router = OutputRouter() output = GenerationOutput(request_id="req_0", status=RequestStatus.FINISHED) router.deliver(output) result = router.output_queue.get_nowait() self.assertEqual(result.request_id, "req_0") self.assertTrue(router.output_queue.empty()) def test_output_router_deliver_to_handler(self): """Test that OutputRouter.deliver forwards to a registered handler instead of the queue.""" router = OutputRouter() received = [] loop = unittest.mock.Mock() with router._lock: router.result_handlers["req_0"] = (lambda out: received.append(out), loop) output = GenerationOutput(request_id="req_0", status=RequestStatus.DECODING) router.deliver(output) loop.call_soon_threadsafe.assert_called_once() self.assertTrue(router.output_queue.empty()) def test_distributed_helper_no_dist(self) -> None: """Test that DistributedHelper falls back to a single-rank, TP-driver setup when distributed is not on.""" helper = DistributedHelper(device_mesh=None, cpu_group_timeout=300, tp_plan={}) self.assertFalse(helper.dist_on) self.assertEqual(helper.global_rank, 0) self.assertEqual(helper.world_size, 1) self.assertEqual(helper.tp_size, 1) self.assertEqual(helper.tp_local_rank, 0) self.assertEqual(helper.dp_rank, 0) self.assertEqual(helper.dp_size, 1) self.assertTrue(helper.is_tp_driver) self.assertIsNone(helper.tp_group) self.assertIsNone(helper.cpu_comm_group) # Tensor and object broadcasts should be no-ops without a TP group tensor = torch.tensor([1.0, 2.0]) self.assertTrue(torch.equal(helper.tp_broadcast_from_rank_0(tensor), tensor)) obj = {"some_request": "payload"} self.assertIs(helper.tp_broadcast_object_from_rank_0(obj), obj) # All-reduce-min should be a no-op without a TP group reduce_tensor = torch.tensor([7, 3], dtype=torch.int64) self.assertIs(helper.tp_all_reduce_min(reduce_tensor), reduce_tensor) self.assertTrue(torch.equal(reduce_tensor, torch.tensor([7, 3], dtype=torch.int64))) def test_distributed_helper_set_tp_seed_no_dist(self) -> None: """Test that set_tp_seed sets a torch seed without distributed initialized, both with and without a user seed.""" helper = DistributedHelper(device_mesh=None, cpu_group_timeout=300, tp_plan={}) # Explicit seed: torch RNG state must be reproducible across calls helper.set_tp_seed(seed=42, model_device=torch.device("cpu")) first = torch.randint(0, 2**31 - 1, (4,)) helper.set_tp_seed(seed=42, model_device=torch.device("cpu")) second = torch.randint(0, 2**31 - 1, (4,)) self.assertTrue(torch.equal(first, second)) # No seed: should not raise and should still set a torch seed helper.set_tp_seed(seed=None, model_device=torch.device("cpu")) def test_continuous_batching_config_disables_nccl_graph_mixing(self) -> None: """Test that ContinuousBatchingConfig sets NCCL_GRAPH_MIXING_SUPPORT=0 only under a distributed launch (WORLD_SIZE > 1) and respects the disable_nccl_graph_mixing flag.""" original_nccl = os.environ.pop("NCCL_GRAPH_MIXING_SUPPORT", None) original_ws = os.environ.pop("WORLD_SIZE", None) try: # Single-GPU launch (no WORLD_SIZE): env var is left untouched ContinuousBatchingConfig() self.assertNotIn("NCCL_GRAPH_MIXING_SUPPORT", os.environ) # Distributed launch (WORLD_SIZE > 1): env var is set to "0" os.environ["WORLD_SIZE"] = "2" ContinuousBatchingConfig() self.assertEqual(os.environ.get("NCCL_GRAPH_MIXING_SUPPORT"), "0") # Explicitly disabled flag: env var is left untouched even under a distributed launch os.environ.pop("NCCL_GRAPH_MIXING_SUPPORT", None) ContinuousBatchingConfig(disable_nccl_graph_mixing=False) self.assertNotIn("NCCL_GRAPH_MIXING_SUPPORT", os.environ) # setdefault semantics: a pre-existing value is preserved os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "1" ContinuousBatchingConfig() self.assertEqual(os.environ.get("NCCL_GRAPH_MIXING_SUPPORT"), "1") finally: if original_nccl is None: os.environ.pop("NCCL_GRAPH_MIXING_SUPPORT", None) else: os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = original_nccl if original_ws is None: os.environ.pop("WORLD_SIZE", None) else: os.environ["WORLD_SIZE"] = original_ws class ContinuousBatchingPauseTest(unittest.TestCase): """Tests for the pause API in states a real single-rank generation loop cannot reach. The generation loop is replaced by a plain thread driving `BackgroundThreadStatus` directly, so no model or accelerator is needed. Pausing a real generation loop is tested in `ContinuousBatchingWithAcceleratorTest`.""" @staticmethod def _get_minimal_manager(status: BackgroundThreadStatus, thread: threading.Thread | None): """Builds a bare manager with just enough states for `pause` and `stop` to work.""" manager = ContinuousBatchingManager.__new__(ContinuousBatchingManager) manager._wake_up_loop = threading.Event() manager.background_thread_status = status manager._generation_thread = thread manager.distributed_helper = SimpleNamespace(cpu_comm_group=None, global_rank=0) # Attributes torn down by `stop` manager.batch_processor = None manager._original_attn_impl = None return manager def test_pause_on_a_peer_initiated_pause(self) -> None: """Under TP the pause request is MAX-reduced across ranks, so a loop can be paused before a local thread asked for it. A thread arriving on an already paused loop must enter the pause directly.""" status = BackgroundThreadStatus() resumed = threading.Event() def generation_loop() -> None: """Simulates a loop that was paused by a thread on another rank.""" try: self.assertFalse(status.is_pause_requested()) # should be false because this rank did not pause status.pause_and_wait() # enters the pause on behalf of the other rank resumed.set() # signals that the pause ended finally: status.mark_as_stopped() # Setup the loop thread and the manager loop_thread = threading.Thread(target=generation_loop) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) # Enter an already paused loop and check that the loop did not resume yet entered = [] with manager.pause(): entered.append(True) self.assertFalse(resumed.is_set(), "the loop resumed while the pause was held") loop_thread.join(timeout=3) # should leave enough for the loop to resume self.assertEqual(entered, [True], "the thread did not enter the already paused loop") self.assertTrue(resumed.is_set(), "the loop did not resume after the pause was released") def test_pause_request_is_dropped_when_the_wait_is_interrupted(self) -> None: """If an exception (e.g. KeyboardInterrupt) is raised while waiting for the loop to pause, the pause request must be dropped. Otherwise `is_pause_requested` stays true and the loop pauses forever with nobody to release it. The exception is injected through the wait predicate, so it is raised inside `wait_for` with the lock held, like a real interrupt would be.""" status = BackgroundThreadStatus() loop_resumed = threading.Event() def generation_loop() -> None: try: status.pause_and_wait() loop_resumed.set() finally: status.mark_as_stopped() loop_thread = threading.Thread(target=generation_loop) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) # Wait for the loop to pause, so that the interrupted request is the only one keeping it paused deadline = time.monotonic() + 3 while not status.is_pause_requested() and time.monotonic() < deadline: time.sleep(0.01) def interrupted_predicate() -> bool: raise KeyboardInterrupt("interrupted while waiting for the pause") # Replace the pause predicate with a function that raises an exception, so that it is raised while waiting status._pause_predicate = interrupted_predicate with self.assertRaises(KeyboardInterrupt): with manager.pause(): pass self.assertFalse(status.is_pause_requested(), "the interrupted pause request was not withdrawn") self.assertFalse(status.is_pause_requested(local=True), "the local pause counter was not decremented") loop_thread.join(timeout=3) self.assertTrue(loop_resumed.is_set(), "the loop stayed paused after the interrupted request") def test_stop_from_inside_a_pause_raises(self) -> None: """`stop` waits for the generation loop to exit, and the loop cannot exit while a pause is held. So a thread calling `stop` from inside its own pause would wait forever: it must raise instead, whatever `block` and `timeout` are.""" status = BackgroundThreadStatus() loop_thread = threading.Thread(target=status.pause_and_wait) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) try: with manager.pause(): for kwargs in [{}, {"block": False}, {"block": True, "timeout": 5.0}]: with self.assertRaises(RuntimeError) as context: manager.stop(**kwargs) self.assertIn("from inside a pause", str(context.exception), f"wrong error for {kwargs}") finally: loop_thread.join(timeout=10) self.assertFalse(loop_thread.is_alive()) def test_stop_from_another_thread_while_paused(self) -> None: """A thread that does not hold a pause can call `stop` while another thread does: the stop waits for the pause to be released and the loop to exit. Only a blocking stop is tested, since a non-blocking stop tears down the manager while the loop is still running, which is a pre-existing issue of `stop` unrelated to pausing.""" status = BackgroundThreadStatus() loop_thread = threading.Thread(target=status.pause_and_wait) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) holding = threading.Event() release = threading.Event() def holder() -> None: with manager.pause(): holding.set() release.wait(timeout=10) holder_thread = threading.Thread(target=holder) holder_thread.start() try: self.assertTrue(holding.wait(timeout=10), "the holder did not get its pause") # A blocking stop should wait for the holder to release the pause stopped = threading.Event() stopper = threading.Thread(target=lambda: (manager.stop(), stopped.set())) stopper.start() self.assertFalse(stopped.wait(timeout=0.5), "stop returned while the pause was still held") release.set() stopper.join(timeout=10) self.assertTrue(stopped.is_set(), "stop did not complete after the pause was released") finally: release.set() holder_thread.join(timeout=10) loop_thread.join(timeout=10) self.assertFalse(loop_thread.is_alive()) @staticmethod def _pause_in_thread(manager, timeout: float = 10) -> tuple[bool, Exception | None]: """Enters and exits `pause` in a separate thread, so a hang shows up as a timeout instead of blocking the test suite. Returns whether the thread finished and the exception it raised, if any.""" raised: list[Exception] = [] def enter_and_leave() -> None: try: with manager.pause(): pass except Exception as e: raised.append(e) thread = threading.Thread(target=enter_and_leave) thread.start() thread.join(timeout=timeout) finished = not thread.is_alive() return finished, (raised[0] if raised else None) def test_pause_raises_when_the_loop_cannot_pause(self) -> None: """When the loop will never pause, `pause` must raise instead of blocking. This covers a loop that died on a fatal error, a loop that hard-stopped right after seeing the request, a manager that was never started and a manager that was already stopped.""" # 1. The loop dies with a fatal error while a pause is pending: the error is chained as `__cause__` fatal_error = RuntimeError("boom in the forward pass") status = BackgroundThreadStatus() keep_unwinding = threading.Event() def dying_loop() -> None: try: while not status.is_pause_requested(): time.sleep(0.001) # Same order as `_handle_critical_error`: request the hard stop, then record the error status.request_stop(BackgroundThreadStatus.HARD_STOP, 0) status.record_fatal_error(fatal_error) keep_unwinding.wait(timeout=10) # keep the thread alive while `pause` gives up finally: status.mark_as_stopped() loop_thread = threading.Thread(target=dying_loop) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) finished, error = self._pause_in_thread(manager) self.assertTrue(finished, "pause blocked after the generation loop died") self.assertIsInstance(error, RuntimeError) self.assertIs(error.__cause__, fatal_error) # The manager should be stopped rather than left half-dead self.assertEqual(status.local_status, BackgroundThreadStatus.HARD_STOP) keep_unwinding.set() loop_thread.join(timeout=10) # 2. The loop hard-stops with a pause pending. `_update_tp_group_state` returns on HARD_STOP before reaching the # pause window, so the loop sees the request but never pauses. There is no fatal error, but `pause` must still # return. status = BackgroundThreadStatus() def hard_stopping_loop() -> None: try: while not status.is_pause_requested(): time.sleep(0.001) finally: status.mark_as_stopped() loop_thread = threading.Thread(target=hard_stopping_loop) loop_thread.start() manager = self._get_minimal_manager(status, loop_thread) finished, error = self._pause_in_thread(manager) loop_thread.join(timeout=10) self.assertTrue(finished, "pause blocked after the generation loop hard-stopped") self.assertIsInstance(error, RuntimeError) self.assertGreaterEqual(status.local_status, BackgroundThreadStatus.HARD_STOP) # 3. A manager that was never started manager = self._get_minimal_manager(BackgroundThreadStatus(), None) finished, error = self._pause_in_thread(manager) self.assertTrue(finished, "pause blocked on a manager that was never started") self.assertIsInstance(error, RuntimeError) # 4. A manager whose generation thread has already exited status = BackgroundThreadStatus() loop_thread = threading.Thread(target=status.mark_as_stopped) loop_thread.start() loop_thread.join(timeout=10) manager = self._get_minimal_manager(status, loop_thread) finished, error = self._pause_in_thread(manager) self.assertTrue(finished, "pause blocked on an already stopped manager") self.assertIsInstance(error, RuntimeError) @require_torch_accelerator class ContinuousBatchingWithAcceleratorTest(unittest.TestCase): # -----------------------------------------------Parity tests----------------------------------------------- # # Ensure continuous batching and non-continuous batching generation produce the same outputs # # ---------------------------------------------------------------------------------------------------------- # @require_deterministic_for_xpu @with_flush_memory def _test_continuous_batching_parity( self, model_id: str, continuous_batching_config: ContinuousBatchingConfig, attn_implementation: str, max_new_tokens: int = 20, num_repeat_prompts: int = 1, upcast_lm_head: bool = False, ) -> None: """Tests the parity between continuous batching and non-continuous batching generation.""" # Skip the test if Flash Attention 2 or 3 is required but not available. flexible_flash_skip(self, attn_implementation) # Skip the test if cuda graph is on but the device does not support graph capture device_type = torch_device.type if isinstance(torch_device, torch.device) else torch_device if any(continuous_batching_config.cuda_graph_booleans) and device_type not in DEVICE_TYPE_TO_GRAPH_NAME: supported = list(DEVICE_TYPE_TO_GRAPH_NAME.keys()) self.skipTest(f"CUDA graph is only supported on {supported}, but {device_type = }. Skipping test.") # If the config turns on compile, change the generation config to use the default mode instead of # max-autotune-no-cudagraphs which can change the kernels between generate_batch and generate if continuous_batching_config.default_compile_level > 0: # Paged attention is wrapped in @torch.compiler.disable so fullgraph is not possible compile_config = CompileConfig(mode="default", fullgraph=False, dynamic=True) continuous_batching_config.varlen_compile_config = compile_config # Eager and SDPA implementations get a precision boost to account for the fact that an attention mask is used in # continuous batching but not in generate is_fa = is_flash_attention_requested(requested_attention_implementation=attn_implementation) dtype = "auto" if is_fa else torch.float32 # Disable auto-switch to flash: if requested by the test, we keep SDPA or eager continuous_batching_config.auto_switch_to_flash = False # Prepare inputs paged_attn_implem = "paged|eager" if attn_implementation == "eager" else attn_implementation tokenizer, model = get_tokenizer_and_model(model_id, paged_attn_implem, torch_device, dtype, upcast_lm_head) if ( attn_implementation == "flash_attention_2" and torch_device == "cpu" and getattr(model.config, "sliding_window", None) is not None and model.config.sliding_window > 0 ): self.skipTest("Flash Attention 2 with sliding window attention is not supported on CPU. Skipping test.") user_messages = _DEFAULT_USER_MESSAGES * num_repeat_prompts input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) model.generation_config.max_new_tokens = max_new_tokens model.generation_config.do_sample = False # Generation with continuous batching continuous_batching_outputs = model.generate_batch( inputs=input_ids, generation_config=model.generation_config, continuous_batching_config=continuous_batching_config, ) # Prepare non-continuous batching inputs and model inputs = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=False) num_input_tokens = inputs.input_ids.shape[1] # Flush compile cache if CB used compile if continuous_batching_config.default_compile_level > 0: flush_memory(flush_compile=True) # Generation without continuous batching (reload model to avoid any state contamination) _, model = get_tokenizer_and_model(model_id, attn_implementation, torch_device, dtype, upcast_lm_head) model.generation_config.max_new_tokens = max_new_tokens model.generation_config.do_sample = False model.generation_config.use_cuda_graph = ( any(continuous_batching_config.cuda_graph_booleans) and device_type in DEVICE_TYPE_TO_GRAPH_NAME ) model.generation_config.compile_config = continuous_batching_config.varlen_compile_config # Create a static cache if compile_config is set, because regular generate requires a compileable cache past_key_values = None if model.generation_config.compile_config is not None: max_cache_len = num_input_tokens + max_new_tokens past_key_values = StaticCache(config=model.config, max_cache_len=max_cache_len) generate_outputs = model.generate( **inputs.to(torch_device), generation_config=model.generation_config, past_key_values=past_key_values ) for i, user_message in enumerate(user_messages): # Find the corresponding request in the continuous batching outputs input_tokens = inputs.input_ids[i][inputs.attention_mask[i] == 1].tolist() key_to_pop = None for key, state in continuous_batching_outputs.items(): if state.prompt_ids == input_tokens: key_to_pop = key break if key_to_pop is None: self.fail(f"Request {i} not found in continuous batching outputs") continuous_batching_output = continuous_batching_outputs.pop(key_to_pop).generated_tokens generate_output = generate_outputs[i][num_input_tokens:].tolist() while generate_output[-1] == model.generation_config.pad_token_id: generate_output.pop() if continuous_batching_output != generate_output: decoded_continuous_batching_output = tokenizer.decode(continuous_batching_output) decoded_generate_output = tokenizer.decode(generate_output) msg = f"Test failed for {model_id = } {continuous_batching_config = }, {attn_implementation = }\n" msg += f"User message : {repr(user_message)}\n" msg += f"Continuous batching output: {repr(decoded_continuous_batching_output)}\n" msg += f"Generate output : {repr(decoded_generate_output)}" self.fail(msg) @parameterized.expand( list( itertools.product( [False, True], ["eager", "sdpa", "flash_attention_2"], [False, True], ) ) ) @slow def test_continuous_batching_config_combinations_no_compile( self, allow_block_sharing: bool, attn_implementation: str, use_cuda_graph: bool, ) -> None: # Compiling adds a lot of overhead, so it's better not to include here (2*3*2=12 tests because of cross product) model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" continuous_batching_config = ContinuousBatchingConfig( allow_block_sharing=allow_block_sharing, use_cuda_graph=use_cuda_graph, default_compile_level=0, ) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation=attn_implementation, ) @parameterized.expand([("eager", False), ("sdpa", False), ("sdpa", True), ("flash_attention_2", True)]) @slow def test_continuous_batching_config_combinations_with_compile( self, attn_implementation: str, use_cuda_graph: bool, ) -> None: model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, default_compile_level=1, ) # Flash runs in bf16, where compile and padding can flip near-tied greedy picks: upcast the lm_head self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation=attn_implementation, upcast_lm_head=is_flash_attention_requested(requested_attention_implementation=attn_implementation), ) @parameterized.expand( [ (False, False), # _supports_flash_attn=False: stays on sdpa (False, True), # same, even with auto switch on: stay on SDPA (True, False), # flash-capable model on a non-flash impl, but auto switch off: stay on SDPA (True, True), # flash-capable model on a non-flash impl, with auto switch on: switch to flash ] ) @slow def test_switch_to_cb_friendly_attn(self, supports_flash_attn: bool, auto_switch_to_flash: bool) -> None: """Continuous batching switches to a paged (ideally flash) attention and restores the original on stop.""" flash_available = is_flash_attn_2_available(kernels_fallback_ok=True) flash_available |= is_flash_attn_3_available(kernels_fallback_ok=True) if not flash_available: self.skipTest("Flash attention is unavailable, cannot test the auto-switch to flash.") # Retrieve the model model_id = "Qwen/Qwen2.5-0.5B-Instruct" _, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.bfloat16) model._supports_flash_attn = supports_flash_attn original_attn_impl = model.config._attn_implementation # Creating the manager switches the model to a CB-friendly attention implementation cb_config = ContinuousBatchingConfig( num_blocks=8, block_size=32, use_cuda_graph=False, auto_switch_to_flash=auto_switch_to_flash ) manager = model.init_continuous_batching(continuous_batching_config=cb_config) final_attn_impl = model.config._attn_implementation if supports_flash_attn and auto_switch_to_flash: self.assertTrue(is_flash_attention_requested(requested_attention_implementation=final_attn_impl)) else: self.assertEqual(final_attn_impl, "sdpa") # Stopping the manager restores the original attention implementation manager.start() manager.stop(block=True) self.assertEqual(model.config._attn_implementation, original_attn_impl) # FIXME: Qwen2.5-0.5B-Instruct is not here because it's broken (it uses a repetition penalty logits processor) # TODO: replace gemma2 with a tiny version of GPT-OSS? That way we can test sliding window AND attention sink @parameterized.expand( list( itertools.product( ["google/gemma-2-2b-it"], [False, True], [False, True], ) ) ) @slow def test_continuous_batching_diverse_models(self, model_id: str, use_cuda_graph: bool, use_compile: bool) -> None: continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, default_compile_level=1 if use_compile else 0 ) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="flash_attention_2", ) @slow @require_torch_accelerator def test_hybrid_sliding_window_crossing_parity(self) -> None: """Tests that a hybrid full+sliding model matches generate() token-for-token on prompts that cross the sliding window, both during decode (prompt just under the window) and during prefill (prompt over the window). This exercises the rolling-buffer read and write indices past the window boundary.""" model_id = "google/gemma-3-1b-it" max_new_tokens = 30 # Use float32 so any mismatch is a real indexing bug and not a numerical tie-flip tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, dtype=torch.float32) sliding_window = model.config.sliding_window # Build prompts around the sliding window size base_text = "The mitochondria is the powerhouse of the cell. " * 200 long_ids = tokenizer(base_text).input_ids prompts = [ tokenizer("The capital of France is").input_ids, long_ids[: sliding_window - 10], # crosses the window during decode long_ids[: sliding_window + 100], # crosses the window during prefill ] # Reference: plain generate, one prompt at a time to avoid any padding effect references = [] for ids in prompts: with torch.no_grad(): out = model.generate( torch.tensor([ids], device=torch_device), max_new_tokens=max_new_tokens, do_sample=False, use_cache=True, ) references.append(out[0, len(ids) :].tolist()) # Continuous batching, all prompts at once tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, dtype=torch.float32) gen_config = GenerationConfig(max_new_tokens=max_new_tokens, do_sample=False, eos_token_id=None) cb_config = ContinuousBatchingConfig( use_cuda_graph=False, use_async_batching=False, auto_switch_to_flash=False ) results = model.generate_batch( inputs=prompts, generation_config=gen_config, continuous_batching_config=cb_config, progress_bar=False, ) ordered_keys = sorted(results.keys(), key=lambda x: int(x.split("_")[1])) for i, (key, reference) in enumerate(zip(ordered_keys, references)): self.assertEqual( results[key].generated_tokens, reference, f"CB output diverges from generate() for prompt {i} (length {len(prompts[i])}, {sliding_window = })", ) @parameterized.expand([(True, False), (False, True)]) @require_flash_attn_3 @slow def test_continuous_batching_tuple_cuda_graph(self, varlen_cg: bool, decode_cg: bool) -> None: """Tests that use_cuda_graph can be a tuple to independently control varlen and decode paths.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=(varlen_cg, decode_cg), use_async_batching=False, ) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="flash_attention_3", ) def test_continuous_batching_fast(self) -> None: model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=False, allow_block_sharing=False, use_async_batching=False, default_compile_level=0, ) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", ) @with_flush_memory def test_memory_footprint_respects_max_memory_percent(self) -> None: """Runs a short batched generation with an auto-inferred cache size and checks that the real peak memory the cache, IOs and activations add on top of the model stays within the configured max_memory_percent of the free device memory (within a tolerance, to account for allocator rounding and the spare cache blocks).""" if not (torch.cuda.is_available() or is_torch_xpu_available()): self.skipTest("Peak memory tracking requires CUDA or XPU") accelerator = torch.cuda if torch.cuda.is_available() else torch.xpu model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" max_memory_percent = 0.5 tolerance = 0.05 tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.float16) input_ids = get_generation_inputs(_DEFAULT_USER_MESSAGES, tokenizer, for_continuous_batching=True) model.generation_config.max_new_tokens = 20 model.generation_config.do_sample = False cb_config = ContinuousBatchingConfig( max_memory_percent=max_memory_percent, use_cuda_graph=False, use_async_batching=False, auto_switch_to_flash=False, ) # Budget = max_memory_percent of the free device memory, computed exactly as the memory handler does _, total, reserved, allocated = get_device_and_memory_breakdown() budget = max_memory_percent * (total - max(allocated, reserved)) # Everything continuous batching allocates (cache + IOs + activations) lives on top of the already-loaded model accelerator.reset_peak_memory_stats() model_footprint = accelerator.memory_allocated() model.generate_batch( inputs=input_ids, generation_config=model.generation_config, continuous_batching_config=cb_config, ) cb_footprint = accelerator.max_memory_allocated() - model_footprint self.assertLessEqual( cb_footprint, budget * (1 + tolerance), f"Continuous batching peak footprint {cb_footprint / 1024**2:.0f} MB exceeds the {max_memory_percent:.0%} " f"budget of {budget / 1024**2:.0f} MB (+{tolerance:.0%} tolerance)", ) def test_continuous_batching_long_generate(self) -> None: model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=True, allow_block_sharing=True, use_async_batching=False, default_compile_level=1 ) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", max_new_tokens=80, ) @require_torch_multi_accelerator @with_flush_memory def test_cuda_graph_capture_on_non_default_device(self) -> None: """The background generation thread does not inherit the main thread's current CUDA device: it's a thread-local variable. A freshly spawned thread starts on device 0. A model living on any other device therefore used to bind its CUDA graph memory pool to device 0 while capturing on the model's device, which tripped an allocator assert (`use_count > 0`) on the first in-thread capture. Regression test for #48312. The graphs must be captured by the background thread, so this test must not call `warmup()` before `start()`: warming up captures on the main thread instead, which is a documented workaround for the bug. """ if torch_device != "cuda": self.skipTest("CUDA graph is only supported on CUDA devices. Skipping test.") # The whole point of the test is that the model does NOT live on device 0, where the background thread starts device = f"{torch_device}:1" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", device, torch.bfloat16) gen_config = GenerationConfig(max_new_tokens=5, do_sample=False, eos_token_id=tokenizer.eos_token_id) cb_config = ContinuousBatchingConfig( use_cuda_graph=True, use_async_batching=False, max_memory_percent=0.2, auto_switch_to_flash=False ) manager = model.init_continuous_batching(generation_config=gen_config, continuous_batching_config=cb_config) manager.start() try: input_ids = get_generation_inputs(_DEFAULT_USER_MESSAGES, tokenizer, for_continuous_batching=True) request_ids = manager.add_requests(input_ids, max_new_tokens=5) for _ in request_ids: output = manager.get_result(timeout=120) self.assertIsNotNone(output, "The background generation thread died before returning a result") self.assertIsNone(output.error, f"Request failed during in-thread graph capture: {output.error}") self.assertTrue(output.generated_tokens, "Request returned no generated tokens") # The cache and the graph memory pool must have been created on the model's device, not on device 0 self.assertEqual(manager.batch_processor.cache.device, torch.device(device)) finally: manager.stop(block=True, timeout=60) @parameterized.expand([(False, False), (False, True), (True, False), (True, True)]) @slow def test_continuous_batching_log_probs(self, use_cuda_graph: bool, use_async_batching: bool) -> None: """Test that log probabilities match between continuous batching and regular generate.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Retrieve tokenizer, model and eos_token_id (required otherwise logits will be misaligned) tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.float32) eos_token_id = model.config.eos_token_id # type: ignore[attr-defined] # Run CB generation user_messages = ["What is 2+2?", "Hello world"] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) gen_config = GenerationConfig(max_new_tokens=10, do_sample=False, eos_token_id=eos_token_id) continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching, return_logprobs=True, auto_switch_to_flash=False, ) cb_outputs = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=continuous_batching_config ) # Load fresh model for regular generate tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.float32) # Run regular generate regular_outputs, regular_logprobs = regular_generate( model=model, tokenizer=tokenizer, user_messages=user_messages, max_new_tokens=10, do_sample=False, output_scores=True, eos_token_id=eos_token_id, ) # Compare log_probs for each request, matching by prompt_ids for i, cb_output in enumerate(cb_outputs.values()): # Compare Cb and regular generate outputs cb_output_ids = cb_output.generated_tokens regular_output_ids = regular_outputs[i] self.assertEqual(len(cb_output_ids), len(regular_output_ids)) self.assertEqual(cb_output_ids, regular_output_ids) # Retrieve logprobs from CB and regular generate cb_logprobs = cb_output.logprobs expected_logprobs = regular_logprobs[i] # Because of padding, we need to truncate to the same length min_len = min(len(cb_logprobs), len(expected_logprobs)) cb_logprobs = cb_logprobs[:min_len] expected_logprobs = expected_logprobs[:min_len] self.assertEqual(len(cb_logprobs), len(expected_logprobs)) # Compare with tolerance for floating point differences (because of padding, tol is higher for cuda graphs) delta = 2e-5 if use_cuda_graph else 1e-5 for j, (cb_lp, exp_lp) in enumerate(zip(cb_logprobs, expected_logprobs)): error_msg = f"logprob mismatch at position {j} for request {i}: CB={cb_lp}, expected={exp_lp}" self.assertAlmostEqual(cb_lp, exp_lp, delta=delta, msg=error_msg) def test_continuous_batching_few_blocks(self) -> None: """This test verifies that generation works with a very small number of blocks, ie. small enough that we need to offload a request at some point. To add more complexity, we repeat the same prompt 4 times and enable prefix sharing.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # 4 blocks of 32 tokens: small enough to force offloading continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=True, allow_block_sharing=True, use_async_batching=False, num_blocks=4, page_size=32 ) # Patch offload_requests to verify it's called at least once original_offload = OffloadingManager.offload_requests with patch.object( OffloadingManager, "offload_requests", autospec=True, side_effect=original_offload ) as mock_offload: self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", max_new_tokens=30, num_repeat_prompts=4, ) self.assertTrue(mock_offload.called, "Offload method was not called.") # ---------------------------------------Streaming tests--------------------------------------- # # Ensures the requests have the right behavior with and without streaming # # --------------------------------------------------------------------------------------------- # def _test_streaming_or_not_request(self, with_streaming: bool, with_non_streaming: bool) -> None: model_id = "Qwen/Qwen2.5-0.5B-Instruct" max_new_tokens = 3 tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device) cb_config = ContinuousBatchingConfig(auto_switch_to_flash=False) manager = model.init_continuous_batching(continuous_batching_config=cb_config) manager.logit_processor.clear() manager.start() user_messages = ["What is the Transformers library known for?"] inputs = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True)[0] # Test with non-streaming if with_non_streaming: request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=False) # In non-streaming mode, the total number of generated tokens is equal to the max new tokens chunk = next(manager.request_id_iter(request_id)) self.assertEqual(len(chunk.generated_tokens), max_new_tokens) # Test with streaming if with_streaming: request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=True) # In streaming mode, the total number of generated tokens is incremented by 1 on each iteration chunk_1 = next(manager.request_id_iter(request_id)) self.assertEqual(len(chunk_1.generated_tokens), 1) chunk_2 = next(manager.request_id_iter(request_id)) self.assertEqual(len(chunk_2.generated_tokens), 2) chunk_3 = next(manager.request_id_iter(request_id)) self.assertEqual(len(chunk_3.generated_tokens), 3) manager.stop(block=True) def test_streaming_request(self) -> None: self._test_streaming_or_not_request(with_streaming=True, with_non_streaming=False) def test_non_streaming_request(self) -> None: self._test_streaming_or_not_request(with_streaming=False, with_non_streaming=True) def test_streaming_and_non_streaming_requests_can_alternate(self) -> None: self._test_streaming_or_not_request(with_streaming=True, with_non_streaming=True) def test_register_result_handler(self) -> None: """Test that register_result_handler receives streaming outputs through the OutputRouter.""" import asyncio model_id = "Qwen/Qwen2.5-0.5B-Instruct" max_new_tokens = 3 tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device) cb_config = ContinuousBatchingConfig(auto_switch_to_flash=False) manager = model.init_continuous_batching(continuous_batching_config=cb_config) manager.logit_processor.clear() manager.start() user_messages = ["What is the Transformers library known for?"] inputs = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True)[0] async def collect_results(): token_counts = [] future = asyncio.get_running_loop().create_future() def on_result(output): token_counts.append(len(output.generated_tokens)) if output.is_finished(): self.assertEqual(output.status, RequestStatus.FINISHED) # fail if the request failed future.set_result(True) request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=True) manager.register_result_handler(request_id, on_result) await asyncio.wait_for(future, timeout=30) return token_counts token_counts = asyncio.run(collect_results()) # Streaming via handler: incremental token count, same as request_id_iter self.assertEqual(token_counts, [1, 2, 3]) # Queue should be empty — everything went through the handler self.assertTrue(manager.output_router.output_queue.empty()) manager.stop(block=True) # -----------------------------------------Misc. tests----------------------------------------- # # Various tests that don't fit into the other categories # # --------------------------------------------------------------------------------------------- # def _test_block_sharing(self, model_id: str, expected_attn_types: set[str], input_msg: str) -> None: # Use float32 for SDPA to handle precision differences from attention masks (same as parity test) tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, dtype=torch.float32) model._supports_flash_attn = False # Configure generation for parity: disable processors not supported by CB (like repetition_penalty) model.generation_config.max_new_tokens = 32 model.generation_config.do_sample = False model.generation_config.repetition_penalty = None # Get expected output from regular generate for parity check expected_output_tokens, _ = regular_generate(model, tokenizer, [input_msg]) # Prefix sharing only activates when every allocator can share, i.e. on full-attention-only models prefix_sharing_expected = expected_attn_types == {"full_attention"} # Spy on prefix matching to observe how many prompt tokens get matched matched_tokens: list[int] = [] original_search = PagedAttentionCache.search_prefix_match def spy_search(cache_self, request_id, prompt_ids): matched = original_search(cache_self, request_id, prompt_ids) matched_tokens.append(matched) return matched cb_context_manager = model.continuous_batching_context_manager( generation_config=model.generation_config, continuous_batching_config=ContinuousBatchingConfig(page_size=32), ) with ( patch.object(PagedAttentionCache, "search_prefix_match", autospec=True, side_effect=spy_search), cb_context_manager as manager, ): # The prompt must span 2 blocks (>= 33 tokens so decode completes the 2nd block: the prompt plus 31 # written decode tokens must reach 64) and its 2nd block must not be complete at prefill (< 64 tokens) inputs = get_generation_inputs([input_msg], tokenizer, for_continuous_batching=True)[0] self.assertGreaterEqual(len(inputs), 33, f"Input length is {len(inputs)} instead of at least 33") self.assertLess(len(inputs), 64, f"Input length is {len(inputs)} instead of less than 64") # First request, which populates the cache with 2 complete blocks (1 by prefill, 1 during decode) request_id = manager.add_request(inputs, max_new_tokens=32) chunk_no_reuse = next(manager.request_id_iter(request_id)) if manager.batch_processor is None: raise RuntimeError("Batch processor is None even after a request was added.") cache = manager.batch_processor.cache # Assert there is one allocator per expected attention type, and prefix sharing is on when possible self.assertEqual(set(cache.cache_allocators.keys()), expected_attn_types) self.assertEqual(cache.use_prefix_sharing, prefix_sharing_expected) # On a full-attention-only model, the first request leaves 2 hashed blocks; on a hybrid model, prefix # sharing is off so no block is ever hashed num_hashed_blocks = sum(len(ca.ledger.hash_to_block) for ca in cache.cache_allocators.values()) expected_hashes = 2 if prefix_sharing_expected else 0 self.assertEqual(num_hashed_blocks, expected_hashes, f"{num_hashed_blocks = }, {expected_hashes = }") self.assertEqual(sum(matched_tokens), 0) # Second request, which should match the first complete block of the first request as a prefix request_id = manager.add_request(inputs, max_new_tokens=32) chunk_with_reuse = next(manager.request_id_iter(request_id)) # De-duplication: the second request computes identical content, so no new hash is registered num_hashed_blocks = sum(len(ca.ledger.hash_to_block) for ca in cache.cache_allocators.values()) self.assertEqual(num_hashed_blocks, expected_hashes, f"{num_hashed_blocks = }, {expected_hashes = }") # The prompt has less than 64 tokens, so only its first block can be matched (32 tokens) expected_total_prefix_length = 32 if prefix_sharing_expected else 0 self.assertEqual( sum(matched_tokens), expected_total_prefix_length, f"Expected {expected_total_prefix_length} matched prefix tokens, got {sum(matched_tokens)}", ) # Check the outputs were the same (block sharing should produce identical results) self.assertEqual(chunk_no_reuse.generated_tokens, chunk_with_reuse.generated_tokens) # Verify parity with regular generate self.assertEqual(chunk_no_reuse.generated_tokens, expected_output_tokens[0]) # ------------------------------------------------- Pause tests ------------------------------------------------- # def _started_manager_with_requests(self, max_new_tokens: int = 200): """Returns a started manager with a few requests in flight, along with the model and the number of requests.""" tokenizer, model = get_tokenizer_and_model("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "sdpa", torch_device) input_ids = get_generation_inputs(_DEFAULT_USER_MESSAGES, tokenizer, for_continuous_batching=True) cb_config = ContinuousBatchingConfig( use_cuda_graph=False, use_async_batching=False, auto_switch_to_flash=False ) manager = model.init_continuous_batching(continuous_batching_config=cb_config) manager.logit_processor.clear() manager.warmup() # so that no graph capture happens on the loop thread during the test manager.start() for ids in input_ids: manager.add_request(ids, max_new_tokens=max_new_tokens) return model, manager, len(input_ids) def _wait_until_generating(self, manager, timeout: float = 60) -> None: """Blocks until the generation loop has run at least one step. Requests are admitted to the scheduler after the pause window in `_update_tp_group_state`, so a pause taken on the first iteration would pause an empty scheduler and the step counter checks would be meaningless.""" # `current_batch` is only set once the loop thread has started deadline = time.monotonic() + timeout while getattr(manager, "current_batch", 0) == 0 and time.monotonic() < deadline: time.sleep(0.01) self.assertGreater(getattr(manager, "current_batch", 0), 0, "the generation loop never started generating") @staticmethod def _collect_finished(manager, expected: int, timeout: float = 120): """Drains the manager's outputs until `expected` requests have finished or the timeout expires.""" outputs = {} deadline = time.monotonic() + timeout while len(outputs) < expected and time.monotonic() < deadline: result = manager.get_result(timeout=1.0) if result is not None and result.is_finished(): outputs[result.request_id] = result return outputs @with_flush_memory def test_pause_during_generation(self) -> None: """Pauses the manager several times while requests are in flight, like a trainer sharing the model would. The generation loop must not run a step while the pause is held, and all requests must still complete.""" num_pauses = 5 model, manager, num_requests = self._started_manager_with_requests() try: self._wait_until_generating(manager) for i in range(num_pauses): with manager.pause(): batches_at_pause = manager.current_batch # The step counter check below only means something if the loop still has work to do self.assertTrue( manager.batch_processor.has_pending_requests(), f"generation finished before pause {i}", ) # Stands in for a trainer updating the weights in place with torch.no_grad(): for parameter in model.parameters(): parameter.add_(0.0) # A generation step takes a few ms, so hold the pause long enough for a loop that ignored it to # run several steps time.sleep(0.3) self.assertEqual(manager.current_batch, batches_at_pause, f"the loop ran during pause {i}") outputs = self._collect_finished(manager, num_requests) finally: manager.stop(block=True) self.assertEqual(len(outputs), num_requests, "some requests did not complete") for request_id, output in outputs.items(): self.assertGreater(len(output.generated_tokens), 0, f"{request_id} generated no tokens") @with_flush_memory def test_pause_with_several_holders(self) -> None: """Several threads can hold the pause at once, and the loop must stay paused until the last one releases it. The holders are released one at a time and the step counter is checked after each release.""" num_holders = 3 model, manager, num_requests = self._started_manager_with_requests() # The main thread is the extra party: it runs its checks once every holder is inside its pause all_inside = threading.Barrier(num_holders + 1, timeout=60) may_leave = [threading.Event() for _ in range(num_holders)] has_left = [threading.Event() for _ in range(num_holders)] holder_errors = [] def holder(index: int) -> None: try: with manager.pause(): all_inside.wait() may_leave[index].wait(timeout=60) has_left[index].set() except Exception as error: # a holder that fails to enter the pause would break the barrier holder_errors.append(error) threads = [threading.Thread(target=holder, args=(index,)) for index in range(num_holders)] try: self._wait_until_generating(manager) for thread in threads: thread.start() all_inside.wait() batches_at_pause = manager.current_batch # The step counter checks below only mean something if the loop still has work to do self.assertTrue( manager.batch_processor.has_pending_requests(), "generation finished before the holders entered their pause", ) for index in range(num_holders): may_leave[index].set() self.assertTrue(has_left[index].wait(timeout=30), f"holder {index} did not exit its pause") # Give a loop that resumed too early time to run a few steps time.sleep(0.3) if index < num_holders - 1: self.assertEqual( manager.current_batch, batches_at_pause, f"the loop resumed with {num_holders - index - 1} holders still inside", ) for thread in threads: thread.join(timeout=30) outputs = self._collect_finished(manager, num_requests) finally: for event in may_leave: event.set() manager.stop(block=True) self.assertEqual(holder_errors, [], f"some holders failed to enter the pause: {holder_errors}") self.assertGreater( manager.current_batch, batches_at_pause, "the loop did not resume after the last holder left" ) self.assertEqual(len(outputs), num_requests, "some requests did not complete") def test_prefix_sharing(self) -> None: model_id = "Qwen/Qwen2.5-0.5B-Instruct" input_msg = "What is the Transformers library known for?" return self._test_block_sharing(model_id, {"full_attention"}, input_msg) def test_block_sharing_with_hybrid_model(self) -> None: model_id = "google/gemma-3-1b-it" input_msg = "I am a software engineer looking to use open source software to build a new AI agent. What is the Transformers library known for?" return self._test_block_sharing(model_id, {"full_attention", "sliding_attention"}, input_msg) @parameterized.expand([True, False]) @require_flash_attn # otherwise the test can fail because attention bias has a very slight impact on SDPA and eager def test_num_return_sequences(self, allow_block_sharing: bool) -> None: model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" tokenizer, model = get_tokenizer_and_model(model_id, "flash_attention_2", torch_device) user_messages = _DEFAULT_USER_MESSAGES[:1] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) model.generation_config.max_new_tokens = 30 model.generation_config.do_sample = False # Generation with continuous batching cb_context_manager = model.continuous_batching_context_manager( continuous_batching_config=ContinuousBatchingConfig(allow_block_sharing=allow_block_sharing), block=True, timeout=5, ) # Main loop results = [] with cb_context_manager as manager: manager.num_return_sequences = 2 manager.add_requests(inputs=input_ids, max_new_tokens=30) requests_left = 2 while requests_left: result = manager.get_result(timeout=1) if result and result.is_finished(): self.assertEqual(result.status, RequestStatus.FINISHED) # fail if the request failed results.append(result) requests_left -= 1 else: if not manager.is_running(): break self.assertEqual(len(results), 2, f"Expected 2 results, but got {len(results) = }") self.assertEqual(results[0].generated_tokens, results[1].generated_tokens) # ----------------------------------Additional features tests---------------------------------- # # Tests to check addtional features of CB do not change its results # # --------------------------------------------------------------------------------------------- # @parameterized.expand( [ # SDPA: basic features or full features ("sdpa", False, False), ("sdpa", True, True), # FA2: full coverage ("flash_attention_2", False, False), ("flash_attention_2", False, True), ("flash_attention_2", True, False), ("flash_attention_2", True, True), # FA3: always turn on CUDA graphs ("flash_attention_3", True, False), ("flash_attention_3", True, True), ] ) @slow def test_continuous_batching_async( self, attn_implementation: str, use_cuda_graph: bool, use_compile: bool ) -> None: # Again, we try to not overly use_compile because it adds a lot of overhead model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Flash runs in bf16, where compile and padding can flip near-tied greedy picks: upcast the lm_head is_fa = is_flash_attention_requested(requested_attention_implementation=attn_implementation) self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=ContinuousBatchingConfig( allow_block_sharing=True, use_cuda_graph=use_cuda_graph, use_async_batching=True, default_compile_level=1 if use_compile else 0, ), attn_implementation=attn_implementation, upcast_lm_head=is_fa and use_compile, ) @parameterized.expand([(False, False), (False, True), (True, False), (True, True)]) @slow @require_flash_attn def test_flash_attn_with_kvcache_parity(self, use_cuda_graph: bool, use_async: bool) -> None: """Test that flash attention with kvcache (flash_attn_with_kvcache path) produces same outputs as varlen.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Upcast the lm_head: in bf16, near-tied logits round to the same value and padding (CUDA graphs) flips greedy picks tokenizer, model = get_tokenizer_and_model( model_id, "flash_attention_2", torch_device, torch.bfloat16, upcast_lm_head=True ) user_messages = _DEFAULT_USER_MESSAGES[:] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) gen_config = GenerationConfig(do_sample=False, max_new_tokens=20) continuous_batching_config = ContinuousBatchingConfig( page_size=256, num_blocks=64, max_batch_tokens=16, use_cuda_graph=use_cuda_graph, use_async_batching=use_async, ) # Generate with varlen path only continuous_batching_config.max_blocks_per_request = 0 outputs_varlen = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=continuous_batching_config ) # Generate with flash_attn_with_kvcache path for decode continuous_batching_config.max_blocks_per_request = 16 # Wrap the kvcache function to check the block table path is actually used og_kvcache_fn = fa_utils._flash_with_kvcache_fn called = [False] @functools.wraps(og_kvcache_fn) def spy_kvcache_fn(*args, **kwargs): called[0] = True return og_kvcache_fn(*args, **kwargs) fa_utils._flash_with_kvcache_fn = spy_kvcache_fn try: outputs_kvcache = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=continuous_batching_config ) # Check the global was not swapped back self.assertIs(fa_utils._flash_with_kvcache_fn, spy_kvcache_fn) finally: fa_utils._flash_with_kvcache_fn = og_kvcache_fn self.assertTrue(called[0], "flash_attn_with_kvcache decode path was not used.") self.assertEqual(len(outputs_varlen), len(outputs_kvcache)) for (_, out_varlen), (_, out_kvcache) in zip(outputs_varlen.items(), outputs_kvcache.items()): text_varlen = tokenizer.decode(out_varlen.generated_tokens, skip_special_tokens=True) text_kvcache = tokenizer.decode(out_kvcache.generated_tokens, skip_special_tokens=True) self.assertEqual(text_varlen, text_kvcache, f"Mismatch:\nvarlen: {text_varlen}\nkvcache: {text_kvcache}") @parameterized.expand([(False, False), (False, True), (True, True)]) @slow @require_flash_attn def test_upcast_lm_head_safety(self, use_cuda_graph: bool, use_async: bool) -> None: """Safety net for `upcast_lm_head`: CB with a fp32 lm_head must still match the outputs of a regular bf16 generate, with expectations generated by the latter. A few tokens suffice to catch an upcast that badly breaks generation.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" tokenizer, model = get_tokenizer_and_model( model_id, "flash_attention_2", torch_device, torch.bfloat16, upcast_lm_head=True ) # Skip the "Josh" prompt: its 5th token is an exact tie in bf16, which the fp32 lm_head breaks differently user_messages = [_DEFAULT_USER_MESSAGES[0], _DEFAULT_USER_MESSAGES[2]] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True) gen_config = GenerationConfig(do_sample=False, max_new_tokens=8) continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, use_async_batching=use_async ) outputs = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=continuous_batching_config ) # Generated with a regular generate in bf16, without CB nor upcast lm_head expected_texts = Expectations( { ("cuda", (10, 0)): ["The robe takes 2 bolts", "The basket contains 25 oranges"], # Default expectation was not tested but needs to be there to avoid crashes. Feel free to change. (None, None): ["The robe takes 2 bolts", "The basket contains 25 oranges"], } ).get_expectation() self.assertEqual(len(outputs), len(expected_texts)) for i, expected_text in enumerate(expected_texts): text = tokenizer.decode(outputs[f"req_{i}"].generated_tokens, skip_special_tokens=True) self.assertEqual(text, expected_text, f"Request {i} mismatch") @parameterized.expand([("flash_attention_2",), ("flash_attention_3",)]) @slow def test_decode_fast_path_wide_batch_parity(self, attn_implementation: str) -> None: """Decode-fast-path output must match varlen when more requests decode concurrently than `max_blocks_per_request` (regression test for the `pad_to_pow2` cap truncating the decode batch).""" # Skip the test if Flash Attention 2 or 3 is required but not available. flexible_flash_skip(self, attn_implementation) model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" tokenizer, model = get_tokenizer_and_model(model_id, attn_implementation, torch_device, torch.bfloat16) # 12 requests but only 4 blocks per request: the decode batch is wider than max_blocks_per_request input_ids = get_generation_inputs(_DEFAULT_USER_MESSAGES * 4, tokenizer, for_continuous_batching=True) gen_config = GenerationConfig(do_sample=False, max_new_tokens=20) # CUDA graphs enable input padding, which is where the truncation happened cb_config = ContinuousBatchingConfig(page_size=256, num_blocks=64, use_cuda_graph=True) cb_config.max_blocks_per_request = 0 # varlen reference outputs_varlen = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=cb_config ) cb_config.max_blocks_per_request = 4 # decode fast path, narrower than the batch outputs_fast = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=cb_config ) self.assertEqual(len(outputs_varlen), len(outputs_fast)) for (_, out_varlen), (_, out_fast) in zip(outputs_varlen.items(), outputs_fast.items()): text_varlen = tokenizer.decode(out_varlen.generated_tokens, skip_special_tokens=True) text_fast = tokenizer.decode(out_fast.generated_tokens, skip_special_tokens=True) self.assertEqual(text_varlen, text_fast, f"Mismatch:\nvarlen: {text_varlen}\nfast: {text_fast}") @parameterized.expand([(False, False), (False, True), (True, False), (True, True)]) @slow def test_per_request_logits_processors(self, use_cuda_graph: bool, use_async_batching: bool) -> None: """Tests that per-request logits processor kwargs (temperature, top_k, top_p) work correctly in generation.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" max_new_tokens = 10 temperatures = [1.0, 1.0] top_ks = [10, 50] top_ps = [0.9, 0.99] # Run in fp32: in bf16, CB and regular generate logprobs differ by ~1e-3, far above the tolerance used below tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.float32) eos_token_id = model.config.eos_token_id # type: ignore[attr-defined] # Same prompt for both requests user_messages = ["Write a random number:"] input_ids = get_generation_inputs(user_messages, tokenizer, for_continuous_batching=True)[0] # Use the context manager to add requests with different per-request kwargs generation_config = GenerationConfig( do_sample=True, temperature=max(temperatures) + 1, # enables temperature warping top_k=max(top_ks) + 1, top_p=min(top_ps) - 0.01, max_new_tokens=max_new_tokens, eos_token_id=eos_token_id, ) continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching, per_request_processors=True, return_logprobs=True, q_padding_interval_size=16, # allows for exact comparison between CB and regular generation auto_switch_to_flash=False, ) manager = model.init_continuous_batching( generation_config=generation_config, continuous_batching_config=continuous_batching_config, ) # Trick to have temperature, top-k, top-p ... without randomness: diable sampling after manager creation manager.generation_config.do_sample = False manager.start() try: # Request 0: low temperature (more deterministic) req0_id = manager.add_request( input_ids, max_new_tokens=max_new_tokens, temperature=temperatures[0], top_k=top_ks[0], top_p=top_ps[0] ) # Request 1: high temperature (more random) req1_id = manager.add_request( input_ids, max_new_tokens=max_new_tokens, temperature=temperatures[1], top_k=top_ks[1], top_p=top_ps[1] ) # Collect results results = {} while len(results) < 2: result = manager.get_result(timeout=1) if result is not None and result.is_finished(): self.assertEqual(result.status, RequestStatus.FINISHED) # fail if the request failed results[result.request_id] = result elif not manager.is_running(): break finally: manager.stop(block=True) # Both requests should complete and have logprobs self.assertEqual(len(results), 2, f"Expected 2 results, got {len(results)}") self.assertGreater(len(results[req0_id].logprobs), 0) self.assertGreater(len(results[req1_id].logprobs), 0) # Also ensure the logprobs were not the same self.assertNotEqual(results[req0_id].logprobs, results[req1_id].logprobs) # Compare each request with regular generation # Build logits processor with do_sample=True (so temperature is included), then set do_sample=False for # deterministic generation, which is the same trick that CB uses delta = 2e-5 if use_cuda_graph else 1e-5 for i, req_id in enumerate([req0_id, req1_id]): tokenizer, model = get_tokenizer_and_model(model_id, "sdpa", torch_device, torch.float32) gen_config = GenerationConfig( do_sample=True, temperature=temperatures[i], top_k=top_ks[i], top_p=top_ps[i], max_new_tokens=max_new_tokens, eos_token_id=eos_token_id, ) logits_processor = model._get_logits_processor(gen_config) gen_config.do_sample = False regular_generated_tokens, regular_logprobs = regular_generate( model=model, tokenizer=tokenizer, user_messages=user_messages, logits_processor=logits_processor, max_new_tokens=max_new_tokens, do_sample=False, output_scores=True, eos_token_id=eos_token_id, ) self.assertEqual(results[req_id].generated_tokens, regular_generated_tokens[0]) for j, (cb_lp, exp_lp) in enumerate(zip(results[req_id].logprobs, regular_logprobs[0])): error_msg = f"Request {i}: logprob mismatch at position {j}: CB={cb_lp}, expected={exp_lp}" self.assertAlmostEqual(cb_lp, exp_lp, delta=delta, msg=error_msg) # ---------------------------------- CPU offloading tests ---------------------------------- # @require_torch_accelerator def test_cpu_offloading_parity(self) -> None: """Test that CPU offloading produces the same results as the legacy soft-reset path, and that it is actually called at least once. Uses a very small cache (few blocks) to force offloading.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # 4 blocks of 32 tokens: small enough to force offloading continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=True, allow_block_sharing=True, use_async_batching=False, num_blocks=4, page_size=32, cpu_offload_space=1.0, ) original_offload = OffloadingManager._offload_to_cpu with patch.object( OffloadingManager, "_offload_to_cpu", autospec=True, side_effect=original_offload ) as mock_offload: self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", max_new_tokens=30, num_repeat_prompts=4, ) self.assertTrue(mock_offload.called, "_offload_to_cpu was not called despite few blocks being available.") @require_torch_accelerator def test_cpu_offloading_parity_async(self) -> None: """Same as test_cpu_offloading_parity but with async batching, where offloading can evict requests that are in flight in the previous batch, exercising the rollback-on-restore path.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # 4 blocks of 32 tokens: small enough to force offloading continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=True, allow_block_sharing=True, use_async_batching=True, num_blocks=4, page_size=32, cpu_offload_space=1.0, ) # Spy on _offload_to_cpu to check it ran and that at least one victim was in flight (rollback path) original_offload = OffloadingManager._offload_to_cpu in_flight_victim_seen = False def spy_offload(manager, victims, victim_block_tables): nonlocal in_flight_victim_seen in_flight_victim_seen |= any( state.position_offset == len(state.initial_tokens) + len(state.generated_tokens) for state in victims ) return original_offload(manager, victims, victim_block_tables) with patch.object(OffloadingManager, "_offload_to_cpu", autospec=True, side_effect=spy_offload) as mock: self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", max_new_tokens=30, num_repeat_prompts=4, ) self.assertTrue(mock.called, "_offload_to_cpu was not called despite few blocks being available.") self.assertTrue(in_flight_victim_seen, "No in-flight victim was offloaded: rollback path not exercised.") @require_torch_accelerator def test_cpu_offloading_disabled_when_zero(self) -> None: """Test that cpu_offload_space=0 produces the same output as the legacy path.""" model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # 4 blocks of 32 tokens: small enough to force offloading continuous_batching_config = ContinuousBatchingConfig( use_cuda_graph=True, allow_block_sharing=True, use_async_batching=False, num_blocks=4, page_size=32, cpu_offload_space=0.0, ) # Should work identically to the existing test_continuous_batching_few_blocks self._test_continuous_batching_parity( model_id=model_id, continuous_batching_config=continuous_batching_config, attn_implementation="sdpa", max_new_tokens=30, num_repeat_prompts=4, ) @require_torch_accelerator class TestMemoryHandlerPrediction(unittest.TestCase): """Verifies that ``PagedAttentionMemoryHandler.compute_memory_footprint`` matches real accelerator memory usage. For each configuration we allocate tensors at the *idealized* sizes modeled by the handler (same shapes, same dtypes, no alignment padding) and compare the CUDA memory delta to the handler's prediction. The handler reserves the largest of the two M-proportional activation peaks (LM head and attention) and always accounts for the cache reads (N-proportional), so we allocate the dominant M peak plus the cache read tensors. """ NUM_BLOCKS = 4 MAX_BATCH_TOKENS = 64 # Each tuple fully specifies a synthetic model config plus the CB knobs; block_size = head_dim * num_kv_heads. # fmt: off # (block_size, head_dim, num_kv_heads, num_attention_heads, hidden_size, vocab_size, group_types, group_size, attn_impl, max_bpr, logprobs, dtype, use_async) CONFIGS = [ (32, 64, 4, 8, 512, 32000, ["full_attention"], 22, "sdpa", 0, False, torch.float16, False), # sdpa-like, 1 mask (256, 64, 4, 8, 512, 32000, ["full_attention"], 22, "flash_attention_2", 0, False, torch.float16, False), # flash-like, no mask (32, 64, 4, 8, 512, 32000, ["full_attention", "sliding_attention"], 14, "sdpa", 0, False, torch.bfloat16, False), # hybrid, 2 groups + 2 masks (32, 64, 2, 8, 512, 32000, ["full_attention"], 16, "sdpa", 8, True, torch.float16, False), # block_table + logprobs (32, 64, 2, 8, 512, 32000, ["full_attention"], 16, "sdpa", 8, True, torch.float16, True), # block_table + logprobs + async ] # fmt: on @parameterized.expand(CONFIGS) def test_memory_prediction( self, block_size: int, head_dim: int, num_kv_heads: int, num_attention_heads: int, hidden_size: int, vocab_size: int, group_types: list[str], group_size: int, attn_impl: str, max_bpr: int, logprobs: bool, dtype: torch.dtype, use_async: bool, ) -> None: config = SimpleNamespace( # rather than creating a full config that we would only use for the namespace head_dim=head_dim, num_key_value_heads=num_kv_heads, num_attention_heads=num_attention_heads, hidden_size=hidden_size, vocab_size=vocab_size, _attn_implementation=attn_impl, ) cb_config = ContinuousBatchingConfig( page_size=block_size, num_blocks=self.NUM_BLOCKS, max_blocks_per_request=max_bpr, return_logprobs=logprobs, use_async_batching=use_async, max_memory_percent=0.9, ) # Compute the sector geometry the same way PagedAttentionCache does, with the layers evenly split across the # attention types layers_per_group = group_size // len(group_types) bytes_per_token = 2 * num_kv_heads * head_dim * dtype.itemsize bytes_per_block = bytes_per_token * block_size * layers_per_group bytes_per_sector = math.lcm(bytes_per_block, 128) tokens_per_sector = bytes_per_sector // bytes_per_block * block_size bytes_per_model_block = bytes_per_token * block_size * group_size handler = PagedAttentionMemoryHandler( config=config, cb_config=cb_config, dtype=dtype, bytes_per_sector=bytes_per_sector, tokens_per_sector=tokens_per_sector, bytes_per_block=bytes_per_model_block, attn_types=group_types, ) num_groups = len(group_types) num_attn_masks = handler.num_attention_masks num_output_rows = 2 if logprobs else 1 elems_per_kv_token = head_dim * num_kv_heads # K/V read back from the cache: flash reads GQA heads natively, eager/sdpa repeat K and V while V is referenced kv_read_heads = 2 * num_kv_heads if "flash" in attn_impl else 2 * num_attention_heads + num_kv_heads elems_per_kv_read_token = head_dim * kv_read_heads q_dim = num_attention_heads * head_dim # NUM_BLOCKS converts to a number of sectors, and reads are sized by the readable tokens those sectors hold num_sectors = handler.num_sectors_from_config() if num_sectors is None: raise ValueError("handler could not find num_sector, which should be impossible in this test.") N = num_sectors * tokens_per_sector # max number of readable cache tokens M = self.MAX_BATCH_TOKENS k = handler.io_multiplier # 1 sync, 2 async -- scales IO tensors only predicted = handler.compute_memory_footprint(M, num_sectors) # -- Allocate tensors at the exact idealized sizes the handler models -- device = torch_device backend_empty_cache(device) baseline = backend_memory_allocated(device) # Tensors present regardless of which activation peak is live fixed = [] # kv_cache: a single flat byte tensor holding the data sectors plus the two trash sectors (not scaled by k) fixed.append(torch.empty((num_sectors + 2) * bytes_per_sector, dtype=torch.uint8, device=device)) # IO tensors below are allocated k times (once per IO instance) for _ in range(k): fixed.append(torch.empty((7, M), dtype=torch.int32, device=device)) # bulk_input fixed.append(torch.empty((num_output_rows, M), dtype=torch.int32, device=device)) # output_ids for _ in range(num_attn_masks): # attention_mask: [1, 1, M, N + M] per mask type fixed.append(torch.empty((1, 1, M, N + M), dtype=dtype, device=device)) if max_bpr > 0: # block_table: [num_groups, M, max_bpr] (skipped when max_bpr == 0) fixed.append(torch.empty((num_groups, M, max_bpr), dtype=torch.int32, device=device)) fixed.append(torch.empty((num_groups, M), dtype=torch.int64, device=device)) # write_index fixed.append(torch.empty((num_groups, N + M), dtype=torch.int64, device=device)) # read_index # Old K/V read from the whole readable cache: always reserved along with the sectors (not scaled by k) fixed.append(torch.empty((N, elems_per_kv_read_token), dtype=dtype, device=device)) # M-proportional activation peaks: only one is live at a time, so the footprint reserves the largest peaks = { # LM head: hidden states [M, hidden] turned into logits [M, vocab] (always fp32) "lm_head": [ torch.empty((M, hidden_size), dtype=dtype, device=device), torch.empty((M, vocab_size), dtype=torch.float32, device=device), ], # Attention: hidden + Q + new K/V over M, plus their copy in the K/V read back from the cache "attention": [ torch.empty((M, hidden_size), dtype=dtype, device=device), torch.empty((M, q_dim), dtype=dtype, device=device), torch.empty((M, elems_per_kv_token), dtype=dtype, device=device), torch.empty((M, elems_per_kv_token), dtype=dtype, device=device), torch.empty((M, elems_per_kv_read_token), dtype=dtype, device=device), ], } peak_nbytes = {name: sum(t.nbytes for t in ts) for name, ts in peaks.items()} dominant = max(peak_nbytes, key=peak_nbytes.get) # Free the non-dominant peak so the accelerator delta reflects only the live one for name in [n for n in peaks if n != dominant]: del peaks[name] actual_accelerator = backend_memory_allocated(device) - baseline expected_nbytes = sum(t.nbytes for t in fixed) + peak_nbytes[dominant] num_allocations = len(fixed) + len(peaks[dominant]) # The CUDA caching allocator rounds allocations larger than 10MB up to a multiple of 2MB large_alloc_rounding = sum( -t.nbytes % (2 * 1024**2) for t in fixed + peaks[dominant] if t.nbytes > 10 * 1024**2 ) del fixed, peaks backend_empty_cache(device) # 1) Exact check: prediction must equal the sum of tensor nbytes. This validates the handler's cost # coefficients against the tensor shapes, with zero tolerance. self.assertEqual( predicted, expected_nbytes, f"Prediction ({predicted}) != sum of tensor nbytes ({expected_nbytes})", ) # 2) GPU memory check, which accounts for the CUDA caching allocator's rounding up of allocations. if baseline > 0: self.skipTest(f"CUDA allocator already holds {baseline} bytes: the memory delta check would be polluted") max_accelerator_overhead = num_allocations * 512 + large_alloc_rounding self.assertLessEqual( abs(actual_accelerator - predicted), max_accelerator_overhead, f"Accelerator delta ({actual_accelerator}) too far from prediction ({predicted}), " f"allowed overhead = {max_accelerator_overhead} ({num_allocations} allocs × 512B)", ) # Worker functions for the TP continuous batching tests, spawned through `_init_distributed`. def _tp_continuous_batching_worker( rank: int, model_id: str, attn_implementation: str, max_new_tokens: int, do_sample: bool, seed: int, use_cuda_graph: bool, use_async_batching: bool, ) -> None: """Loads `model_id` with `DistributedConfig(tp_size=...)`, checks three TP-specific paths in the same process: (a) direct broadcasts via `DistributedHelper`, (b) per-rank parity of CB-generated tokens via `dist.all_gather_object`, and (c) reproducibility across two CB runs sharing the same seed. Rank 0 owns all the assertions; the other ranks only need to participate in the collectives.""" import torch import torch.distributed as dist from transformers.distributed import DistributedConfig from transformers.generation.continuous_batching.distributed import DistributedHelper tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") if not hasattr(tokenizer, "pad_token") and hasattr(tokenizer, "eos_token"): tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_id, attn_implementation=attn_implementation, distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), dtype="auto", ).eval() # Direct broadcast tests: only rank 0's value should propagate to every TP rank helper = DistributedHelper(device_mesh=model._device_mesh, cpu_group_timeout=300, tp_plan=model._tp_plan) received_obj = helper.tp_broadcast_object_from_rank_0({"src_rank": rank}) assert received_obj == {"src_rank": 0}, f"tp_broadcast_object: rank {rank} got {received_obj}" sent_tensor = torch.tensor([float(rank)], device=model.device) helper.tp_broadcast_from_rank_0(sent_tensor) assert sent_tensor.item() == 0.0, f"tp_broadcast_from_rank_0: rank {rank} got {sent_tensor.item()}" # CB runs: same seed twice, assert reproducibility AND cross-rank parity user_messages = [ "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?" ] chats = [[{"role": "user", "content": m}] for m in user_messages] tokenized = [tokenizer.apply_chat_template(chat, add_generation_prompt=True) for chat in chats] input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized] cb_config_kwargs = {"use_cuda_graph": use_cuda_graph, "use_async_batching": use_async_batching, "seed": seed} gen_config = GenerationConfig(do_sample=do_sample, max_new_tokens=max_new_tokens) first_outputs = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=ContinuousBatchingConfig(**cb_config_kwargs), ) second_outputs = model.generate_batch( inputs=input_ids, generation_config=gen_config, continuous_batching_config=ContinuousBatchingConfig(**cb_config_kwargs), ) # Cross-rank parity: every TP rank must produce the same tokens, otherwise the seed broadcast / TP collectives are # diverging silently. Gather the first run's tokens onto all ranks and let rank 0 compare. local_tokens = [out.generated_tokens for out in first_outputs.values()] gathered_tokens = [None] * helper.tp_size dist.all_gather_object(gathered_tokens, local_tokens, group=helper.tp_group) if rank != 0: return assert len(first_outputs) == len(input_ids), f"Expected {len(input_ids)} CB outputs, got {len(first_outputs)}" for i, (_, output) in enumerate(first_outputs.items()): assert len(output.generated_tokens) > 0, f"Request {i} got no generated tokens" for src_rank, src_tokens in enumerate(gathered_tokens): if src_tokens != gathered_tokens[0]: raise AssertionError( f"TP continuous batching diverges across ranks: rank {src_rank} got {src_tokens}, rank 0 got " f"{gathered_tokens[0]}" ) second_tokens = [out.generated_tokens for out in second_outputs.values()] if local_tokens != second_tokens: raise AssertionError( f"TP continuous batching is not reproducible across runs with the same seed\n" f"First run : {local_tokens}\n" f"Second run: {second_tokens}" ) def _tp_cancellation_worker( rank: int, model_id: str, attn_implementation: str, use_cuda_graph: bool = False, use_async_batching: bool = False, ) -> None: """Loads `model_id` with `DistributedConfig(tp_size=...)`, submits a long-running streaming request, and cancels it mid-flight. The cancellation goes through the cancel-queue + `tp_broadcast_object` path: if the broadcast were broken, the non-driver rank's scheduler would not learn about the cancellation and the test would hang or crash on the next TP forward pass. Rank 0 owns the assertions.""" import time from transformers.distributed import DistributedConfig cb_config = ContinuousBatchingConfig(use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching) tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") if not hasattr(tokenizer, "pad_token") and hasattr(tokenizer, "eos_token"): tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_id, attn_implementation=attn_implementation, distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), dtype="auto", ).eval() chat = [{"role": "user", "content": "Tell me a long story about a robot exploring the galaxy."}] tokenized = tokenizer.apply_chat_template(chat, add_generation_prompt=True) inputs = tokenized if isinstance(tokenized, list) else tokenized["input_ids"] max_new_tokens = 200 cancel_after_n_chunks = 3 manager = model.init_continuous_batching(continuous_batching_config=cb_config) manager.logit_processor.clear() # Warm up synchronously so CUDA-graph capture doesn't eat the streaming-loop deadline below manager.warmup() manager.start() try: request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=True) chunks_seen = 0 cancelled = False deadline = time.time() + 60 while time.time() < deadline: chunk = manager.get_result(request_id=request_id, timeout=2.0) if chunk is None: # No new chunks for 2s after cancel — cancellation took effect on every rank break chunks_seen += 1 if chunks_seen >= cancel_after_n_chunks and not cancelled: manager.cancel_request(request_id) cancelled = True if rank == 0: assert cancelled, "Test setup did not reach the cancel call" assert chunks_seen < max_new_tokens, ( f"Cancellation did not stop generation early: saw {chunks_seen} chunks " f"for max_new_tokens={max_new_tokens}" ) finally: manager.stop(block=True) def _tp_pause_generation_worker( rank: int, model_id: str, attn_implementation: str, max_new_tokens: int, num_pauses: int, skew_seconds: float, use_async_batching: bool = False, use_cuda_graph: bool = False, ) -> None: """Loads `model_id` with `DistributedConfig(tp_size=...)` and runs a generation loop on every rank, plus a second thread that repeatedly pauses it. Each rank starts pausing at a different time, so the ranks reach the pause window at different iterations. Since the pause request is MAX-reduced across ranks, a rank ends up paused before its own thread asked for it, which a single-rank test cannot reproduce. Inside the pause, every rank issues a collective on the TP group. If a rank resumed generation instead of staying paused, the collectives of its forward pass would pair up with this probe on the other ranks, and the probe would return a wrong value or hang. Rank 0 owns the assertions.""" import threading import time import torch import torch.distributed as dist from transformers.distributed import DistributedConfig tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left") if not hasattr(tokenizer, "pad_token") and hasattr(tokenizer, "eos_token"): tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_id, attn_implementation=attn_implementation, distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])), dtype="auto", ).eval() chats = [[{"role": "user", "content": message}] for message in _DEFAULT_USER_MESSAGES] tokenized = [tokenizer.apply_chat_template(chat, add_generation_prompt=True) for chat in chats] input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized] # A dedicated group for the test's own synchronization: the manager's groups are used by the generation loop # thread, and issuing collectives on them from the main thread at the same time would interleave. sync_group = dist.new_group(backend="gloo") cb_config = ContinuousBatchingConfig( use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching, cpu_group_timeout=120 ) manager = model.init_continuous_batching(continuous_batching_config=cb_config) manager.logit_processor.clear() # Warm up synchronously so that no CUDA-graph capture happens on the loop thread during the test manager.warmup() manager.start() tp_group = manager.distributed_helper.tp_group tp_size = manager.distributed_helper.tp_size expected_probe = float(sum(range(1, tp_size + 1))) pauses_taken = [] pause_errors = [] def pauser() -> None: """Stands in for a trainer thread updating the model between generation steps.""" try: time.sleep(skew_seconds * rank) # so that the ranks do not ask for the pause at the same iteration for _ in range(num_pauses): with manager.pause(): # A collective on the TP group, like a training step would issue. It only returns the right value # if every rank is in the same pause window. probe = torch.full((16,), float(rank + 1), device=model.device) dist.all_reduce(probe, group=tp_group) if probe[0].item() != expected_probe: raise AssertionError( f"TP collective inside the pause returned {probe[0].item()}, expected {expected_probe}: " f"the ranks are not in the same pause window" ) pauses_taken.append(1) time.sleep(0.02) # let the generation loop progress between pauses except Exception as error: # re-raised on the main thread below pause_errors.append(error) pause_thread = threading.Thread(target=pauser) pause_thread.start() try: # `add_request` only enqueues on the TP driver and returns None on the other ranks. The requests are then # broadcast, so every rank produces outputs for the driver's request ids: collect by count and key on the ids # that come back. for ids in input_ids: manager.add_request(ids, max_new_tokens=max_new_tokens, streaming=False) outputs = {} deadline = time.time() + 300 while len(outputs) < len(input_ids) and time.time() < deadline: result = manager.get_result(timeout=1.0) if result is not None and result.is_finished(): outputs[result.request_id] = result pause_thread.join(timeout=120) finally: # The pause thread must be done before stopping: `stop` joins the generation thread, which cannot exit while a # pause is held. Also, no rank may stop while another rank is still pausing, since the stop goes through the # same all-reduce as the pause and the late rank would be left waiting on a loop that no longer exists. if not pause_thread.is_alive(): dist.barrier(group=sync_group) manager.stop(block=True) assert not pause_thread.is_alive(), f"rank {rank}: the pausing thread did not finish" if pause_errors: # `is_running()` does not distinguish a dead loop from a stopped one, so chain the fatal error if there is one loop_error = manager.background_thread_status.fatal_error raise pause_errors[0] from loop_error assert len(pauses_taken) == num_pauses, f"rank {rank}: took {len(pauses_taken)} pauses, expected {num_pauses}" assert len(outputs) == len(input_ids), f"rank {rank}: got {len(outputs)}/{len(input_ids)} requests back" # Pausing must not change the generated tokens nor make the ranks diverge. Sorting by request id gives every rank # the same order. local_tokens = [outputs[request_id].generated_tokens for request_id in sorted(outputs)] gathered_tokens = [None] * tp_size dist.all_gather_object(gathered_tokens, local_tokens, group=tp_group) if rank != 0: return for index, tokens in enumerate(local_tokens): assert len(tokens) > 0, f"Request {index} got no generated tokens" for src_rank, src_tokens in enumerate(gathered_tokens): if src_tokens != gathered_tokens[0]: raise AssertionError( f"Generation diverges across ranks when pausing: rank {src_rank} got {src_tokens}, rank 0 got " f"{gathered_tokens[0]}" ) @require_torch_multi_accelerator class ContinuousBatchingTensorParallelTest(unittest.TestCase): """Integration tests for continuous batching with tensor parallelism. Each test spawns a TP-sized process group via `_init_distributed` (see `tests/test_tensor_parallel_mixin.py`) with the NCCL backend.""" @property def tp_size(self) -> int: return min(torch.cuda.device_count(), 2) def _run_cb_worker(self, max_new_tokens: int = 20, **worker_kwargs) -> None: """Spawn `_tp_continuous_batching_worker` on `tp_size` NCCL processes with sensible defaults.""" defaults = { "model_id": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "attn_implementation": "flash_attention_2", "max_new_tokens": max_new_tokens, "do_sample": False, "seed": 42, "use_cuda_graph": False, "use_async_batching": False, } defaults.update(worker_kwargs) flexible_flash_skip(self, defaults["attn_implementation"]) _init_distributed(tp=self.tp_size, backend="nccl")(_tp_continuous_batching_worker)(**defaults) def test_continuous_batching_tp_fast(self) -> None: """Test that continuous batching with `DistributedConfig(tp_size=...)` produces non-empty, reproducible greedy outputs and that all TP ranks agree on the generated tokens.""" self._run_cb_worker(max_new_tokens=4) @require_flash_attn def test_continuous_batching_tp_pause(self) -> None: """Test that `pause` keeps the TP ranks in the same pause window even when they request it at different iterations, and that pausing repeatedly mid-generation loses no request and does not make the ranks diverge.""" _init_distributed(tp=self.tp_size, backend="nccl")(_tp_pause_generation_worker)( model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="flash_attention_2", max_new_tokens=20, num_pauses=5, skew_seconds=0.25, ) @slow def test_continuous_batching_tp_greedy(self) -> None: """Test that continuous batching with `DistributedConfig(tp_size=...)` produces non-empty, reproducible greedy outputs and that all TP ranks agree on the generated tokens.""" self._run_cb_worker() @slow def test_continuous_batching_tp_with_sampling(self) -> None: """Test that continuous batching with TP and sampling is reproducible across runs with the same seed and that all TP ranks agree on the sampled tokens — implicitly validating the seed broadcast from rank 0.""" self._run_cb_worker(do_sample=True, seed=123) @slow def test_continuous_batching_tp_with_cuda_graph(self) -> None: """Test that continuous batching with TP and CUDA graphs is reproducible across runs and that all TP ranks agree on the generated tokens — captured-graph collectives must stay in sync across ranks.""" self._run_cb_worker(use_cuda_graph=True) @slow def test_continuous_batching_tp_with_cuda_graph_and_async(self) -> None: """Test that continuous batching with TP, CUDA graphs, and async batching is reproducible across runs and that all TP ranks agree on the generated tokens — the toughest combination, exercising both captured-graph collectives and the async producer/consumer split.""" self._run_cb_worker(use_cuda_graph=True, use_async_batching=True) @slow @require_flash_attn def test_continuous_batching_tp_cancellation(self) -> None: """Test that `cancel_request` propagates across the TP group: the driver enqueues the cancellation, broadcasts it to non-driver ranks via `tp_broadcast_object`, and generation stops well before `max_new_tokens`.""" _init_distributed(tp=self.tp_size, backend="nccl")(_tp_cancellation_worker)( model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="flash_attention_2", ) @slow @require_flash_attn def test_continuous_batching_tp_cancellation_realistic(self) -> None: """Test that `cancel_request` propagates across the TP group: the driver enqueues the cancellation, broadcasts it to non-driver ranks via `tp_broadcast_object`, and generation stops well before `max_new_tokens`.""" _init_distributed(tp=self.tp_size, backend="nccl")(_tp_cancellation_worker)( model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="flash_attention_2", use_async_batching=True, use_cuda_graph=True, )