ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-04-21 06:52:20 +00:00
commit ff75e5c96e
1 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

8
scripts/benchmarks/qwen3_flex_inference.py
View file

@ -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)