Motivation

einops.rearrange performs Python-level string parsing and validation on every call. In Qwen3_5GatedDeltaNet.forward(), the pattern rearrange(x, "... h d -> ... (h d)") is a simple flatten of the last two dimensions — equivalent to torch.flatten(-2), which is a zero-copy view operation with no parsing overhead.

This call runs on every forward pass through every GDN layer in the model, so the overhead accumulates significantly.

Changes

• Replace rearrange(core_attn_out, "... h d -> ... (h d)") with core_attn_out.flatten(-2) in Qwen3_5GatedDeltaNet.forward()
• Remove unused from einops import rearrange import

Profiling Results

Profiled on Qwen3.5-9B, H100 80GB, TP=1, CUDA graphs disabled, using sglang.bench_serving with torch profiler. Perfetto SQL queries filtered to only GatedDeltaNet ancestor slices.

Baseline (einops)

Optimized (torch.flatten)

Summary

• Per-call: 2.67x faster (12.67 us → 4.74 us)
• Total time: 63% reduction (9125 us → 3410 us across 720 calls)
• 5715 us saved per profiled window

Profiling Reproduction Commands

1. Launch server

export SGLANG_TORCH_PROFILER_DIR=/path/to/traces

python -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-9B \
  --port 30000 --tp 1 --disable-cuda-graph

2. Run profiling benchmark

export PYTHONPATH=/path/to/sglang/python:$PYTHONPATH
export SGLANG_TORCH_PROFILER_DIR=/path/to/traces

python -m sglang.bench_serving \
  --backend sglang \
  --model Qwen/Qwen3.5-9B \
  --num-prompts 10 --random-input-len 256 --random-output-len 32 \
  --dataset-name random --profile

3. Perfetto SQL queries

Baseline (einops rearrange in GatedDeltaNet):

SELECT s.name, COUNT(*) AS call_count,
       AVG(s.dur)/1000.0 AS avg_dur_us, SUM(s.dur)/1000.0 AS total_dur_us
FROM slice s
WHERE s.name LIKE '%rearrange%'
  AND EXISTS (SELECT 1 FROM ancestor_slice(s.id) a WHERE a.name LIKE '%GatedDeltaNet%')
GROUP BY s.name;

Optimized (flatten in GatedDeltaNet):

SELECT s.name, COUNT(*) AS call_count,
       AVG(s.dur)/1000.0 AS avg_dur_us, SUM(s.dur)/1000.0 AS total_dur_us,
       MIN(s.dur)/1000.0 AS min_dur_us, MAX(s.dur)/1000.0 AS max_dur_us
FROM slice s
WHERE s.name LIKE '%flatten%'
  AND EXISTS (SELECT 1 FROM ancestor_slice(s.id) a WHERE a.name LIKE '%GatedDeltaNet%')
GROUP BY s.name;

Correctness

flatten(-2) is semantically identical to rearrange("... h d -> ... (h d)") — both merge the last two dimensions into one. No behavior change.

Ran GSM8k to test performance and accuracy

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.