Repository: huggingface/transformers
File: tests/generation/test_paged_attention.py (Python)

Code before:
    def setUpClass(cls):
        cls.model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-3b-Instruct", dtype="bfloat16", device_map="auto"
        ).eval()

        cls.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3b-Instruct", padding_side="left")

        if cls.tokenizer.pad_token is None:
            cls.tokenizer.pad_token = cls.tokenizer.eos_token
            cls.model.config.pad_token_id = cls.model.config.eos_token_id

        cls.model.use_cache = False

    @parameterized.expand(
        [
            ("paged|eager", 64, 128, 64),
            ("paged|sdpa", 32, 256, 128),
            ("paged|flash_attention_2", 16, 512, 256),
            ("paged|flex_attention", 64, 128, 64),
        ]
    )
    def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max_batch_tokens):
        if (
            attn_impl
            in [
                "paged|flex_attention",
            ]
            and torch_device == "cpu"
        ):
            self.skipTest(
                f"CPU only support flash/sdpa/eager paged attention for now, but found {attn_impl}. Skipping test."
            )

        self.model.config.attn_implementation = attn_impl

        generation_config = GenerationConfig(
            max_new_tokens=50,
            top_k=0,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.tokenizer.pad_token_id,
            use_cache=False,
            num_blocks=num_blocks,
            block_size=block_size,
            max_batch_tokens=max_batch_tokens,
        )

        tokenized = self.tokenizer(_TEST_PROMPTS, truncation=True, max_length=512)
        batch_inputs = list(tokenized["input_ids"])

        start = time.time()
        batch_outputs = self.model.generate_batch(

Review comment (suggestion):
You can remove the `paged|flex_attention` case and the associated checks for it 

Code after:
class TestBatchGeneration(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-3b-Instruct", dtype="bfloat16", device_map="auto"
        ).eval()

        cls.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3b-Instruct", padding_side="left")

        if cls.tokenizer.pad_token is None:
            cls.tokenizer.pad_token = cls.tokenizer.eos_token
            cls.model.config.pad_token_id = cls.model.config.eos_token_id

        cls.model.use_cache = False

    @parameterized.expand(
        [
            ("paged|eager", 64, 128, 64),
            ("paged|sdpa", 32, 256, 128),
            ("paged|flash_attention_2", 16, 512, 256),
            ("paged|flex_attention", 64, 128, 64),
        ]
    )
    def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max_batch_tokens):
        self.model.config.attn_implementation = attn_impl

        generation_config = GenerationConfig(
            max_new_tokens=50,
            top_k=0,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.tokenizer.pad_token_id,
            use_cache=False,
            num_blocks=num_blocks,
            block_size=block_size,
            max_batch_tokens=max_batch_tokens,
        )

        tokenized = self.tokenizer(_TEST_PROMPTS, truncation=True, max_length=512)
        batch_inputs = list(tokenized["input_ids"])

        start = time.time()
        batch_outputs = self.model.generate_batch(
            inputs=batch_inputs,
            generation_config=generation_config,
        )
        end = time.time()
        print(
            f"\n[{attn_impl}] Batch took {end - start:.2f}s with config: blocks={num_blocks}, block_size={block_size}, max_batch_tokens={max_batch_tokens}"
        )

        expected_outputs = _EXPECTED_OUTPUTS.get_expectation()