chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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-11 18:51:27 -07: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

3
python/sglang/srt/models/qwen3_5.py
View File

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