[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

scripts/benchmarks/cb_sync_driver.py

@@ -141,14 +141,17 @@ class SyncCBDriver:
         # CUDA graph replay path.
         if cfg.compile_mode:
             import torch._dynamo
+
             torch._dynamo.config.cache_size_limit = cfg.dynamo_cache_size_limit
             # GRPO's `requires_grad_` issue doesn't apply here (eval mode).
             try:
                 torch._dynamo.config.allow_unspec_int_on_nn_module = True
             except AttributeError:
                 pass
-            print(f"[cb_sync] torch.compile(model, mode='{cfg.compile_mode}', "
-                  f"dynamic=True)")
+            print(
+                f"[cb_sync] torch.compile(model, mode='{cfg.compile_mode}', "
+                f"dynamic=True)"
+            )
             self.model.forward = torch.compile(
                 self.model.forward,
                 mode = cfg.compile_mode,
@@ -169,8 +172,10 @@ class SyncCBDriver:
         """
         results: dict[str, list[int]] = {}
         while True:
-            if (self.manager.input_queue.empty()
-                    and not self.batch_processor.has_pending_requests()):
+            if (
+                self.manager.input_queue.empty()
+                and not self.batch_processor.has_pending_requests()
+            ):
                 break
             if not self.batch_processor.prepare_next_batch():
                 break
@@ -228,9 +233,17 @@ if __name__ == "__main__":
     parser.add_argument("--n_rounds", type = int, default = 2)
     parser.add_argument("--max_new_tokens", type = int, default = 512)
     parser.add_argument("--attn_impl", default = "paged_attention")
-    parser.add_argument("--compile_mode", default = None,
-                        choices = [None, "default", "reduce-overhead",
-                                   "max-autotune", "max-autotune-no-cudagraphs"])
+    parser.add_argument(
+        "--compile_mode",
+        default = None,
+        choices = [
+            None,
+            "default",
+            "reduce-overhead",
+            "max-autotune",
+            "max-autotune-no-cudagraphs",
+        ],
+    )
     parser.add_argument("--max_batch_tokens", type = int, default = 8192)
     parser.add_argument("--num_blocks", type = int, default = 8192)
     parser.add_argument("--lora_adapter", default = None)
@@ -251,6 +264,7 @@ if __name__ == "__main__":
 
     if args.lora_adapter:
         from peft import PeftModel
+
         model = PeftModel.from_pretrained(
             model, str(Path(args.lora_adapter).resolve()), is_trainable = False
         )
@@ -266,12 +280,16 @@ if __name__ == "__main__":
     ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
     ds = ds.shuffle(seed = 3407).select(range(args.n_prompts))
     messages = [
-        [{"role": "system", "content": SYSTEM_PROMPT},
-         {"role": "user", "content": x["prompt"]}]
+        [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": x["prompt"]},
+        ]
         for x in ds
     ]
-    prompt_ids = [tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
-                  for m in messages]
+    prompt_ids = [
+        tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
+        for m in messages
+    ]
 
     gc_cfg = GenerationConfig(
         max_new_tokens = args.max_new_tokens,
@@ -312,8 +330,10 @@ if __name__ == "__main__":
         torch.cuda.synchronize()
         wall_times.append(time.perf_counter() - t0)
         total_decoded = sum(len(v) for v in results.values())
-        print(f"[cb_sync] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, "
-              f"{total_decoded / wall_times[-1]:.1f} tok/s")
+        print(
+            f"[cb_sync] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, "
+            f"{total_decoded / wall_times[-1]:.1f} tok/s"
+        )
 
     med = sorted(wall_times)[len(wall_times) // 2]
     out = {

scripts/benchmarks/flex_paged_attention.py

@@ -28,21 +28,27 @@ class PagedKVCache(torch.nn.Module):
     def __init__(self, page_table, n_heads, head_dim, dtype):
         super().__init__()
         cache_shape = (1, n_heads, page_table.n_pages * page_table.page_size, head_dim)
-        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))
-        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))
+        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype = dtype))
+        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype = dtype))
 
         self.page_table = page_table
 
-    def update(self, input_pos, k_val, v_val, batch_idx=None):
-        assert batch_idx is not None, "batch_idx is required for paged kv cache, are you using non-paged attention?"
+    def update(self, input_pos, k_val, v_val, batch_idx = None):
+        assert (
+            batch_idx is not None
+        ), "batch_idx is required for paged kv cache, are you using non-paged attention?"
 
         if batch_idx.ndim == 1:
             # batch_idx should be [B] (decode)
-            return self.page_table.assign(batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache)
+            return self.page_table.assign(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
         else:
             assert batch_idx.ndim == 2, "batch_idx must be 1D or 2D"
             # batch_idx should be [1, L] (batch prefill)
-            return self.page_table.assign_prefill_no_paging(batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache)
+            return self.page_table.assign_prefill_no_paging(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
 
 
 class PageTable:
@@ -69,25 +75,40 @@ class PageTable:
         self.device = device
 
         # page table: [logical_batch_idx, logical_block_idx] -> physical_page_idx
-        self.page_table = -torch.ones((max_batch_size, self.n_pages), dtype=torch.int64, device=device)
-        self.page_table[0, :] = 0  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        self.page_table = -torch.ones(
+            (max_batch_size, self.n_pages), dtype = torch.int64, device = device
+        )
+        self.page_table[0, :] = (
+            0  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        )
         self.page_table_cpu = [[] for _ in range(max_batch_size)]
 
-        self.capacity = [0 for _ in range(max_batch_size)]  # capacity: batch_idx -> number of pages allocated * page size
-        self.free_pages = list(reversed(range(1, n_pages)))  # page 0 is reserved for simpler code in assign_prefill_no_paging
-        self.free_batch_idx = list(reversed(range(1, max_batch_size)))  # batch_idx 0 is reserved for no-op
+        self.capacity = [
+            0 for _ in range(max_batch_size)
+        ]  # capacity: batch_idx -> number of pages allocated * page size
+        self.free_pages = list(
+            reversed(range(1, n_pages))
+        )  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        self.free_batch_idx = list(
+            reversed(range(1, max_batch_size))
+        )  # batch_idx 0 is reserved for no-op
 
         # [logical_batch_idx, physical_page_idx] -> logical_page_idx
-        self.physical_to_logical = -torch.ones((max_batch_size, n_pages), dtype=torch.int64, device=device)
+        self.physical_to_logical = -torch.ones(
+            (max_batch_size, n_pages), dtype = torch.int64, device = device
+        )
 
     def can_reserve(self, size: int, batch_idx_int: int | None = None) -> bool:
         """check if we can reserve new pages for an existing request or a new request, without gpu operations"""
         if batch_idx_int is None:
             # check if we can schedule a new request
-            return self.pages_available * self.page_size >= size and len(self.free_batch_idx) > 0
+            return (
+                self.pages_available * self.page_size >= size
+                and len(self.free_batch_idx) > 0
+            )
         else:
             # check if we can reserve new pages for an existing request
-            return self.reserve(batch_idx_int, None, size, dry_run=True)
+            return self.reserve(batch_idx_int, None, size, dry_run = True)
 
     def allocate(self) -> int:
         """allocate a new batch"""
@@ -102,7 +123,13 @@ class PageTable:
     def pages_available(self) -> int:
         return len(self.free_pages)
 
-    def reserve(self, batch_idx_int: int, batch_idx: torch.Tensor, seq_len: int, dry_run: bool = False) -> bool:
+    def reserve(
+        self,
+        batch_idx_int: int,
+        batch_idx: torch.Tensor,
+        seq_len: int,
+        dry_run: bool = False,
+    ) -> bool:
         """
         Requests the capacity of a given batch to be at least enough to
         hold `seq_len` elements.
@@ -119,7 +146,9 @@ class PageTable:
         if seq_len <= self.capacity[batch_idx_int]:
             return True
 
-        num_pages_to_allocate = _cdiv(seq_len - self.capacity[batch_idx_int], self.page_size)
+        num_pages_to_allocate = _cdiv(
+            seq_len - self.capacity[batch_idx_int], self.page_size
+        )
 
         can_allocate = num_pages_to_allocate <= self.pages_available
         if dry_run:
@@ -137,7 +166,7 @@ class PageTable:
 
         # find empty physical pages
         allocated_pages_list = self.free_pages[-num_pages_to_allocate:]
-        allocated_pages = torch.tensor(allocated_pages_list, device=self.device)
+        allocated_pages = torch.tensor(allocated_pages_list, device = self.device)
         # update page table
         self.page_table[batch_idx, start_page_idx:end_page_idx] = allocated_pages
 
@@ -145,7 +174,7 @@ class PageTable:
         self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
             start_page_idx,
             end_page_idx,
-            device=self.device,
+            device = self.device,
         )
         # update cpu side metadata
         self.page_table_cpu[batch_idx_int] += allocated_pages_list
@@ -198,24 +227,38 @@ class PageTable:
 
         V_D = v_val.shape[3]
         if B != batch_idx.shape[0]:
-            raise RuntimeError(f"Expect val and batch_idx have the same batch size but got B={B} and B={batch_idx.shape[0]}.")
+            raise RuntimeError(
+                f"Expect val and batch_idx have the same batch size but got B={B} and B={batch_idx.shape[0]}."
+            )
         if H != k_cache.shape[1]:
-            raise RuntimeError(f"Expect val and cache has the same number of heads but got H={H} and H={k_cache.shape[1]}.")
+            raise RuntimeError(
+                f"Expect val and cache has the same number of heads but got H={H} and H={k_cache.shape[1]}."
+            )
         if S != input_pos.shape[1]:
-            raise RuntimeError(f"Expect val and input_pos has the same length but got S={S} and S={input_pos.shape[0]}.")
+            raise RuntimeError(
+                f"Expect val and input_pos has the same length but got S={S} and S={input_pos.shape[0]}."
+            )
         if K_D != k_cache.shape[3]:
-            raise RuntimeError(f"Expect k_val and k_cache has the same hidden dim but got D={K_D} and D={k_cache.shape[3]}.")
+            raise RuntimeError(
+                f"Expect k_val and k_cache has the same hidden dim but got D={K_D} and D={k_cache.shape[3]}."
+            )
         if V_D != v_cache.shape[3]:
-            raise RuntimeError(f"Expect v_val and v_cache has the same hidden dim but got D={V_D} and D={v_cache.shape[3]}.")
+            raise RuntimeError(
+                f"Expect v_val and v_cache has the same hidden dim but got D={V_D} and D={v_cache.shape[3]}."
+            )
 
         # find address
         logical_block_idx = input_pos // self.page_size  # [B, S]
         logical_block_offset = input_pos % self.page_size  # [B, S]
 
         # NOTE: this code path is only used for decoding. For batch prefill, use assign_prefill_no_paging() instead
-        physical_block_idx = torch.gather(self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)).to(torch.int32)  # [B, S]
+        physical_block_idx = torch.gather(
+            self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
+        ).to(torch.int32)  # [B, S]
 
