SID-GR Inference

July 22, 2026 ยท View on GitHub

1. Background and Motivation

Algorithm Background

Semantic ID based generative recommender modeling is one of the main directions for recommendation, search, and advertising systems. The core workflow is:

offline clustering
-> map each real item ID to a multi-level cluster ID tuple
-> autoregressively generate a short semantic ID sequence at inference time
-> map the generated semantic ID tuple back to real item IDs

This pattern has direct implications for inference systems:

  • The cluster depth is usually small, for example 3 or 5, so autoregressive decode is short.
  • User history can be long, so prefill/context computation can dominate cost.
  • Recommendation and search systems often need result diversity. Larger and sometimes dynamic beam widths are a practical way to improve diversity.

Inference Problem

The resulting inference workload is:

long context + short decode + large beam width

This is different from the common chat LLM serving workload. vLLM, SGLang, and TensorRT-LLM are optimized primarily for:

many user requests
dynamic batching
paged KV
long decode
OpenAI/chat APIs

Those capabilities matter, but they are not the full bottleneck for SID-GR inference. In recommendation and search workloads, a typical request has a long user or candidate context, a large beam width such as 128 or 256, and only a few decode steps. Many beams share the same request-level context.

Current general-purpose LLM inference frameworks do not fully match this pattern:

  • vLLM: does not provide a stable production beam-search serving path. Users often implement beam search by repeatedly calling vLLM from business logic.
  • TensorRT-LLM: does not natively expose logprobs for this path, and large beam widths can create memory pressure. Its decode attention and postprocess kernels are also not specialized for short SID-GR decode.
  • SGLang: is currently the most usable open-source baseline for this workload, but large-beam support is still in a feature PR and is not merged to upstream main.

General LLM serving frameworks are also large and complex. Adapting that full stack for a specialized recommender inference workload can work for one business case, but it is not the best long-term path for maintainability or for reaching the speed-of-light target of this workload.

2. Goals and Approach

Goals

  • Optimize Qwen-family models toward speed-of-light performance for the SID-GR workload: long context, short decode, and large beam width.
  • Provide a compact framework that supports practical SID-GR serving needs, including both feature requirements and performance requirements.

Approach

The implementation keeps SID-GR specific runtime contracts while selectively reusing mature ideas from open-source serving systems:

  • Keep SID-GR native abstractions. ContextKV, BeamKV, BeamPath, dynamic beam policies, item-constrained decode, and request x active-beam batching are represented directly by the runtime.
  • Reuse proven serving ideas selectively. The project borrows concepts from vLLM and SGLang for continuous batching, paged KV, HTTP serving, benchmark tooling, and APIs. It also borrows from TensorRT-LLM for kernels, CUDA graph, operator fusion, and model-layer optimization. These pieces support the SID-GR runtime rather than defining it.
  • Drive changes with benchmarks. Correctness, performance, and Nsight breakdowns are used to decide which optimizations should enter the hot path.
  • Keep the framework small and specialized. General-purpose LLM serving is used as a source of ideas, while KV ownership, beam state, decode attention, batching units, and business constraints remain SID-GR specific.

Current Status

The repository has a single-node alpha path with real model weights:

Qwen3-1.7B real weights
+ SID-GR native ContextKV / BeamKV / BeamPath
+ real gr-decode_atten backend
+ continuous batching
+ BeamKV / ContextKV dense pools
+ direct pool-view decode CUDA graph
+ HTTP /generate
+ SGLang-equivalent beam_results output
+ offline and online SID-GR vs SGLang benchmarks

This path validates the main value proposition: for the tested long-context, short-decode, large-beam matrix, SID-GR offline performance is consistently faster than the SGLang beam-search PR branch. Online serving also runs through the same HTTP client benchmark. CUDA graph capture has been stabilized as a startup warmup path, with replay-only execution during the measured serving window.

3. Design Highlights

