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

Code before:
        b_size = padded_q_size if use_padding else self.actual_batch_size
        # If there is padding, the size of the KV is the nb of padded Q tokens + the size padded of the padded KV cache
        padded_kv_size = padded_q_size + padded_kv_cache_size

        # Prepare the kwargs, the attributes that are either tensors or dict of tensors are initialized to empty dicts
        kwargs = PagedAttentionArgs(
            input_ids=self.input_ids[:, :q_len],
            position_ids=self.position_ids[:, :q_len],
            cu_seq_lens_q=self.cumulative_seqlens_q[: b_size + 1],
            max_seqlen_q=self.max_seqlen_q,
            logits_indices=self.logits_indices[:q_len],
            cu_seq_lens_k={},
            max_seqlen_k={},
            attention_mask={},
            read_index=[],
            write_index=[],
            cache=self.cache,
            use_cache=False,
        )

        # If we use constant-sized slicing, there are some "padding" queries tokens which FA has some issues with. In
        # some models like Qwen3-4B-Instruct-2507, if we don't include these tokens in cumulative_seqlens_q, there are
        # some NaNs in the output logits even for non-padded tokens.
        if use_padding:
            self.max_seqlen_q = max(self.max_seqlen_q, q_len - self.total_seqlen_q)
            kwargs["max_seqlen_q"] = self.max_seqlen_q  # Update kwargs with the new value
            self.cumulative_seqlens_q[self.actual_batch_size + 1 :] = q_len
            # FIXME: is there another way to avoid this? It has a very slight impact on performance (~5 tok/s)

        # For the attributes that are lists of tensors, we construct list of tensor references
        for i, (read_index_size, write_index_size) in enumerate(self.actual_index_sizes):
            read_index_size = padded_kv_size if use_padding else read_index_size
            write_index_size = padded_q_size if use_padding else write_index_size
            kwargs.read_index.append(self.read_index_storage[i][:read_index_size])
            kwargs.write_index.append(self.write_index_storage[i][:write_index_size])

        # For the attributes that are dict of tensors, we replace the dict with a tensor if there is only one entry
        layer_types = list(self.cumulative_seqlens_k.keys())
        if len(layer_types) > 1:
            kwargs.max_seqlen_k: dict[str, int] = {}
            kwargs.cu_seq_lens_k: dict[str, torch.Tensor] = {}
            kwargs.attention_mask: dict[str, torch.Tensor] = {}
            for layer_type, seqlens_k in self.cumulative_seqlens_k.items():
                kwargs.cu_seq_lens_k[layer_type] = seqlens_k[: b_size + 1]
                kwargs.max_seqlen_k[layer_type] = self.max_seqlen_k[layer_type]
                if self.attention_mask is not None:
                    k_len = padded_kv_size if use_padding else seqlens_k[b_size]
                    kwargs.attention_mask[layer_type] = self.attention_mask[layer_type][..., :q_len, :k_len]
        else:
            layer_type = layer_types[0]
            kwargs.cu_seq_lens_k = self.cumulative_seqlens_k[layer_type][: b_size + 1]

Review comment (suggestion):
```suggestion
            kwargs.max_seqlen_q = self.max_seqlen_q  # Update kwargs with the new value
```

Code after:
        b_size = padded_q_size if use_padding else self.actual_batch_size
        # If there is padding, the size of the KV is the nb of padded Q tokens + the size padded of the padded KV cache
        padded_kv_size = padded_q_size + padded_kv_cache_size

        # Prepare the kwargs, the attributes that are either tensors or dict of tensors are initialized to empty dicts
        kwargs = PagedAttentionArgs(
            input_ids=self.input_ids[:, :q_len],
            position_ids=self.position_ids[:, :q_len],
            cu_seq_lens_q=self.cumulative_seqlens_q[: b_size + 1],
            max_seqlen_q=self.max_seqlen_q,
            logits_indices=self.logits_indices[:q_len],
            cu_seq_lens_k={},
            max_seqlen_k={},
            attention_mask={},
            read_index=[],
            write_index=[],
            cache=self.cache,
            use_cache=False,
        )

        # If we use constant-sized slicing, there are some "padding" queries tokens which FA has some issues with. In
        # some models like Qwen3-4B-Instruct-2507, if we don't include these tokens in cumulative_seqlens_q, there are
        # some NaNs in the output logits even for non-padded tokens.
        if use_padding:
            self.max_seqlen_q = max(self.max_seqlen_q, q_len - self.total_seqlen_q)
            kwargs.max_seqlen_q = self.max_seqlen_q  # Update kwargs with the new value
            self.cumulative_seqlens_q[self.actual_batch_size + 1 :] = q_len
            # FIXME: is there another way to avoid this? It has a very slight impact on performance (~5 tok/s)

        # For the attributes that are lists of tensors, we construct list of tensor references
        for i, (read_index_size, write_index_size) in enumerate(self.actual_index_sizes):
            read_index_size = padded_kv_size if use_padding else read_index_size
            write_index_size = padded_q_size if use_padding else write_index_size
            kwargs.read_index.append(self.read_index_storage[i][:read_index_size])
            kwargs.write_index.append(self.write_index_storage[i][:write_index_size])

        # For the attributes that are dict of tensors, we replace the dict with a tensor if there is only one entry
        layer_types = list(self.cumulative_seqlens_k.keys())
        if len(layer_types) > 1:
            kwargs.max_seqlen_k: dict[str, int] = {}
            kwargs.cu_seq_lens_k: dict[str, torch.Tensor] = {}
            kwargs.attention_mask: dict[str, torch.Tensor] = {}
            for layer_type, seqlens_k in self.cumulative_seqlens_k.items():
                kwargs.cu_seq_lens_k[layer_type] = seqlens_k[: b_size + 1]
                kwargs.max_seqlen_k[layer_type] = self.max_seqlen_k[layer_type]
                if self.attention_mask is not None:
                    k_len = padded_kv_size if use_padding else seqlens_k[b_size]
                    kwargs.attention_mask[layer_type] = self.attention_mask[layer_type][..., :q_len, :k_len]
        else:
            layer_type = layer_types[0]
            kwargs.cu_seq_lens_k = self.cumulative_seqlens_k[layer_type][: b_size + 1]