-        addr = (physical_block_idx * self.page_size + logical_block_offset).view(-1)  # [B*S]
+        addr = (physical_block_idx * self.page_size + logical_block_offset).view(
+            -1
+        )  # [B*S]
 
         k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
         v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
@@ -252,11 +295,15 @@ class PageTable:
         device = block_mask.kv_num_blocks.device
 
         if batch_idx is None:
-            batch_idx = torch.arange(B, device=device)
+            batch_idx = torch.arange(B, device = device)
 
         assert batch_idx.ndim == 1, "batch_idx must be a 1D tensor"
-        assert batch_idx.shape[0] == B, "batch_idx must have the same shape as block_mask"
-        assert B <= self.max_batch_size, "batch_idx must be less than or equal to max_batch_size"
+        assert (
+            batch_idx.shape[0] == B
+        ), "batch_idx must have the same shape as block_mask"
+        assert (
+            B <= self.max_batch_size
+        ), "batch_idx must be less than or equal to max_batch_size"
 
         page_table = self.page_table[batch_idx]
 
@@ -271,14 +318,22 @@ class PageTable:
             if num_blocks is None:
                 return None, None
             new_kv_num_blocks = num_blocks.clone()
-            new_kv_indices = torch.zeros((B, H, ROWS, self.n_pages), dtype=torch.int32, device=device)
+            new_kv_indices = torch.zeros(
+                (B, H, ROWS, self.n_pages), dtype = torch.int32, device = device
+            )
             new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