DimensionGeneral LLM serving pathSID-GR Inference path
KV abstractionSequence, token block, and paged KV centricExplicit request-level ContextKV plus short BeamKV
Large-beam decodeFlattens batch * beam into many decode rowsProcesses shared context by request and beam tile
Attention kernelGeneral paged decode attentiongr-decode_atten receives ContextKV + BeamKV + BeamPath directly
Batch CUDA graphGeneral batch buckets and paged KV constraintsFixed SID-GR shapes and stable pool slices for replay
OutputGeneral API output and beam managementFast path returns beam_results; debug paths can enable beam_details

Core implementation points:

  • SID-GR native KV layout. A request's long context is stored once in ContextKV. Short decode history is stored in BeamKV. BeamPath records logical parent-child relations.
  • Dense ContextKV hot path. The current ContextKV layout is dense and contiguous. This supports kernel-friendly decode attention and stable pool slices for CUDA graph replay.
  • BeamKV and ContextKV pools. Continuous serving uses dense pools with leases, capacity tracking, high-water marks, utilization metrics, and leak checks.
  • Specialized decode attention. gr-decode_atten understands request-level shared context and short beam history, instead of treating every beam as an independent generic decode row.
  • SID-GR continuous batching. The scheduler groups by decode step, beam width, and context shape. The batching unit is request x active beams.
  • Direct pool-view decode CUDA graph. Fixed-shape graphs bind stable ContextKV and BeamKV pool slices. Captured graphs replay on the serving path; dynamic non-contiguous KV slices fall back to eager execution.
  • Last-token logits only. Serving prefill computes only the last-position logits needed for the next token.
  • Dynamic beam policy support. The runtime supports fixed, scheduled, and score-margin beam policies, including request-level HTTP configuration.
  • Item-constrained generation support. Runtime and HTTP paths include item tries, masks, constrained topK, catalog reload/rollback, and item metadata.
  • Correctness and performance alignment. The fast benchmark path returns SGLang-equivalent beam_results; debug-rich beam_details is enabled only for correctness and debugging.

4. Default Production Path

  • Continuous serving uses the real gr-decode_atten backend when --decode-backend real is selected.
  • Eligible continuous decode batches use decode CUDA graph by default. Set GR_INFERENCE_DISABLE_DECODE_CUDA_GRAPH=1 to fall back to eager decode.
  • ContextKV and BeamKV use pool slices directly. Graph replay updates only small inputs such as beam token IDs and topK indices.
  • Decode graph cache uses an entry limit, LRU eviction, and pointer guards. A graph is not reused if pool slice addresses do not match.
  • Decode-graph captures share a single CUDA-graph private pool and one logits output buffer per (beam_width, bucket), so capture memory stays roughly flat as the number of captured graphs grows.
  • Online serving warms up the common batch, pool-window, and /generate ignore_eos shapes at startup, then freezes new graph capture by default.
  • QK norm and RoPE prefer the fastest available SGLang-style in-place kernels, with FlashInfer/Torch fallback. Experimental branches are not part of the serving hot path.
  • Prefill uses SGLang-style piecewise CUDA graph by default. The current Qwen3-1.7B bs4/ctx1000 shape captures six graph pieces: embed, four layer chunks, and output.
  • /generate returns SGLang-equivalent beam_results by default. When ignore_eos=true, tokenizer special tokens are suppressed by default to match SGLang fixed-length generation semantics.

Fallback to eager decode:

GR_INFERENCE_DISABLE_DECODE_CUDA_GRAPH=1

Disable SID-GR experimental JIT kernels:

GR_INFERENCE_GR_TRTLLM_KERNELS_JIT=0

Decode graph cache size can be configured with:

GR_INFERENCE_DECODE_CUDA_GRAPH_MAX_ENTRIES=32

By default, decode-graph (and prefill piecewise) captures share one private graph pool, and decode captures share one [batch, beam, vocab] logits output buffer per (beam_width, bucket). To opt out and give each capture its own private pool / per-entry logits (for debugging or A/B comparison), set:

