Apply lintrunner auto-fixes more

Author: cthotti

examples/models/gemma4_31b/CMakeLists.txt

@@ -74,7 +74,7 @@ target_include_directories(
 )
 target_link_libraries(gemma4_31b_worker PUBLIC ${link_libraries})
 
-# Standalone C++ DFlash benchmark runner using the Module API. 
+# Standalone C++ DFlash benchmark runner using the Module API.
 add_executable(dflash_cpp_driver dflash_cpp_driver.cpp)
 target_include_directories(
   dflash_cpp_driver PUBLIC ${_common_include_directories} ${_json_include}

examples/models/gemma4_31b/dflash_cpp_driver.cpp

@@ -1,4 +1,5 @@
-// C++ DFlash runner for benchmarking the exported Gemma4-31B target and draft models. 
+// C++ DFlash runner for benchmarking the exported Gemma4-31B target and draft
+// models.
 
 #include <gflags/gflags.h>
 
@@ -10,25 +11,34 @@
 #include <algorithm>
 #include <chrono>
 #include <cstdint>
-#include <cstring>
 #include <cstdio>
+#include <cstring>
 #include <set>
 #include <string>
 #include <vector>
 
-using executorch::extension::Module;
 using executorch::extension::make_tensor_ptr;
+using executorch::extension::Module;
 using executorch::extension::TensorPtr;
-using executorch::runtime::EValue;
 using executorch::runtime::Error;
+using executorch::runtime::EValue;
 
-DEFINE_string(target_pte, "gemma4_31b_dflash_exports_mlx/model.pte", "Target .pte path.");
+DEFINE_string(
+    target_pte,
+    "gemma4_31b_dflash_exports_mlx/model.pte",
+    "Target .pte path.");
 DEFINE_string(draft_pte, "gemma4_31b_dflash_draft.pte", "Draft .pte path.");
-DEFINE_string(tokenizer_path, "gemma-4-31B-it-HQQ-INT4/tokenizer.json", "Tokenizer path.");
+DEFINE_string(
+    tokenizer_path,
+    "gemma-4-31B-it-HQQ-INT4/tokenizer.json",
+    "Tokenizer path.");
 DEFINE_string(prompt, "The capital of France is", "Prompt text.");
 DEFINE_int32(max_new_tokens, 64, "Maximum tokens to generate.");
 DEFINE_int32(block_size, 16, "DFlash draft block size.");
-DEFINE_int32(mask_id, 4, "DFlash mask token id (from draft checkpoint config).");
+DEFINE_int32(
+    mask_id,
+    4,
+    "DFlash mask token id (from draft checkpoint config).");
 DEFINE_bool(raw_prompt, false, "Skip chat-template wrapping.");
 DEFINE_bool(verbose, false, "Print per-round timing/acceptance debug output.");
 
@@ -72,7 +82,8 @@ int64_t first_mismatch(
   return static_cast<int64_t>(draft_ids.size());
 }
 
-// Convert BF16 hidden states returned by the target model to FP32 for the draft model. 
+// Convert BF16 hidden states returned by the target model to FP32 for the draft
+// model.
 std::vector<float> bf16_tensor_to_fp32(const executorch::aten::Tensor& t) {
   int64_t numel = t.numel();
   const uint16_t* src = reinterpret_cast<const uint16_t*>(t.const_data_ptr());
@@ -84,7 +95,7 @@ std::vector<float> bf16_tensor_to_fp32(const executorch::aten::Tensor& t) {
   return out;
 }
 
-}
+} // namespace
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
@@ -122,15 +133,20 @@ int main(int argc, char** argv) {
   printf("Prompt tokens: %lld\n", (long long)prompt_len);
 
   std::vector<int64_t> input_pos_vec(prompt_len);
-  for (int64_t i = 0; i < prompt_len; ++i) input_pos_vec[i] = i;
+  for (int64_t i = 0; i < prompt_len; ++i)
+    input_pos_vec[i] = i;
 
-  auto prompt_ids_tensor =
-      make_tensor_ptr({1, (int)prompt_len}, prompt_ids.data(), executorch::aten::ScalarType::Long);
+  auto prompt_ids_tensor = make_tensor_ptr(
+      {1, (int)prompt_len},
+      prompt_ids.data(),
+      executorch::aten::ScalarType::Long);
   auto input_pos_tensor = make_tensor_ptr(
-      {(int)prompt_len}, input_pos_vec.data(), executorch::aten::ScalarType::Long);
+      {(int)prompt_len},
+      input_pos_vec.data(),
+      executorch::aten::ScalarType::Long);
 