-                torch.gather(page_table, 1, indices.view(B, -1).to(torch.int64)).view(block_mask.kv_indices.shape).to(torch.int32)
+                torch.gather(page_table, 1, indices.view(B, -1).to(torch.int64))
+                .view(block_mask.kv_indices.shape)
+                .to(torch.int32)
             )
             return new_kv_num_blocks, new_kv_indices
 
-        new_kv_num_blocks, new_kv_indices = transform(block_mask.kv_num_blocks, block_mask.kv_indices)
-        new_full_kv_num_blocks, new_full_kv_indices = transform(block_mask.full_kv_num_blocks, block_mask.full_kv_indices)
+        new_kv_num_blocks, new_kv_indices = transform(
+            block_mask.kv_num_blocks, block_mask.kv_indices
+        )
+        new_full_kv_num_blocks, new_full_kv_indices = transform(
+            block_mask.full_kv_num_blocks, block_mask.full_kv_indices
+        )
 
         new_mask_mod = self.get_mask_mod(block_mask.mask_mod, batch_idx)
 
@@ -290,20 +345,29 @@ class PageTable:
             new_full_kv_indices,
             block_mask.BLOCK_SIZE,
             new_mask_mod,
-            seq_lengths=seq_lengths,
+            seq_lengths = seq_lengths,
         )
 