GR_INFERENCE_DECODE_CUDA_GRAPH_SEPARATE_POOLS=1
GR_INFERENCE_PREFILL_CUDA_GRAPH_SEPARATE_POOLS=1
GR_INFERENCE_DECODE_CUDA_GRAPH_SEPARATE_LOGITS=1

5. Baseline

The headline numbers compare against the SGLang beam-search PR branch:

repo:   https://github.com/cswuyg/sglang.git
branch: feature/beam_search
PR:     https://github.com/sgl-project/sglang/pull/15645

Test environment and workload:

GPU: NVIDIA H100 80GB HBM3
model: Qwen3-1.7B
context_len: 1000 / 5000
beam_width: 256
effective output length: 3 tokens

6. Offline Performance

Performance measurement uses the default SID-GR output mode, which returns SGLang-equivalent beam_results and does not construct debug-rich beam_details. SID-GR enables prefill CUDA graph and direct pool-view decode CUDA graph. Prefix/prefill cache is disabled, and Qwen special tokens are suppressed when ignore_eos fixes output length. SGLang timing wraps the beam-search PR branch Engine.generate call with radix cache disabled.

Radix cache off, which matches the production no-prefix-reuse setting:

ctxbatchSID-GR msSGLang msSGLang/SID-GRwinnerSID-GR prefillSID-GR decode
1000117.61133.5151.903xSID-GR7.0279.952
1000227.76857.9262.086xSID-GR12.63314.170
1000447.736102.3182.143xSID-GR23.44222.566
1000893.230199.2802.138xSID-GR46.52143.707
5000142.25594.5542.238xSID-GR31.02910.701
5000280.904179.3692.217xSID-GR63.76316.216
50004154.224349.8572.269xSID-GR126.08726.772
50008307.917685.3542.226xSID-GR253.34551.448

Offline conclusion: SID-GR is faster than SGLang in all tested production no-prefix-reuse cases. For ctx=5000, batch=8, SID-GR is 2.23x faster than SGLang.

7. Online In-flight Serving

Online serving uses SGLang bench_serving as the shared HTTP client with request_rate=inf, max_concurrency=4, and requests=64. SID-GR serves the compatible /generate endpoint with default beam_results, prefill CUDA graph, direct pool-view decode CUDA graph, and SGLang-aligned special-token suppression for ignore_eos. SGLang uses the beam-search PR branch.

The measured command uses warmup_requests=0; external server startup warmup and priming are not counted.

SID-GR stable reproduction mode:

  • Warm up online batch sizes and KV pool slot windows during server startup.
  • Run warmup requests through the same ignore_eos=true and special-token suppression path as real /generate requests.
  • Freeze prefill and decode CUDA graph capture after warmup.
  • Skip CUDA graph for dynamic non-contiguous KV composition and use eager execution instead.

Fixed case:

ctx=5000, beam=256, output=3, requests=64, max_concurrency=4

Latest three reruns:

server / output moderoundreq/smedian msp90 msp99 msinput tok/soutput tok/s
SID-GR /generate, beam_results, frozen graph capture119.41199.32240.16323.759704514906
SID-GR /generate, beam_results, frozen graph capture220.10195.62223.38252.4010050515438
SID-GR /generate, beam_results, frozen graph capture319.66199.92235.61304.199829415098
SGLang /generate, beam results, primed steady110.69369.69374.14379.08534508208
SGLang /generate, beam results, primed steady210.65370.59373.94378.19532508177
SGLang /generate, beam results, primed steady310.68370.36373.68375.22534008201

SID-GR graph stability gate:

checkpointprefill capturesdecode capturescaptures enableddynamic graph skips
startup103000
round1103008
round21030011
round31030015

Online conclusion: with the stable reproduction mode, SID-GR reaches 19.41 / 20.10 / 19.66 req/s over three rounds. Decode CUDA graph captures stay fixed at 30 and no longer grow during online scheduling. Compared with the three primed-steady SGLang runs, SID-GR average throughput is about 1.85x higher and median latency is about 46% lower. p99 can still fluctuate due to the HTTP client, Python scheduling, request arrival timing, and batch fill.