-  auto prefill_result =
-      target_module.forward({EValue(prompt_ids_tensor), EValue(input_pos_tensor)});
+  auto prefill_result = target_module.forward(
+      {EValue(prompt_ids_tensor), EValue(input_pos_tensor)});
   if (!prefill_result.ok()) {
     ET_LOG(Error, "Prefill forward failed");
     return 1;
@@ -148,31 +164,36 @@ int main(int argc, char** argv) {
   int64_t accepted_total = 0;
   int64_t emitted_total = 0;
 
-
-
   int64_t hidden_concat_dim = hidden.sizes()[2];
   std::vector<float> hidden_history = bf16_tensor_to_fp32(hidden);
   int64_t hidden_len = prompt_len;
 
-  // Warm up the draft model to avoid including one-time MLX JIT compilation in benchmark timings. 
+  // Warm up the draft model to avoid including one-time MLX JIT compilation in
+  // benchmark timings.
   {
     std::vector<int64_t> warm_input_vec(FLAGS_block_size);
     warm_input_vec[0] = last_token;
-    for (int64_t i = 1; i < FLAGS_block_size; ++i) warm_input_vec[i] = FLAGS_mask_id;
+    for (int64_t i = 1; i < FLAGS_block_size; ++i)
+      warm_input_vec[i] = FLAGS_mask_id;
     auto warm_input_tensor = make_tensor_ptr(
-        {1, (int)FLAGS_block_size}, warm_input_vec.data(), executorch::aten::ScalarType::Long);
+        {1, (int)FLAGS_block_size},
+        warm_input_vec.data(),
+        executorch::aten::ScalarType::Long);
     auto warm_hidden_tensor = make_tensor_ptr(
         {1, (int)hidden_len, (int)hidden_concat_dim},
         hidden_history.data(),
         executorch::aten::ScalarType::Float);
     std::vector<int64_t> warm_pos_vec(hidden_len + FLAGS_block_size);
-    for (int64_t i = 0; i < hidden_len + FLAGS_block_size; ++i) warm_pos_vec[i] = i;
+    for (int64_t i = 0; i < hidden_len + FLAGS_block_size; ++i)
+      warm_pos_vec[i] = i;
     auto warm_pos_tensor = make_tensor_ptr(
         {1, (int)(hidden_len + FLAGS_block_size)},
         warm_pos_vec.data(),
         executorch::aten::ScalarType::Long);
     auto warm_result = draft_module.forward(
-        {EValue(warm_input_tensor), EValue(warm_hidden_tensor), EValue(warm_pos_tensor)});
+        {EValue(warm_input_tensor),
+         EValue(warm_hidden_tensor),
+         EValue(warm_pos_tensor)});
     if (!warm_result.ok()) {
       ET_LOG(Error, "Draft warm-up forward failed");
       return 1;
@@ -188,23 +209,31 @@ int main(int argc, char** argv) {
 
     std::vector<int64_t> draft_input_vec(bs);
     draft_input_vec[0] = last_token;
-    for (int64_t i = 1; i < bs; ++i) draft_input_vec[i] = FLAGS_mask_id;
-    auto draft_input_tensor =
-        make_tensor_ptr({1, (int)bs}, draft_input_vec.data(), executorch::aten::ScalarType::Long);
+    for (int64_t i = 1; i < bs; ++i)
+      draft_input_vec[i] = FLAGS_mask_id;
+    auto draft_input_tensor = make_tensor_ptr(
+        {1, (int)bs},
+        draft_input_vec.data(),
+        executorch::aten::ScalarType::Long);
 
     auto hidden_tensor = make_tensor_ptr(
         {1, (int)hidden_len, (int)hidden_concat_dim},
         hidden_history.data(),
         executorch::aten::ScalarType::Float);
 
     std::vector<int64_t> draft_pos_vec(hidden_len + bs);
-    for (int64_t i = 0; i < hidden_len + bs; ++i) draft_pos_vec[i] = i;
+    for (int64_t i = 0; i < hidden_len + bs; ++i)
+      draft_pos_vec[i] = i;
     auto draft_pos_tensor = make_tensor_ptr(
-        {1, (int)(hidden_len + bs)}, draft_pos_vec.data(), executorch::aten::ScalarType::Long);
+        {1, (int)(hidden_len + bs)},
+        draft_pos_vec.data(),
+        executorch::aten::ScalarType::Long);
 
     double dt0 = now_ms();
     auto draft_result = draft_module.forward(
-        {EValue(draft_input_tensor), EValue(hidden_tensor), EValue(draft_pos_tensor)});
+        {EValue(draft_input_tensor),
+         EValue(hidden_tensor),
+         EValue(draft_pos_tensor)});
     double draft_exec_ms = now_ms() - dt0;
     if (!draft_result.ok()) {
       ET_LOG(Error, "Draft forward failed at round %lld", (long long)rounds);
@@ -215,15 +244,21 @@ int main(int argc, char** argv) {
 
     std::vector<int64_t> verify_input_vec;
     verify_input_vec.push_back(last_token);
-    verify_input_vec.insert(verify_input_vec.end(), draft_ids.begin(), draft_ids.end());
+    verify_input_vec.insert(
+        verify_input_vec.end(), draft_ids.begin(), draft_ids.end());
     int64_t verify_len = static_cast<int64_t>(verify_input_vec.size());
     auto verify_input_tensor = make_tensor_ptr(
-        {1, (int)verify_len}, verify_input_vec.data(), executorch::aten::ScalarType::Long);
+        {1, (int)verify_len},
+        verify_input_vec.data(),
+        executorch::aten::ScalarType::Long);
 
     std::vector<int64_t> verify_pos_vec(verify_len);
-    for (int64_t i = 0; i < verify_len; ++i) verify_pos_vec[i] = pos + i;
+    for (int64_t i = 0; i < verify_len; ++i)
+      verify_pos_vec[i] = pos + i;
     auto verify_pos_tensor = make_tensor_ptr(
-        {(int)verify_len}, verify_pos_vec.data(), executorch::aten::ScalarType::Long);
+        {(int)verify_len},
+        verify_pos_vec.data(),
+        executorch::aten::ScalarType::Long);
 
     double vt0 = now_ms();
     auto verify_result = target_module.forward(
@@ -247,14 +282,16 @@ int main(int argc, char** argv) {
           (long long)hidden_len);
     }
 
-    std::vector<int64_t> new_tokens(draft_ids.begin(), draft_ids.begin() + accepted);
+    std::vector<int64_t> new_tokens(
+        draft_ids.begin(), draft_ids.begin() + accepted);
     new_tokens.push_back(target_ids[accepted]);
 
     bool hit_eos = false;
     for (size_t i = 0; i < new_tokens.size(); ++i) {
       if (kEosIds.count(new_tokens[i])) {
         new_tokens.resize(i + 1);
-        accepted = std::min(accepted, static_cast<int64_t>(new_tokens.size()) - 1);
+        accepted =
+            std::min(accepted, static_cast<int64_t>(new_tokens.size()) - 1);
         hit_eos = true;
         break;
       }
@@ -274,7 +311,8 @@ int main(int argc, char** argv) {
         new_hidden_fp32.begin() + append_len * hidden_concat_dim);
     hidden_len += append_len;
 
-    if (hit_eos) break;
+    if (hit_eos)
+      break;
   }
 
   double total_ms = now_ms() - t0;

examples/models/gemma4_31b/dflash_export.py

@@ -10,11 +10,11 @@
 import torch.nn as nn
 
 from executorch.examples.models.gemma4_31b.mlx_source_transformations import (
+    _replace_attention_forward,
+    _replace_layer_forward,
     MLXKVCache,
     MLXRingKVCache,
     MLXTurboQuantKVCache,
-    _replace_attention_forward,
-    _replace_layer_forward,
 )
 from executorch.examples.models.gemma4_31b.model import Gemma4_31B, Gemma4_31BConfig
 
@@ -32,8 +32,8 @@ def forward(
         tokens: torch.LongTensor,
         input_pos: torch.LongTensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Returns full-sequence logits and hidden states for DFlash block verification. 
-        Unlike the base implementation, which only returns the last token logits for single-token decoding, DFlash returns logits for every position in the drafted block so run_dflash.py can perform greedy verification and first-mismatch acceptance. 
+        """Returns full-sequence logits and hidden states for DFlash block verification.
+        Unlike the base implementation, which only returns the last token logits for single-token decoding, DFlash returns logits for every position in the drafted block so run_dflash.py can perform greedy verification and first-mismatch acceptance.
         """
         x = self.embed_tokens(tokens) * self.embed_normalizer
         sliding_mask, full_mask = self._build_masks(input_pos)
@@ -65,9 +65,7 @@ def _replace_dflash_model_forward(model: nn.Module) -> None:
     The stock mlx_source_transformations' _replace_model_forward overwrites the model's forward entirely, discarding Gemma4_31BWithHidden.forward, so this reinstalls the same hidden-capturing logic on top of the MLX-optimized layers.
     """
 
-    def _mlx_dflash_model_forward(
-        self, tokens: torch.Tensor, input_pos: torch.Tensor
-    ):
+    def _mlx_dflash_model_forward(self, tokens: torch.Tensor, input_pos: torch.Tensor):
         x = self.embed_tokens(tokens) * self.embed_normalizer
 
         layer_id_set = set(self.dflash_layer_ids)
@@ -154,4 +152,4 @@ def dflash_mlx_source_transformations(
         _replace_attention_forward(attn)
         _replace_layer_forward(layer)
 
-    _replace_dflash_model_forward(model)
+    _replace_dflash_model_forward(model)

examples/models/gemma4_31b/export_dflash_draft.py

@@ -5,6 +5,7 @@
 
 import argparse
 import json
+from pathlib import Path
 
 import torch
 
@@ -14,7 +15,6 @@
 )
 from executorch.examples.models.gemma4_31b.quant.quantize import dequantize_weight
 from huggingface_hub import snapshot_download
-from pathlib import Path
 from safetensors import safe_open
 from safetensors.torch import load_file
 from torch.export import Dim
@@ -23,8 +23,7 @@
 def _load_dequantized_tensor(
     safetensors_path: str, logical_name: str, dtype: torch.dtype = torch.bfloat16
 ) -> torch.Tensor:
-    """Loads and dequantizes a named weight from a torchao-quantized safetensors checkpoint into a dense tensor. 
-    """
+    """Loads and dequantizes a named weight from a torchao-quantized safetensors checkpoint into a dense tensor."""
     from torchao.prototype.safetensors.safetensors_support import (
         unflatten_tensor_state_dict,
     )
@@ -43,9 +42,7 @@ def _load_dequantized_tensor(
         prefix = f"{module_fqn}._{weight_name}_"
         partial = {k: f.get_tensor(k) for k in all_keys if k.startswith(prefix)}
         if not partial:
-            raise KeyError(
-                f"No keys found with prefix {prefix!r} for {logical_name!r}"
-            )
+            raise KeyError(f"No keys found with prefix {prefix!r} for {logical_name!r}")
         result, _ = unflatten_tensor_state_dict(partial, metadata)
 
     reconstructed = result[logical_name]
@@ -102,7 +99,7 @@ def main():
     )
     model.eval()
 
-    # Quantize the draft model to match the target model. 
+    # Quantize the draft model to match the target model.
     from executorch.backends.mlx.llm.quantization import quantize_model_
 
     quantize_model_(

examples/models/gemma4_31b/export_dflash_target.py

@@ -5,14 +5,11 @@
 
 import argparse
 import gc
-import json
 import os
 
 import torch
 
-from executorch.examples.models.gemma4_31b.dflash_export import (
-    Gemma4_31BWithHidden,
-)
+from executorch.examples.models.gemma4_31b.dflash_export import Gemma4_31BWithHidden
 from executorch.examples.models.gemma4_31b.export import _pack_for_backend
 from executorch.examples.models.gemma4_31b.model import (
     Gemma4_31BConfig,
@@ -25,8 +22,7 @@ def load_prequantized_dflash_target(
     layer_ids: list,
     max_seq_len: int = 4096,
 ) -> tuple:
-    """Loads the prequantized target model with DFlash hidden-state outputs.
-    """
+    """Loads the prequantized target model with DFlash hidden-state outputs."""
     config = Gemma4_31BConfig.from_hf_config(
         os.path.join(prequantized_dir, "config.json")
     )
@@ -53,11 +49,9 @@ def export_dflash_target_mlx(
     config: Gemma4_31BConfig,
     output_dir: str,
 ) -> None:
-    """Exports the DFlash target model through torch.export and the MLX backend.
-    """
+    """Exports the DFlash target model through torch.export and the MLX backend."""
     import executorch.backends.mlx.custom_kernel_ops.gguf.patterns
-    import executorch.extension.llm.export.gguf
-    import executorch.extension.llm.export.int4
+    import executorch.extension.llm.export.gguf  # noqa: F401
 
     from executorch.backends.mlx import MLXPartitioner
     from executorch.backends.mlx.passes import get_default_passes
@@ -161,7 +155,11 @@ def main() -> None:
             "Comma separated 0-indexed layer IDs matching the draft checkpoint, e.g. 1,12,23,35,46,57."
         ),
     )
-    p.add_argument("--output-dir", required=True, help="Directory for the exported .pte and .ptd files.")
+    p.add_argument(
+        "--output-dir",
+        required=True,
+        help="Directory for the exported .pte and .ptd files.",
+    )
     p.add_argument("--max-seq-len", type=int, default=4096)
     args = p.parse_args()