-    def get_logical_kv_idx(self, physical_batch_idx: torch.Tensor, physical_kv_idx: torch.Tensor, batch_idx: torch.Tensor):
+    def get_logical_kv_idx(
+        self,
+        physical_batch_idx: torch.Tensor,
+        physical_kv_idx: torch.Tensor,
+        batch_idx: torch.Tensor,
+    ):
         logical_batch_idx = batch_idx[physical_batch_idx]
         physical_kv_block = physical_kv_idx // self.page_size
         physical_kv_offset = physical_kv_idx % self.page_size
-        logical_block_idx = self.physical_to_logical[logical_batch_idx, physical_kv_block]
+        logical_block_idx = self.physical_to_logical[
+            logical_batch_idx, physical_kv_block
+        ]
         logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
         is_valid = logical_block_idx >= 0
-        safe_logical_kv_idx = logical_kv_idx.clamp(min=0)
+        safe_logical_kv_idx = logical_kv_idx.clamp(min = 0)
         return is_valid, safe_logical_kv_idx
 
-    def get_mask_mod(self, mask_mod: Optional[_mask_mod_signature], batch_idx: torch.Tensor) -> _mask_mod_signature:
+    def get_mask_mod(
+        self, mask_mod: Optional[_mask_mod_signature], batch_idx: torch.Tensor
+    ) -> _mask_mod_signature:
         """
         Converts a mask_mod based on mapping from the physical block index to the logical
         block index.
@@ -320,13 +384,19 @@ class PageTable:
             q_idx: torch.Tensor,
             physical_kv_idx: torch.Tensor,
         ):
-            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(b, physical_kv_idx, batch_idx)
-            return torch.where(is_valid, mask_mod(b, h, q_idx, safe_logical_kv_idx), False)
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
+            return torch.where(
+                is_valid, mask_mod(b, h, q_idx, safe_logical_kv_idx), False
+            )
 
         return new_mask_mod
 
     # NOTE: not used in the current codebase
-    def get_score_mod(self, score_mod: Optional[_score_mod_signature], batch_idx: torch.Tensor) -> _score_mod_signature:
+    def get_score_mod(
+        self, score_mod: Optional[_score_mod_signature], batch_idx: torch.Tensor
+    ) -> _score_mod_signature:
         """
         Converts a score_mod based on mapping from the physical block index to the logical
         block index.
@@ -344,7 +414,9 @@ class PageTable:
             q_idx: torch.Tensor,
             physical_kv_idx: torch.Tensor,
         ):