Reproduce the SID-GR online stable mode:

BASE_OUT=benchmark_artifacts/sglang_compare/gr_online_repro
mkdir -p "${BASE_OUT}/gr"

env GR_MODEL_DIR=/workspace/models/Qwen3-1.7B \
  GR_CONTEXT_LEN=5000 \
  GR_DECODE_STEPS=3 \
  GR_BEAM_WIDTH=256 \
  GR_MAX_BATCH_SIZE=4 \
  GR_BEAM_KV_POOL_CAPACITY=4 \
  GR_CONTEXT_KV_POOL_CAPACITY=4 \
  GR_HTTP_HOST=0.0.0.0 \
  GR_HTTP_PORT=8000 \
  GR_DECODE_BACKEND=real \
  GR_DEVICE=cuda \
  GR_DECODE_CUDA_GRAPH_BATCH_BUCKETS=1,2,4,8 \
  GR_WARMUP_ONLINE_SHAPES=1 \
  GR_WARMUP_ONLINE_POOL_WINDOWS=1 \
  GR_WARMUP_ONLINE_MAX_CASES=64 \
  GR_FREEZE_CUDA_GRAPHS_AFTER_WARMUP=1 \
  GR_ENABLE_PREFILL_CACHE=0 \
  scripts/serve_qwen3_gr_http.sh \
  > "${BASE_OUT}/gr_server.log" 2>&1 &
SERVER_PID=$!

until curl -fsS http://127.0.0.1:8000/ready >/dev/null; do sleep 2; done
curl -fsS http://127.0.0.1:8000/metrics \
  -o "${BASE_OUT}/metrics_after_startup.json"

for round in 1 2 3; do
  OUT_DIR="${BASE_OUT}/gr/round${round}" \
  REQUESTS=64 CONTEXT_LEN=5000 DECODE_STEPS=3 BEAM_WIDTH=256 \
  REQUEST_RATE=inf MAX_CONCURRENCY=4 WARMUP_REQUESTS=0 \
  scripts/run_gr_sglang_bench_serving_beam_benchmark.sh \
  2>&1 | tee "${BASE_OUT}/gr_round${round}.log"
  curl -fsS http://127.0.0.1:8000/metrics \
    -o "${BASE_OUT}/metrics_after_round${round}.json"
done

kill "${SERVER_PID}"

Using scripts/serve_qwen3_gr_http.sh plus scripts/run_gr_sglang_bench_serving_beam_benchmark.sh follows this stable reproduction path by default. Set GR_FREEZE_CUDA_GRAPHS_AFTER_WARMUP=0 only when debugging graph coverage.

Online /generate top1 correctness smoke:

ctxbeamrequestsmax concurrencytop1 exacttopK overlap
500025664458/64min 0.918 / mean 0.956

This uses the real HTTP /generate path and default beam_results output. Requests whose top1 differs still have high TopK overlap.

8. Correctness

Correctness compares SID-GR default beam_results output against SGLang beam_results. SID-GR enables CUDA graph and suppresses Qwen special tokens to match SGLang ignore_eos=true fixed-length generation semantics.

ctxbatchtop1 exacttopK overlapnote
100011.000.949top1 exact
100021.000.953top1 exact
100041.000.956top1 exact
100081.000.958top1 exact
500011.000.969top1 exact
500021.000.961top1 exact
500041.000.955top1 exact
500081.000.959top1 exact

Offline correctness conclusion: in the production no-prefix-cache mode, all eight fixed beam=256 cases have exact Top1 agreement and high TopK overlap. For the full 24-case matrix, top1 min/mean is 1.000 / 1.000, and TopK overlap min/mean is 0.945 / 0.960.

9. Performance Breakdown

Fixed case:

ctx=1000, beam=256, batch=4, output=3

Nsight profile summary:

