From 9b55a98a67059b30a2997f244221be1b0b4042c3 Mon Sep 17 00:00:00 2001 From: Vedant V Jhaveri Date: Wed, 11 Mar 2026 18:51:27 -0700 Subject: [PATCH] =?UTF-8?q?perf(qwen3=5F5):=20replace=20einops=20rearrange?= =?UTF-8?q?=20with=20torch.flatten=20in=20GatedDe=E2=80=A6=20(#20386)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/sglang/srt/models/qwen3_5.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 5b2128e68..73f22f9be 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -20,7 +20,6 @@ from typing import Iterable, Optional, Set, Tuple, Union import torch import torch.nn as nn -from einops import rearrange # Configs from sglang.srt.configs.qwen3_5 import ( @@ -287,7 +286,7 @@ class Qwen3_5GatedDeltaNet(nn.Module): z = z.reshape(-1, z.shape[-1]) core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) output, _ = self.out_proj(core_attn_out) return output