From ff75e5c96ec5a87f8ee1fefcc8c918530e07cfe4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Apr 2026 06:52:20 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/benchmarks/qwen3_flex_inference.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/scripts/benchmarks/qwen3_flex_inference.py b/scripts/benchmarks/qwen3_flex_inference.py index eb3dfef9e1..e6b8dcf2ea 100644 --- a/scripts/benchmarks/qwen3_flex_inference.py +++ b/scripts/benchmarks/qwen3_flex_inference.py @@ -137,8 +137,12 @@ def make_flex_attention_forward(page_table: PageTable): if hasattr(self, "q_norm"): # Qwen3: RMSNorm on [B, S, H, D] (per-head), then transpose. - q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) - k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose( + 1, 2 + ) + k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose( + 1, 2 + ) else: # Llama: no QK-norm. q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)