Fix DSA compression tail capacity for PD decode request slots

[HanHan009527]

Motivation

GLM-5.3-Flash's NextN draft shares the target's hybrid request pool. PD decode adds preallocation slots beyond max_running_requests, but the plain DSA builder uses that running limit to size compression tails indexed by req_pool_idx.

With 32 running requests and 64 extra slots, slot 96 is valid but the tail has only 33 rows:

index 96 is out of bounds for dimension 0 with size 33

The same model's target hybrid branch already sizes tails from the actual request table. The plain draft branch missed this distinction. Graph replay also preserves these request IDs.

Modifications

• At the plain DSA builder call site, pass req_to_token_pool.req_to_token.shape[0] as max_running_requests, matching the target hybrid branch. Keep the existing parameter name, function signature and constructor +1; the scheduler's running limit stays unchanged.

Accuracy Tests

The results below were collected at 7f66f7e5af with shape[0] - 1. This revision aligns with the target hybrid branch by using shape[0], allocating one additional row; runtime tests were not rerun. No new UT is included.

Check	Before	After
CPU/CUDA capacity check (prior revision)	Old capacity: assert torch.Size([33, 6, 128]) == (97, 6, 128); 4 failures	16/16 pass
Real PD requests	Earlier eager run: torch._dynamo.exc.InternalTorchDynamoError: AcceleratorError: CUDA error: an illegal memory access was encountered	130/130 pass

The regression uses a controlled restoration of the old capacity. Serving used TileLang/BF16, TP8/EP8, EAGLE 1/1/2, Prefill eager and Decode full Graph. Prefill draft Graph still has the separate width issue in #37573. No model-quality benchmark was run.

Deployment commands and full earlier error

Deploy and use

Use an SGLang installation containing this PR, with GLM-5.3-Flash and a working Mooncake/RDMA setup. Start two workers with 8 GPUs each, followed by the router. Set the model path, reachable worker/router addresses and RDMA device list for your deployment. The router needs access to the model configuration as well. Allow the Prefill HTTP/bootstrap ports (32171/28998), Decode HTTP port (32172), and router port (30000) between the relevant components.

Prefill

export MODEL_PATH=/path/to/GLM-5.3-Flash
export IB_DEVICES="<RDMA device list>"
export SGLANG_HOST_IP="<prefill-address>"
export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0

python3 -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --served-model-name GLM-5.3-Flash \
  --chat-template "$MODEL_PATH/chat_template.jinja" \
  --tp-size 8 \
  --ep-size 8 \
  --pp-size 1 \
  --dp-size 1 \
  --dcp-size 1 \
  --mem-fraction-static 0.70 \
  --context-length 69632 \
  --dsa-prefill-backend tilelang \
  --dsa-decode-backend tilelang \
  --kv-cache-dtype bf16 \
  --cuda-graph-backend-prefill disabled \
  --cuda-graph-backend-decode disabled \
  --cuda-graph-max-bs-decode 32 \
  --disable-shared-experts-fusion \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --skip-server-warmup \
  --watchdog-timeout 600 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 1 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 2 \
  --max-running-requests 32 \
  --chunked-prefill-size 8192 \
  --max-prefill-tokens 8192 \
  --moe-runner-backend triton \
  --speculative-moe-runner-backend triton \
  --disaggregation-mode prefill \
  --disaggregation-transfer-backend mooncake \
  --disaggregation-ib-device "$IB_DEVICES" \
  --host 0.0.0.0 \
  --disable-overlap-schedule \
  --disaggregation-bootstrap-port 28998 \
  --port 32171 \
  --nccl-port 22171

Decode

export MODEL_PATH=/path/to/GLM-5.3-Flash
export IB_DEVICES="<RDMA device list>"
export SGLANG_HOST_IP="<decode-address>"
export SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=0

python3 -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --served-model-name GLM-5.3-Flash \
  --chat-template "$MODEL_PATH/chat_template.jinja" \
  --tp-size 8 \
  --ep-size 8 \
  --pp-size 1 \
  --dp-size 1 \
  --dcp-size 1 \
  --mem-fraction-static 0.70 \
  --context-length 69632 \
  --dsa-prefill-backend tilelang \
  --dsa-decode-backend tilelang \
  --kv-cache-dtype bf16 \
  --cuda-graph-backend-prefill disabled \
  --cuda-graph-backend-decode full \
  --cuda-graph-max-bs-decode 32 \
  --disable-shared-experts-fusion \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --skip-server-warmup \
  --watchdog-timeout 600 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 1 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 2 \
  --max-running-requests 32 \
  --chunked-prefill-size 8192 \
  --max-prefill-tokens 8192 \
  --moe-runner-backend triton \
  --speculative-moe-runner-backend triton \
  --disaggregation-mode decode \
  --disaggregation-transfer-backend mooncake \
  --disaggregation-ib-device "$IB_DEVICES" \
  --host 0.0.0.0 \
  --disaggregation-decode-extra-slots 64 \
  --port 32172 \
  --nccl-port 22172

Router

export MODEL_PATH=/path/to/GLM-5.3-Flash
export PREFILL_HOST="<prefill-address>"
export DECODE_HOST="<decode-address>"

python3 -m sglang_router.launch_router \
  --pd-disaggregation \
  --prefill-policy round_robin \
  --decode-policy round_robin \
  --prefill "http://$PREFILL_HOST:32171" 28998 \
  --decode "http://$DECODE_HOST:32172" \
  --host 0.0.0.0 \
  --port 30000 \
  --request-timeout-secs 1800 \
  --model-path "$MODEL_PATH"

