Avoid materializing GDN QKV tensors during target verification

[vedantjh2]

Summary

This PR removes redundant Q/K/V materialization from compatible GDN
speculative target-verification paths.

causal_conv1d_update already produces packed QKV. Previously, every GDN layer
launched fused_qkv_split_gdn_prefill_kernel to copy that output into three
contiguous tensors before Triton target verification. Triton accepts explicit
token strides, so it can consume zero-copy torch.split/view tensors instead.

The optimization is route-aware:

• Triton and ReplaySSM target verification use strided Q/K/V views.
• Prefill, FlashInfer, CuTeDSL, and other routes requiring contiguous inputs
  retain materialization.
• Model arithmetic, recurrent-state updates, speculative acceptance, logits,
  and sampling are unchanged.

This eliminates one memory-copy kernel launch per GDN layer on every compatible
target-verification pass.

Profile evidence

Matched H200 CPU/GPU traces using Qwen3.5-4B, ReplaySSM, and NEXTN T=3:

Target-verification metric	Baseline	Patch	Change
QKV-split launches/pass	24	0	-100%
QKV-split GPU time/pass	56.946 us	0	-56.946 us
CUDA-graph kernels	452	428	-24
CUDA-graph span	8,132.684 us	8,040.898 us	-91.786 us / -1.129%

The graph-span reduction's approximate 95% interval was
79.732-103.841 us.

Fixed input/output benchmark

The throughput results below use sglang.benchmark.one_batch_server, not
GSM8K. Unlike the in-process one_batch path, this runs against a real
ReplaySSM/NEXTN server and therefore exercises the speculative target-
verification route changed by this PR.

The baseline and patch were built from the same PR head; the baseline reverses
only this PR's six-file diff. Two alternating baseline/patch pairs ran on the
same H200 with a fresh server for every phase:

baseline -> patch -> baseline -> patch

Each phase used one warmup followed by five measured single-batch requests per
shape. Inputs were exact random token IDs, outputs used exact lengths with EOS
ignored, and the server and benchmark seeds were fixed. The table aggregates
ten measurements per arm and shape.

Input / output / batch	Baseline output tok/s	Patch output tok/s	Output change	Baseline latency	Patch latency	Latency change	Pair 1 / Pair 2 output change	Acceptance length
256 / 256 / 1	432.089	435.374	+0.760%	0.6297 s	0.6250 s	-0.761%	+0.855% / +0.666%	2.30 / 2.30
256 / 256 / 8	3,097.692	3,132.350	+1.119%	0.6990 s	0.6913 s	-1.096%	+1.689% / +0.554%	3.38 / 3.38
256 / 256 / 32	9,555.933	9,596.790	+0.428%	0.9812 s	0.9774 s	-0.388%	+0.402% / +0.453%	3.66 / 3.66
256 / 256 / 64	14,424.606	14,498.402	+0.512%	1.3813 s	1.3745 s	-0.487%	+0.421% / +0.603%	3.70 / 3.70
1024 / 256 / 64	14,828.231	14,874.086	+0.309%	2.0493 s	2.0476 s	-0.081%	+0.407% / +0.212%	3.79 / 3.79

At the representative 256/256 batch-64 shape, the paired-run 95% intervals
were +0.351% to +0.673% for output throughput and -0.642% to -0.330%
for total latency. Speculative acceptance was identical in every matched run.

Prior GSM8K end-to-end validation

These earlier GSM8K results are retained as secondary end-to-end and quality
context. Because GSM8K has variable output lengths, its throughput numbers are
not used as the primary performance evidence for this PR; the fixed
input/output results above are the primary benchmark.

Three alternating baseline/patch pairs ran on the same H200 with fresh
servers, seed 0, an empty prefix cache, and a 30-second cooldown. The benchmark
used its defaults: 200 questions, five shots, 512 maximum output tokens,
temperature 0, and parallelism 64.

Pair 1