-            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(b, physical_kv_idx, batch_idx)
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
             return torch.where(
                 is_valid,
                 score_mod(score, b, h, q_idx, safe_logical_kv_idx),
@@ -359,9 +431,19 @@ class PageTable:
         def causal(b, h, q_idx, kv_idx):
             return q_idx >= kv_idx
 
-        return create_block_mask(causal, B=B, H=None, Q_LEN=L, KV_LEN=L, BLOCK_SIZE=self.page_size, device=self.device)
+        return create_block_mask(
+            causal,
+            B = B,
+            H = None,
+            Q_LEN = L,
+            KV_LEN = L,
+            BLOCK_SIZE = self.page_size,
+            device = self.device,
+        )
 
-    def create_prefill_blockmask_no_paging(self, batch_idx: Tensor, BLOCK_SIZE: int = 128):
+    def create_prefill_blockmask_no_paging(
+        self, batch_idx: Tensor, BLOCK_SIZE: int = 128
+    ):
         """
         there's no prefix sharing implemented, batch_idx is the document id, batch_idx is not guaranteed to be sorted
         """
@@ -375,7 +457,9 @@ class PageTable:
             document_mask = docs[q_idx] == docs[kv_idx]
             return causal_mask & document_mask
 
-        return create_block_mask(document_causal, B=1, H=None, Q_LEN=L, KV_LEN=L, BLOCK_SIZE=BLOCK_SIZE)
+        return create_block_mask(
+            document_causal, B = 1, H = None, Q_LEN = L, KV_LEN = L, BLOCK_SIZE = BLOCK_SIZE
+        )
 
     # we assign prefill to the cache, similar to assign(), except we don't return the k_cache, v_cache, we only return the k_val, v_val
     def assign_prefill_no_paging(
@@ -408,7 +492,10 @@ class PageTable:
 
         input_pos_block_idx = input_pos // self.page_size
         input_pos_offset_in_block = input_pos % self.page_size
-        physical_kv_idx = self.page_table[batch_idx, input_pos_block_idx] * self.page_size + input_pos_offset_in_block
+        physical_kv_idx = (
+            self.page_table[batch_idx, input_pos_block_idx] * self.page_size
+            + input_pos_offset_in_block
+        )
         k_cache[:, :, physical_kv_idx.view(-1), :] = k_val
         v_cache[:, :, physical_kv_idx.view(-1), :] = v_val
 

scripts/benchmarks/qwen3_flex_inference.py

@@ -58,10 +58,12 @@ flex_attention_compiled = torch.compile(flex_attention, fullgraph = True)
 def _apply_rotary(q, k, cos, sin):
     cos = cos.unsqueeze(1)
     sin = sin.unsqueeze(1)
+
     def rotate_half(x):
         x1 = x[..., : x.shape[-1] // 2]
-        x2 = x[..., x.shape[-1] // 2:]
+        x2 = x[..., x.shape[-1] // 2 :]
         return torch.cat((-x2, x1), dim = -1)
+
     q = (q * cos) + (rotate_half(q) * sin)
     k = (k * cos) + (rotate_half(k) * sin)
     return q, k
@@ -143,11 +145,13 @@ def patch_qwen3_model(model: torch.nn.Module, page_table: PageTable):
         ).to(model.device)
         # Bind as method.
         import types
+
         attn.forward = types.MethodType(fwd, attn)
 
 
 # --- model forward helper that passes flex kwargs through ------------------
 
+
 def call_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs):
     """`model(**inputs, **flex_kwargs)` would error because Qwen3ForCausalLM
     doesn't declare the flex_* kwargs. We walk through the model manually
@@ -175,6 +179,7 @@ def call_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs):
 
 # --- inference engine ------------------------------------------------------
 
+
 @dataclass
 class Sequence:
     text: str = ""
@@ -196,8 +201,16 @@ class Sequence:
 
 
 class FlexInference:
-    def __init__(self, model, tokenizer, max_batch_size = 32, max_seq_length = 2048,
-                 n_pages = 2048, page_size = 128, max_new_tokens = 512):
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        max_batch_size = 32,
+        max_seq_length = 2048,
+        n_pages = 2048,
+        page_size = 128,
+        max_new_tokens = 512,
+    ):
         assert max_seq_length % page_size == 0
         self.model = model
         self.tokenizer = tokenizer
@@ -209,17 +222,21 @@ class FlexInference:
         self.max_new_tokens = max_new_tokens
 
         self.page_table = PageTable(
-            n_pages = n_pages, page_size = page_size,
-            max_batch_size = max_batch_size, device = self.device.type,
+            n_pages = n_pages,
+            page_size = page_size,
+            max_batch_size = max_batch_size,
+            device = self.device.type,
         )
         patch_qwen3_model(model, self.page_table)
 
         # Pre-allocated decode state.
-        self.input_pos_buffer = torch.zeros(max_batch_size, dtype = torch.int32,
-                                            device = self.device)
+        self.input_pos_buffer = torch.zeros(
+            max_batch_size, dtype = torch.int32, device = self.device
+        )
         # Full-length logical causal mask (shared across decode batch).
         self.block_mask_logical = self.page_table.create_causal_blockmask(
-            B = max_batch_size, L = max_seq_length,
+            B = max_batch_size,
+            L = max_seq_length,
         )
 
         self.cudagraph_captured = False
@@ -238,11 +255,16 @@ class FlexInference:
         sequence as [num_seqs, V].
         """
         input_ids_list = [seq.input_ids.to(self.device) for seq in batch]
-        input_pos_list = [torch.arange(seq.input_length, dtype = torch.long,
-                                       device = self.device) for seq in batch]
-        batch_idx_list = [torch.full((seq.input_length,), seq.batch_idx,
-                                     dtype = torch.long, device = self.device)
-                          for seq in batch]
+        input_pos_list = [
+            torch.arange(seq.input_length, dtype = torch.long, device = self.device)
+            for seq in batch
+        ]
+        batch_idx_list = [
+            torch.full(
+                (seq.input_length,), seq.batch_idx, dtype = torch.long, device = self.device
+            )
+            for seq in batch
+        ]
         input_ids = torch.cat(input_ids_list).view(1, -1)
         input_pos = torch.cat(input_pos_list).view(1, -1)
         batch_idx = torch.cat(batch_idx_list).view(1, -1)
@@ -255,8 +277,9 @@ class FlexInference:
             input_pos = F.pad(input_pos, (0, pad), value = 0)
             batch_idx = F.pad(batch_idx, (0, pad), value = 0)
 
-        input_lengths = torch.tensor([s.input_length for s in batch],
-                                     dtype = torch.long, device = self.device)
+        input_lengths = torch.tensor(
+            [s.input_length for s in batch], dtype = torch.long, device = self.device
+        )
         logits_positions = input_lengths.cumsum(dim = 0) - 1  # [num_seqs]
 
         mask = self.page_table.create_prefill_blockmask_no_paging(batch_idx)
@@ -268,8 +291,9 @@ class FlexInference:
             flex_kernel_options = {"FORCE_USE_FLEX_ATTENTION": True},
         )
         position_ids = input_pos  # Qwen3 uses 0-based; unlike Gemma2
-        hidden = call_model_with_flex_kwargs(self.model, input_ids, position_ids,
-                                             flex_kwargs)
+        hidden = call_model_with_flex_kwargs(
+            self.model, input_ids, position_ids, flex_kwargs
+        )
         return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0)
 
     def _decode_block_mask(self, batch_idx: torch.Tensor):
@@ -280,21 +304,33 @@ class FlexInference:
         assert batch_idx.ndim == 1 and input_pos.ndim == 1
         B = batch_idx.shape[0]
         input_block_idx = input_pos // block_mask.BLOCK_SIZE[0]
-        kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view(B, 1, 1)
-        kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view(B, 1, 1, -1)
+        kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view(
+            B, 1, 1
+        )
+        kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view(
+            B, 1, 1, -1
+        )
         full_num = full_idx = None
         if block_mask.full_kv_num_blocks is not None:
-            full_num = block_mask.full_kv_num_blocks[batch_idx, :, input_block_idx].view(B, 1, 1)
-            full_idx = block_mask.full_kv_indices[batch_idx, :, input_block_idx].view(B, 1, 1, -1)
+            full_num = block_mask.full_kv_num_blocks[
+                batch_idx, :, input_block_idx
+            ].view(B, 1, 1)
+            full_idx = block_mask.full_kv_indices[batch_idx, :, input_block_idx].view(
+                B, 1, 1, -1
+            )
 
         def causal_offset(off):
             def offset(b, h, q_idx, kv_idx):
                 return q_idx + off[b] >= kv_idx
+
             return offset
 
         seq_length = (1, block_mask.seq_lengths[1])
         mask = BlockMask.from_kv_blocks(
-            kv_num_blocks, kv_indices, full_num, full_idx,
+            kv_num_blocks,
+            kv_indices,
+            full_num,
+            full_idx,
             BLOCK_SIZE = block_mask.BLOCK_SIZE,
             mask_mod = causal_offset(input_pos),
             seq_lengths = seq_length,
@@ -312,12 +348,14 @@ class FlexInference:
             flex_batch_idx = batch_idx,
             flex_kernel_options = None,
         )
-        hidden = call_model_with_flex_kwargs(self.model, input_ids.view(B, 1),
-                                             position_ids, flex_kwargs)
+        hidden = call_model_with_flex_kwargs(
+            self.model, input_ids.view(B, 1), position_ids, flex_kwargs
+        )
         return self.model.lm_head(hidden[:, -1, :])  # [B, V]
 
-    def _decode_step(self, batch_idx: torch.Tensor, input_ids: torch.Tensor,
-                     input_pos: torch.Tensor):
+    def _decode_step(
+        self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor
+    ):
         self.input_pos_buffer.zero_()
         self.input_pos_buffer[batch_idx] = input_pos
         if not self.cudagraph_captured:
@@ -365,8 +403,11 @@ class FlexInference:
 
         input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device)
         batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device)
