minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

perf(qwen3_5): replace einops rearrange with torch.flatten in GatedDe… (#20386)

Browse Source
This commit is contained in:
Vedant V Jhaveri
2026-03-12 09:51:27 +08:00
committed by GitHub
parent 25bd83033d
commit 9b55a98a67
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/qwen3_5.py
+1 -2
View File
@@ -20,7 +20,6 @@ from typing import Iterable, Optional, Set, Tuple, Union
import torch import torch
import torch.nn as nn import torch.nn as nn
from einops import rearrange
# Configs # Configs
from sglang.srt.configs.qwen3_5 import ( from sglang.srt.configs.qwen3_5 import (
@@ -287,7 +286,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
z = z.reshape(-1, z.shape[-1]) z = z.reshape(-1, z.shape[-1])
core_attn_out = self.norm(core_attn_out, z) core_attn_out = self.norm(core_attn_out, z)
core_attn_out = core_attn_out.reshape(z_shape_og) 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) output, _ = self.out_proj(core_attn_out)
return output return output