Metric	Baseline	Patch	Change
Output throughput	6,433.642 tok/s	6,762.467 tok/s	+5.111%
Latency	5.379 s	5.145 s	-4.350%
Engine TPOT	9.7374 ms	9.4735 ms	-2.710%

Pair 2

Metric	Baseline	Patch	Change
Output throughput	6,202.258 tok/s	6,366.221 tok/s	+2.644%
Latency	5.470 s	5.401 s	-1.261%
Engine TPOT	10.2534 ms	9.7761 ms	-4.655%

Pair 3

Metric	Baseline	Patch	Change
Output throughput	6,496.073 tok/s	6,669.102 tok/s	+2.664%
Latency	5.270 s	5.134 s	-2.581%
Engine TPOT	9.3372 ms	9.4393 ms	+1.093%

Three-pair aggregate

Metric	Baseline	Patch	Change
Output throughput	6,377.324 tok/s	6,599.263 tok/s	+3.480%
Latency	5.373 s	5.227 s	-2.723%
Engine TPOT	9.7760 ms	9.5630 ms	-2.179%
Accuracy	87.667%	87.500%	-0.167 pp

GSM8K reproduction

Start the server:

CUDA_VISIBLE_DEVICES=0 \
sglang serve Qwen/Qwen3.5-4B \
  --port 30000 \
  --dtype bfloat16 \
  --language-only \
  --limit-mm-data-per-request '{"image":0,"video":0,"audio":0}' \
  --context-length 32768 \
  --mem-fraction-static 0.8 \
  --max-running-requests 64 \
  --linear-attn-decode-backend triton \
  --random-seed 0 \
  --speculative-algorithm NEXTN \
  --speculative-draft-model-path Qwen/Qwen3.5-4B \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --enable-linear-replayssm-spec

Run the official SGLang GSM8K benchmark:

python benchmark/gsm8k/bench_sglang.py \
  --data-path /data/gsm8k/test.jsonl

Dataset SHA-256:

3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14

Reproduction

The measurements above used PR head
f0b834be9b89f316f2daed9dc6f78e06df300e4b and base reference
d12b313b93e1547d9b02c3a84426aa88519fc494. The two benchmark trees can be
constructed without checking out an older, otherwise different SGLang tree:

git clone https://github.com/sgl-project/sglang.git pr33778-repro
cd pr33778-repro
git fetch origin pull/33778/head

PR_HEAD=f0b834be9b89f316f2daed9dc6f78e06df300e4b
BASE_REF=d12b313b93e1547d9b02c3a84426aa88519fc494
FILES=(
  python/sglang/srt/layers/attention/linear/gdn_backend.py
  python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py
  python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
  test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py
  test/registered/attention/unittests/gdn/test_triton.py
  test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
)

git diff "${BASE_REF}" "${PR_HEAD}" -- "${FILES[@]}" > ../pr33778.patch
echo "326deb7426602b0fbcf3aa436b3a6b5e34053c9f8be85f9e50d02b7149881006  ../pr33778.patch" \
  | sha256sum --check

git worktree add --detach ../sglang-pr "${PR_HEAD}"
git worktree add --detach ../sglang-base "${PR_HEAD}"
git -C ../sglang-base apply --reverse ../pr33778.patch

# Both trees otherwise remain at the identical PR head.
git -C ../sglang-pr rev-parse HEAD
git -C ../sglang-base rev-parse HEAD
git -C ../sglang-base diff --name-only

The measured software versions were Python 3.12.6, PyTorch 2.13.0+cu130,
Triton 3.7.1, Transformers 5.12.1, and sglang-kernel 0.4.6.post1. A fresh
environment can be prepared from the PR tree with:

python3.12 -m venv ../pr33778-venv
source ../pr33778-venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e '../sglang-pr/python[test]'

Start the server from either the baseline or patch source tree:

TREE=../sglang-base  # Change to ../sglang-pr for the patch phase.

CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH="${TREE}/python" \
python -m sglang.launch_server \
  --model-path Qwen/Qwen3.5-4B \
  --port 30000 \
  --dtype bfloat16 \
  --language-only \
  --limit-mm-data-per-request '{"image":0,"video":0,"audio":0}' \
  --context-length 4096 \
  --mem-fraction-static 0.75 \
  --max-running-requests 64 \
  --cuda-graph-max-bs-decode 64 \
  --disable-radix-cache \
  --linear-attn-decode-backend triton \
  --random-seed 0 \
  --speculative-algorithm NEXTN \
  --speculative-draft-model-path Qwen/Qwen3.5-4B \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --enable-linear-replayssm-spec

Run one warmup and five measured iterations for each fixed shape:

MODEL=Qwen/Qwen3.5-4B
BASE_URL=http://127.0.0.1:30000

for shape in "256 256 1" "256 256 8" "256 256 32" \
             "256 256 64" "1024 256 64"; do
  read -r INPUT_LEN OUTPUT_LEN BATCH_SIZE <<< "${shape}"
  NAME="i${INPUT_LEN}-o${OUTPUT_LEN}-b${BATCH_SIZE}"

  python -m sglang.benchmark.one_batch_server \
    --model-path "${MODEL}" \
    --local-tokenizer-path "${MODEL}" \
    --base-url "${BASE_URL}" \
    --backend sglang \
    --dataset-name random-ids \
    --batch-size "${BATCH_SIZE}" \
    --input-len "${INPUT_LEN}" \
    --output-len "${OUTPUT_LEN}" \
    --temperature 0 \
    --seed 42 \
    --skip-warmup \
    --no-append-to-github-summary \
    --result-filename "${NAME}-warmup.jsonl"

  for REPEAT in 1 2 3 4 5; do
    python -m sglang.benchmark.one_batch_server \
      --model-path "${MODEL}" \
      --local-tokenizer-path "${MODEL}" \
      --base-url "${BASE_URL}" \
      --backend sglang \
      --dataset-name random-ids \
      --batch-size "${BATCH_SIZE}" \
      --input-len "${INPUT_LEN}" \
      --output-len "${OUTPUT_LEN}" \
      --temperature 0 \
      --seed 42 \
      --skip-warmup \
      --no-append-to-github-summary \
      --result-filename "${NAME}-${REPEAT}.jsonl"
  done
done

Restart the server between baseline and patch phases and run the order shown
above.

The following is the complete four-phase command used to automate that process.
It saves and terminates only the server process group it starts:

cat > /tmp/run-pr33778-fixed.sh <<'BASH'
#!/usr/bin/env bash
set -euo pipefail

MODEL=Qwen/Qwen3.5-4B
PORT=30000
RESULT_ROOT=${RESULT_ROOT:-"$PWD/pr33778-results"}
PHASES=(base pr base pr)
TREES=(
  "$PWD/../sglang-base"
  "$PWD/../sglang-pr"
  "$PWD/../sglang-base"
  "$PWD/../sglang-pr"
)
SHAPES=(
  "256 256 1"
  "256 256 8"
  "256 256 32"
  "256 256 64"
  "1024 256 64"
)
SERVER_PID=

stop_server() {
  if [[ -n "${SERVER_PID}" ]] &&
     ps -o pid= -g "${SERVER_PID}" 2>/dev/null | grep -q '[0-9]'; then
    kill -TERM -- "-${SERVER_PID}"
    for _ in $(seq 1 60); do
      ps -o pid= -g "${SERVER_PID}" 2>/dev/null | grep -q '[0-9]' || break
      sleep 1
    done
    if ps -o pid= -g "${SERVER_PID}" 2>/dev/null | grep -q '[0-9]'; then
      kill -KILL -- "-${SERVER_PID}"
    fi
    wait "${SERVER_PID}" 2>/dev/null || true
  fi
  SERVER_PID=
}
trap stop_server EXIT INT TERM