Send a request

Once both workers and the router are ready, send a request to the router's OpenAI-compatible endpoint:

export ROUTER_HOST="<router-address>"

curl --fail-with-body --max-time 600 \
  "http://$ROUTER_HOST:30000/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "GLM-5.3-Flash",
    "messages": [{"role": "user", "content": "List three things to check when an HTTP service returns 503."}],
    "temperature": 0,
    "max_tokens": 1024,
    "stream": false
  }'

Full error before the fix

One complete TP0/EP0 traceback from the earlier dcebe8c PD eager baseline. Checkout/install path prefixes and the timestamp are normalized; all stack frames and the final error are retained.

[TP0 EP0] Scheduler hit an exception: Traceback (most recent call last):
  File "python/sglang/srt/managers/scheduler.py", line 5770, in run_scheduler_process
    scheduler.run_event_loop()
  File "python/sglang/srt/managers/scheduler.py", line 1890, in run_event_loop
    dispatch_event_loop(self)
  File "python/sglang/srt/managers/scheduler.py", line 5630, in dispatch_event_loop
    scheduler.event_loop_overlap_disagg_decode()
  File "<site-packages>/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/disaggregation/decode.py", line 2562, in event_loop_overlap_disagg_decode
    batch_result = self.run_batch(batch)
                   ^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/observability/scheduler_stage_metrics.py", line 155, in wrapper
    return profiled_func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/utils/nvtx_utils.py", line 109, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/managers/scheduler.py", line 4241, in run_batch
    batch_result = self.model_worker.forward_batch_generation(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/speculative/eagle_worker_v2.py", line 1292, in forward_batch_generation
    self.draft_worker._draft_extend_for_decode(batch, batch_output)
  File "python/sglang/srt/speculative/eagle_worker_v2.py", line 1006, in _draft_extend_for_decode
    draft_logits_output = self.draft_runner.forward(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/model_executor/model_runner.py", line 1654, in forward
    output = self._forward_raw(
             ^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/model_executor/model_runner.py", line 1841, in _forward_raw
    ret = self.eager_runner.execute(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/model_executor/runner/eager_runner.py", line 228, in execute
    return self._execute_extend(forward_batch, pp_proxy_tensors)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/model_executor/runner/eager_runner.py", line 372, in _execute_extend
    ret = model_runner.model.forward(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_nextn.py", line 402, in forward
    hidden_states = self.model(input_ids, positions, forward_batch)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_nextn.py", line 273, in forward
    hidden_states, residual, topk_indices = self.decoder(
                                            ^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_v2.py", line 2479, in forward
    hidden_states = self.self_attn(
                    ^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_v2.py", line 2065, in forward
    s = self.forward_prepare(
        ^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_v2.py", line 2123, in forward_prepare
    inner_state = self.forward_absorb_prepare(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py", line 448, in forward_absorb_prepare
    topk_indices = self.indexer(
                   ^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/kernels/fused_op.py", line 659, in forward
    result = method(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py", line 1395, in forward_cuda
    return self._forward_cuda_target_verify(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py", line 1345, in _forward_cuda_target_verify
    weights = self._get_logits_head_gate(x, q_scale)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/eval_frame.py", line 1166, in compile_wrapper
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 2628, in __call__
    result = self._torchdynamo_orig_backend(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 2319, in __call__
    result = self._inner_convert(
             ^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 777, in __call__
    result = _compile(
             ^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 2168, in _compile
    raise InternalTorchDynamoError(
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 2102, in _compile
    guarded_code, tracer_output = compile_inner(code, one_graph, hooks)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 1687, in compile_inner
    result = _compile_inner(code, one_graph, hooks)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 1747, in _compile_inner
    dynamo_output = compile_frame(
                    ^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 1584, in compile_frame
    bytecode, tracer_output = transform_code_object(code, transform)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/bytecode_transformation.py", line 1821, in transform_code_object
    tracer_output = transformations(instructions, code_options)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 1555, in transform
    tracer_output = trace_frame(
                    ^^^^^^^^^^^^
  File "<site-packages>/torch/_dynamo/convert_frame.py", line 386, in _fn
    torch.cuda.set_rng_state(cuda_rng_state)
  File "<site-packages>/torch/cuda/random.py", line 76, in set_rng_state
    _lazy_call(cb)
  File "<site-packages>/torch/cuda/__init__.py", line 460, in _lazy_call
    callable()
  File "<site-packages>/torch/cuda/random.py", line 74, in cb
    default_generator.set_state(new_state)
torch._dynamo.exc.InternalTorchDynamoError: AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1

Speed Tests and Profiling

Not run; this is an allocation-correctness fix, with no performance claim.

Checklist

• Pre-commit checks passed.
• New unit tests: not included.
• SGLang code style followed.
• Documentation: no public interface change; usage commands are included above.
• Model-quality and speed benchmarks: not run.

Review and Merge Process

Awaiting review and CI; the PR remains a draft.

---

CI States

Latest PR Test (Base): ⏳ Run #34206300050
Latest PR Test (Extra): ❌ Run #34206299712
Latest PR Test (AMD ROCm 7.2): ⏳ Run #34206300540

[HanHan009527]

@JustinTong0323 Please review this when you have time.

[HanHan009527]

/tag-and-rerun-ci