-        outputs = torch.zeros((max_bs, self.model.config.vocab_size),
-                              dtype = self.model.dtype, device = self.device)
+        outputs = torch.zeros(
+            (max_bs, self.model.config.vocab_size),
+            dtype = self.model.dtype,
+            device = self.device,
+        )
         self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
         pool = None
         for bs in reversed(self.graph_bs):
@@ -386,7 +427,9 @@ class FlexInference:
         # Release the scratch batches; real requests will re-allocate them.
         for bi in reserved_batches:
             self.page_table.erase(bi)
-        self.graph_vars = dict(input_ids = input_ids, batch_idx = batch_idx, outputs = outputs)
+        self.graph_vars = dict(
+            input_ids = input_ids, batch_idx = batch_idx, outputs = outputs
+        )
 
     @torch.inference_mode()
     def generate(self, sequences: list[Sequence], capture_cudagraph = False):
@@ -406,7 +449,8 @@ class FlexInference:
                 seq = waiting.popleft()
                 bi = self.page_table.allocate()
                 self.page_table.reserve(
-                    bi, torch.tensor([bi], device = self.device, dtype = torch.long),
+                    bi,
+                    torch.tensor([bi], device = self.device, dtype = torch.long),
                     seq.total_length,
                 )
                 seq.batch_idx = bi