start_server() {
  local tree=$1
  local phase_dir=$2
  setsid env \
    CUDA_VISIBLE_DEVICES=0 \
    PYTHONPATH="${tree}/python" \
    python -m sglang.launch_server \
      --model-path "${MODEL}" \
      --port "${PORT}" \
      --dtype bfloat16 \
      --language-only \
      --limit-mm-data-per-request '{"image":0,"video":0,"audio":0}' \
      --context-length 4096 \
      --mem-fraction-static 0.75 \
      --max-running-requests 64 \
      --cuda-graph-max-bs-decode 64 \
      --disable-radix-cache \
      --linear-attn-decode-backend triton \
      --random-seed 0 \
      --speculative-algorithm NEXTN \
      --speculative-draft-model-path "${MODEL}" \
      --speculative-num-steps 3 \
      --speculative-eagle-topk 1 \
      --speculative-num-draft-tokens 4 \
      --enable-linear-replayssm-spec \
      > "${phase_dir}/server.log" 2>&1 &
  SERVER_PID=$!

  for _ in $(seq 1 240); do
    curl -fsS --max-time 2 \
      "http://127.0.0.1:${PORT}/health_generate" >/dev/null && return
    kill -0 "${SERVER_PID}" 2>/dev/null || {
      tail -100 "${phase_dir}/server.log"
      return 1
    }
    sleep 5
  done
  echo "Server startup timed out" >&2
  return 1
}

run_case() {
  local tree=$1 input_len=$2 output_len=$3 batch_size=$4 output=$5
  env PYTHONPATH="${tree}/python" \
    python -m sglang.benchmark.one_batch_server \
      --model-path "${MODEL}" \
      --local-tokenizer-path "${MODEL}" \
      --base-url "http://127.0.0.1:${PORT}" \
      --backend sglang \
      --dataset-name random-ids \
      --batch-size "${batch_size}" \
      --input-len "${input_len}" \
      --output-len "${output_len}" \
      --temperature 0 \
      --seed 42 \
      --skip-warmup \
      --no-append-to-github-summary \
      --result-filename "${output}"
}

mkdir -p "${RESULT_ROOT}"
for phase_index in "${!PHASES[@]}"; do
  source_label=${PHASES[phase_index]}
  tree=${TREES[phase_index]}
  phase_number=$((phase_index + 1))
  phase_dir=$(printf '%s/phase-%02d-%s' \
    "${RESULT_ROOT}" "${phase_number}" "${source_label}")
  mkdir -p "${phase_dir}"
  start_server "${tree}" "${phase_dir}"

  for shape in "${SHAPES[@]}"; do
    read -r input_len output_len batch_size <<< "${shape}"
    shape_dir="${phase_dir}/i${input_len}-o${output_len}-b${batch_size}"
    mkdir -p "${shape_dir}"

    run_case "${tree}" "${input_len}" "${output_len}" "${batch_size}" \
      "${shape_dir}/warmup.jsonl"
    for repeat in 1 2 3 4 5; do
      run_case "${tree}" "${input_len}" "${output_len}" "${batch_size}" \
        "$(printf '%s/result-%02d.jsonl' "${shape_dir}" "${repeat}")"
      sleep 2
    done
  done

  stop_server
  sleep 30
done
BASH

chmod +x /tmp/run-pr33778-fixed.sh
/tmp/run-pr33778-fixed.sh

Aggregate the ten measurements per arm and verify acceptance parity:

python - <<'PY'
import json
import statistics
from collections import defaultdict
from pathlib import Path

root = Path("pr33778-results")
grouped = defaultdict(list)
for path in root.glob("phase-*/i*-o*-b*/result-*.jsonl"):
    phase = path.parts[-3]
    source = phase.rsplit("-", 1)[-1]
    shape = path.parts[-2]
    grouped[source, shape].append(json.loads(path.read_text()))