metricSID-GRSGLang
active CUDA window46.856 ms99.944 ms
kernel total43.168 ms78.859 ms
CUDA runtime API total42.886 ms42.975 ms
CPU runtime gaps >50us2.186 ms28.305 ms
CUDA graph launches, active window total829
kernel launch count12611620

Prefill stage:

metricSID-GRSGLang
stage total24.404 ms27.549 ms
attention kernels1.367 ms1.389 ms
non-attention kernels20.825 ms19.102 ms
CPU overhead2.213 ms7.057 ms
CUDA graph launches629

Decode stage:

metricSID-GRSGLang
stage total23.318 ms55.663 ms
attention kernels1.593 ms35.235 ms
non-attention kernels19.384 ms13.228 ms
CPU overhead2.341 ms7.200 ms
CUDA graph launches20

Additional kernel buckets:

metricSID-GRSGLang
topK / beam selection4.027 ms4.716 ms
attention bucket total2.960 ms40.951 ms

The main latency tables should be used for end-to-end latency. Nsight is used here to explain the active CUDA window, stage split, and kernel breakdown. The raw Nsight output is in:

benchmark_artifacts/sglang_compare/prod_breakdown_ctx1000_beam256_b4_20260525_121548/

For batch 4, the main gap is not topK sorting. It comes from the large-beam decode attention path and general serving overhead.

decode CUDA graph shortens the fixed-shape path;
SID-GR decode attention reduces core compute.

In this case, decode attention kernels differ by about 35.235 - 1.593 = 33.6 ms, and the active CUDA window differs by about 99.944 - 46.856 = 53.1 ms. The largest gain still comes from SID-GR specific decode attention and KV structure.

SGLang decode attention sees:

batch * beam = 4 * 256 = 1024 decode rows

Each row follows a general paged decode attention path.

SID-GR decode attention preserves the workload structure:

4 requests
256 beams per request
one shared ContextKV per request
short BeamKV history per beam

The context part is processed by request and beam tile. The short BeamKV part attends only a few decode tokens. Decode CUDA graph further reduces launch and CPU scheduling overhead for fixed-shape decode steps.

In short, SGLang is a general beam serving path; SID-GR is specialized for long context, large beam, and short decode.

10. TODO and Roadmap

The single-node alpha core path is complete. The remaining work is productionization, broader validation, and framework maintainability.

AreaExisting foundationFollow-up work
Online serving hardeningHTTP /generate, background worker, continuous batching, pool metrics, online correctness/perf benchmarks, frozen CUDA graph capture after warmupImprove admission, batch fill, and tail latency; add request-rate, max-concurrency, arrival-pattern, and long soak regressions
ContextKV memory strategyDense ContextKV pool on offline/online hot path with stable pool slices for CUDA graphAdd multiple context buckets based on real context length distribution; integrate page-backed ContextKV storage; eventually support native page tables in decode attention
CUDA graph productionizationDirect pool-view decode graph enabled by default; startup warmup covers batch, pool window, and /generate ignore_eos shapesExpand warmup shapes; add graph coverage, fallback, eviction, and metric regressions
Beam selection graphDecode forward is already in graph; beam selection remains outside graphMove log_softmax + topK + beam selection into graph once pool ownership, item masks, special-token suppression, and output trimming are safe
SID-GR vs SGLang benchmarkFinal offline mode covers ctx=1000/5000, beam=256, batch=1/2/4/8; online HTTP benchmark has a stable recipeExtend context lengths, beam widths, dtypes, model sizes, and GPUs; provide one-command final offline/online summary scripts
Beam result output/generate returns SGLang-equivalent beam_results; debug-rich beam_details is opt-inFurther optimize Python construction and JSON serialization; define score normalization, length penalty, and tie-break policy
Dynamic beam policiesFixed, scheduled, and score-margin policies exist and are configurable over HTTPUse real quality metrics to set default policies, score margins, shrinking rules, and quality regression gates
Item-constrained generationItem trie, legal-token mask, constrained topK, catalog reload/rollback, and item metadata exist in runtime/HTTPTest with real large catalogs; add item-level correctness, illegal-token checks, constrained topK optimization, and serving semantics regressions
Memory admission and reclamationKV budget, dense pool metrics, high-water marks, leak checks, and cancel/timeout lifecycle existCombine page/offload support with finer reclamation; improve high-concurrency admission policy; connect memory estimates to serving decisions
Model and backend matrixMain path validates Qwen3-0.6B / H100 / BF16Expand Qwen-family sizes, dtypes, quantization, head configs, checkpoint compatibility, and backend fallback tables
Multi-GPU and scale-outCurrent main path is single-node / single-GPUDesign TP/PP, multi-replica scheduling, cross-GPU KV and beam ownership, load balancing, and deployment orchestration
Tests and documentationSmoke tests, offline/online benchmarks, Nsight breakdown, and memory estimator existConsolidate stable entrypoints; add one-command final offline/online runs; add CI/nightly correctness and performance regressions