@@ -417,8 +461,10 @@ class FlexInference:
                 for i, seq in enumerate(batch):
                     seq.last_token_id = next_ids[i]
                     seq.output_ids.append(next_ids[i])
-                    if (seq.last_token_id == self.eos_token_id
-                        or len(seq.output_ids) >= seq.max_new_tokens):
+                    if (
+                        seq.last_token_id == self.eos_token_id
+                        or len(seq.output_ids) >= seq.max_new_tokens
+                    ):
                         seq.finished = True
                         done.append(seq)
                         self.page_table.erase(seq.batch_idx)
@@ -432,12 +478,14 @@ class FlexInference:
                 seq = running.popleft()
                 if self.page_table.capacity[seq.batch_idx] >= seq.total_length:
                     decode_batch.append(seq)
-                elif self.page_table.can_reserve(seq.total_length,
-                                                  batch_idx_int = seq.batch_idx):
+                elif self.page_table.can_reserve(
+                    seq.total_length, batch_idx_int = seq.batch_idx
+                ):
                     self.page_table.reserve(
                         seq.batch_idx,
-                        torch.tensor([seq.batch_idx], device = self.device,
-                                     dtype = torch.long),
+                        torch.tensor(
+                            [seq.batch_idx], device = self.device, dtype = torch.long
+                        ),
                         seq.total_length,
                     )
                     decode_batch.append(seq)
@@ -450,19 +498,30 @@ class FlexInference:
                 continue
 
             B = len(decode_batch)