print("shape,base_output_tps,patch_output_tps,output_change_pct,"
      "base_latency_s,patch_latency_s,latency_change_pct,"
      "base_acceptance,patch_acceptance")
shapes = sorted(shape for source, shape in grouped if source == "base")
for shape in shapes:
    base = grouped["base", shape]
    patch = grouped["pr", shape]
    assert len(base) == len(patch) == 10

    def mean(records, key):
        return statistics.fmean(record[key] for record in records)

    base_output = mean(base, "output_throughput")
    patch_output = mean(patch, "output_throughput")
    base_latency = mean(base, "latency")
    patch_latency = mean(patch, "latency")
    base_accept = mean(base, "acc_length")
    patch_accept = mean(patch, "acc_length")
    assert base_accept == patch_accept

    print(
        f"{shape},{base_output:.3f},{patch_output:.3f},"
        f"{100 * (patch_output / base_output - 1):.3f},"
        f"{base_latency:.4f},{patch_latency:.4f},"
        f"{100 * (patch_latency / base_latency - 1):.3f},"
        f"{base_accept:.2f},{patch_accept:.2f}"
    )
PY

Generate a fixed-shape CPU/GPU trace against either running phase with:

mkdir -p pr33778-traces
PYTHONPATH="${TREE}/python" \
python -m sglang.benchmark.one_batch_server \
  --model-path Qwen/Qwen3.5-4B \
  --local-tokenizer-path Qwen/Qwen3.5-4B \
  --base-url http://127.0.0.1:30000 \
  --backend sglang \
  --dataset-name random-ids \
  --batch-size 64 \
  --input-len 256 \
  --output-len 256 \
  --temperature 0 \
  --seed 42 \
  --skip-warmup \
  --profile \
  --profile-activities CPU GPU \
  --profile-steps 5 \
  --profile-prefix pr33778-b64 \
  --profile-output-dir "$PWD/pr33778-traces" \
  --no-append-to-github-summary \
  --result-filename pr33778-profile.jsonl

Run the focused H200 correctness suite from the patch tree:

cd ../sglang-pr
CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH="$PWD/python" \
python -m pytest -q \
  test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py \
  test/registered/attention/unittests/gdn/test_triton.py \
  test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py

git diff --check \
  d12b313b93e1547d9b02c3a84426aa88519fc494...\
f0b834be9b89f316f2daed9dc6f78e06df300e4b

Correctness validation

• Exact target-verification output and recurrent-state parity.
• ReplaySSM ring/fold parity and routing coverage.
• Contiguous-input behavior retained for unsupported routes.
• 30 focused tests and 55 subtests passed on H200.
• The prior GSM8K section is retained as secondary end-to-end and quality
  context; the fixed input/output benchmark is the primary performance
  evidence.

---

CI States

Latest PR Test (Base): ✅ Run #35563025383
Latest PR Test (Extra): ❌ Run #35563025315
Latest PR Test (AMD ROCm 10): ⏳ Run #35563025321

[Qiaolin-Yu]

/tag-and-rerun-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-skipped-ci

[vedantjh2]

/rerun-skipped-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[Qiaolin-Yu]

when qkv_dim > MAX_FUSED_QKV_SPLIT_DIM, will the consumer side silently get wrong output?

[vedantjh2]

yes this can happen for larger models at low TP but it is pre-existing behavior; the  qkv_dim > limit  path already used strided split views before this change and this PR does not change that path.

[vedantjh2]

@Qiaolin-Yu should i add something like this here?

if not use_strided_target_verify_qkv:
    query = query.contiguous()
    key = key.contiguous()
    value = value.contiguous()

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[b8zhong]

Hi, for performance benchmarking, we need fixed input/output benchmark, like bench_one_batch or similar, since we can't take the throughput number from GSM8K

[vedantjh2]

hi @b8zhong, thanks for the review! Just added additional evidence for bench_one_batch with multiple batch sizes showing positive gains. please let me know if you need anything else!

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci

[vedantjh2]

/rerun-failed-ci