Repository: huggingface/transformers
File: src/transformers/generation/continuous_batching/continuous_api.py (Python)

Code before:
        self.kv_padding_intervals = kv_padding_intervals
        self._graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] | None = {} if use_cuda_graph else None
        # Compile-related arguments
        self.compile_config: CompileConfig | None = getattr(generation_config, "compile_config", None)
        self._forward_process_and_sample_is_compiled = False

        self._pad_inputs = use_cuda_graph or (self.compile_config is not None and not self.compile_config.dynamic)

        # Set up metrics collector
        self.max_batch_tokens = cache.max_batch_tokens
        self.metrics = ContinuousBatchProcessorMetrics(cache.max_batch_tokens)

        # Setup static tensors
        self.actual_query_length = 0  # This is the actual number of queries tokens in the batch
        self.actual_key_length = 0  # This is the actual number of keys/values tokens in the batch
        self.actual_batch_size = 0  # This is the actual number of requests in the batch
        self.actual_index_sizes = [(0, 0) for _ in range(cache.num_groups)]
        self.setup_static_tensors(cache.num_groups)

    @traced(standalone=True)
    def setup_static_tensors(self, num_groups: int) -> None:
        """Setup the static tensors that are used for storage during the generation step. No other tensor will be
        allowed for the inputs or the outputs of the generation step."""
        self.num_pages = self.cache.num_blocks * self.cache.block_size
        self.tensor_metadata = {"dtype": torch.int32, "device": self.model_device}

        # Some tensors always have the same shape regardless of the model
        self.input_ids = torch.empty((1, self.max_batch_tokens), dtype=torch.int32, device=self.model_device)
        self.position_ids = torch.empty((1, self.max_batch_tokens), dtype=torch.int32, device=self.model_device)
        self.cumulative_seqlens_q = torch.empty(
            (self.max_batch_tokens + 1,), dtype=torch.int32, device=self.model_device
        )
        self.max_seqlen_q = 0
        self.logits_indices = torch.empty((self.max_batch_tokens,), dtype=torch.int32, device=self.model_device)
        self.output_ids = torch.empty((self.max_batch_tokens,), dtype=torch.int32, device=self.model_device)

        # For some kwargs, we have a dict of tensors with as many items as there are attention types
        layer_types = getattr(self.config, "layer_types", None)
        if layer_types is None:
            sliding_window = getattr(self.config, "sliding_window", 1)
            layer_types = ["full_attention"] if sliding_window in [1, None] else ["sliding_attention"]
        layer_types = list(set(layer_types))

        self.cumulative_seqlens_k = {}
        for l_type in layer_types:
            self.cumulative_seqlens_k[l_type] = torch.empty(
                (self.max_batch_tokens + 1), dtype=torch.int32, device=self.model_device
            )
        self.max_seqlen_k = dict.fromkeys(layer_types, 0)

        if attn_mask_is_needed(self.config):

Review comment (suggestion):
you no longer use these no? you could type them as type dict if you still wanted to use them 

Code after:
        self.kv_padding_intervals = kv_padding_intervals
        self._graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] | None = {} if use_cuda_graph else None
        # Compile-related arguments
        self.compile_config: CompileConfig | None = getattr(generation_config, "compile_config", None)
        self._forward_process_and_sample_is_compiled = False

        self._pad_inputs = use_cuda_graph or (self.compile_config is not None and not self.compile_config.dynamic)

        # Set up metrics collector
        self.max_batch_tokens = cache.max_batch_tokens
        self.metrics = ContinuousBatchProcessorMetrics(cache.max_batch_tokens)

        # Setup static tensors
        self.actual_query_length = 0  # This is the actual number of queries tokens in the batch
        self.actual_key_length = 0  # This is the actual number of keys/values tokens in the batch
        self.actual_batch_size = 0  # This is the actual number of requests in the batch
        self.actual_index_sizes = [(0, 0) for _ in range(cache.num_groups)]
        self.setup_static_tensors(cache.num_groups)

    @traced(standalone=True)
    def setup_static_tensors(self, num_groups: int) -> None:
        """Setup the static tensors that are used for storage during the generation step. No other tensor will be
        allowed for the inputs or the outputs of the generation step."""
        self.num_pages = self.cache.num_blocks * self.cache.block_size

        # Some tensors always have the same shape regardless of the model
        self.input_ids = torch.empty((1, self.max_batch_tokens), dtype=torch.int32, device=self.model_device)
        self.position_ids = torch.empty((1, self.max_batch_tokens), dtype=torch.int32, device=self.model_device)
        self.cumulative_seqlens_q = torch.empty(
            (self.max_batch_tokens + 1,), dtype=torch.int32, device=self.model_device
        )
        self.max_seqlen_q = 0
        self.logits_indices = torch.empty((self.max_batch_tokens,), dtype=torch.int32, device=self.model_device)
        self.output_ids = torch.empty((self.max_batch_tokens,), dtype=torch.int32, device=self.model_device)

        # For some kwargs, we have a dict of tensors with as many items as there are attention types
        layer_types = getattr(self.config, "layer_types", None)
        if layer_types is None:
            sliding_window = getattr(self.config, "sliding_window", 1)
            layer_types = ["full_attention"] if sliding_window in [1, None] else ["sliding_attention"]
        layer_types = list(set(layer_types))

        self.cumulative_seqlens_k = {}
        for l_type in layer_types:
            self.cumulative_seqlens_k[l_type] = torch.empty(
                (self.max_batch_tokens + 1), dtype=torch.int32, device=self.model_device
            )
        self.max_seqlen_k = dict.fromkeys(layer_types, 0)

        if attn_mask_is_needed(self.config):
            self.attention_mask = {}