11. Repository Structure

src/gr_inference/gr_models/      Qwen-family model integration
src/gr_inference/gr_kv/          ContextKV, BeamKV, BeamPath
src/gr_inference/gr_kernels/     kernel wrappers and backend selection
src/gr_inference/gr_runtime/     beam search runtime and logits processing
src/gr_inference/gr_serving/     continuous batching, memory pools, HTTP serving
tools/                           benchmarks, comparison, profiling utilities
scripts/                         reproducible benchmark and serving entrypoints
tests/                           runtime, serving, model, kernel selection tests

12. Quickstart

Default Environment and Model

Default Docker image: lmsysorg/sglang:dev-cu13

Default model: Qwen/Qwen3-1.7B

From a recsys-examples checkout:

cd examples/sid-gr-inference

Or clone this branch directly:

git clone --recurse-submodules -b merge_gr_inference_to_main git@github.com:cb521/recsys-examples.git
cd recsys-examples/examples/sid-gr-inference

Enter the container:

scripts/run_container.sh

Select a Model

Hugging Face model:

MODEL=Qwen/Qwen3-0.6B scripts/run_container.sh scripts/quickstart_offline.sh

Existing local model:

MODEL_ROOT=/path/to/models MODEL_DIR=/workspace/models/Qwen3-1.7B scripts/run_container.sh scripts/quickstart_offline.sh

Pinned Hugging Face revision:

MODEL=Qwen/Qwen3-1.7B MODEL_REVISION=main scripts/run_container.sh scripts/quickstart_offline.sh

Common Commands

Quick performance and accuracy check:

RUN_ACCURACY=1 scripts/run_container.sh scripts/quickstart_offline.sh

Full offline performance and accuracy matrix:

CONTEXT_LENS="1000 5000" \
BATCH_SIZES="1 2 4 8" \
REPEAT=3 \
RUN_ACCURACY=1 \
scripts/run_container.sh scripts/quickstart_offline.sh

Results are written to:

benchmark_artifacts/sglang_compare/offline_perf_YYYYmmdd_HHMMSS/summary.md
benchmark_artifacts/sglang_compare/offline_accuracy_YYYYmmdd_HHMMSS/summary.md

Other Reproduction Entrypoints

In a container where dependencies are already installed, run:

# Full offline performance comparison.
scripts/run_offline_perf_benchmark.sh

# Full offline correctness alignment.
scripts/run_offline_accuracy_benchmark.sh

Run a fair evaluation with radix on/off, performance, and correctness:

OUT_DIR=benchmark_artifacts/sglang_compare/fair_eval_correctness_quick \
CONTEXT_LENS="1000" \
BEAM_WIDTHS="256" \
BATCH_SIZES="1 4" \
PERF_REPEAT=1 \
CORRECTNESS_REPEAT=1 \
scripts/run_gr_sglang_fair_eval.sh

Run Nsight breakdown for a fixed case:

CONTEXT_LEN=5000 \
BEAM_WIDTH=256 \
REQUESTS=4 \
MAX_BATCH_SIZE=4 \
scripts/run_short_context_nsys_compare.sh