-            bi_tensor = torch.tensor([s.batch_idx for s in decode_batch],
-                                     dtype = torch.long, device = self.device)
-            last_ids = torch.tensor([s.last_token_id for s in decode_batch],
-                                    dtype = torch.long, device = self.device)
-            cur_pos = torch.tensor([s.total_length - 1 for s in decode_batch],
-                                   dtype = torch.int32, device = self.device)
+            bi_tensor = torch.tensor(
+                [s.batch_idx for s in decode_batch],
+                dtype = torch.long,
+                device = self.device,
+            )
+            last_ids = torch.tensor(
+                [s.last_token_id for s in decode_batch],
+                dtype = torch.long,
+                device = self.device,
+            )
+            cur_pos = torch.tensor(
+                [s.total_length - 1 for s in decode_batch],
+                dtype = torch.int32,
+                device = self.device,
+            )
             logits = self._decode_step(bi_tensor, last_ids, cur_pos)
             next_ids = torch.argmax(logits, dim = -1).tolist()
             for i, seq in enumerate(decode_batch):
                 seq.last_token_id = next_ids[i]
                 seq.output_ids.append(next_ids[i])
-                if (seq.last_token_id == self.eos_token_id
-                    or len(seq.output_ids) >= seq.max_new_tokens):
+                if (
+                    seq.last_token_id == self.eos_token_id
+                    or len(seq.output_ids) >= seq.max_new_tokens
+                ):
                     seq.finished = True
                     done.append(seq)
                     self.page_table.erase(seq.batch_idx)
@@ -488,19 +547,25 @@ def main():
     args = p.parse_args()
 
     from transformers import AutoModelForCausalLM, AutoTokenizer
+
     tok = AutoTokenizer.from_pretrained(args.model_name)
     if tok.pad_token is None:
         tok.pad_token = tok.eos_token
     # Load eager; we swap attention forward below.
     model = AutoModelForCausalLM.from_pretrained(
-        args.model_name, dtype = torch.bfloat16, attn_implementation = "eager",
+        args.model_name,
+        dtype = torch.bfloat16,
+        attn_implementation = "eager",
     ).to("cuda")
     model.eval()
 
     if args.lora_adapter:
         from peft import PeftModel
+
         model = PeftModel.from_pretrained(
-            model, str(Path(args.lora_adapter).resolve()), is_trainable = False,
+            model,
+            str(Path(args.lora_adapter).resolve()),
+            is_trainable = False,
         )
         # Merge so attention forward below sees merged weights without the
         # PEFT wrapper mangling `self.q_proj` etc.
@@ -508,24 +573,35 @@ def main():
         model.eval()
 
     from unsloth_grpo_common import (
-        SYSTEM_PROMPT, apply_chat_template_to_tokenizer,
+        SYSTEM_PROMPT,
+        apply_chat_template_to_tokenizer,
     )
     from datasets import load_dataset
+
     apply_chat_template_to_tokenizer(tok)
     ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
     ds = ds.shuffle(seed = 3407).select(range(args.n_prompts))
-    messages = [[{"role": "system", "content": SYSTEM_PROMPT},
-                 {"role": "user", "content": x["prompt"]}] for x in ds]
-    texts = [tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False)
-             for m in messages]
+    messages = [
+        [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": x["prompt"]},
+        ]
+        for x in ds
+    ]
+    texts = [
+        tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False)
+        for m in messages
+    ]
 
     # Make sure the base HF model that Qwen3Attention belongs to isn't wrapped
     # by PeftModel anymore (we merged); `.model` should be Qwen3ForCausalLM.
     inference = FlexInference(
-        model, tok,
+        model,
+        tok,
         max_batch_size = args.max_batch_size,
         max_seq_length = args.max_seq_length,
-        n_pages = args.n_pages, page_size = args.page_size,
+        n_pages = args.n_pages,
+        page_size = args.page_size,
         max_new_tokens = args.max_new_tokens,
     )
 
@@ -547,8 +623,10 @@ def main():
         torch.cuda.synchronize()
         wall_times.append(time.perf_counter() - t0)
         total_decoded = sum(len(s.output_ids) for s in out)
-        print(f"[flex] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, "
-              f"{total_decoded / wall_times[-1]:.1f} tok/s")
+        print(
+            f"[flex] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, "
+            f"{total_decoded / wall_times[-1]:.1f} tok/s"
+        )
 
     med = sorted(wall_times)[len(wall_times) // 2]
     peak = torch.cuda.max_memory_allocated() / 1024**3