[AMD] [GLM5] Guard cuda_runtime.h for ROCm in fused_metadata_copy

[Raiden-Makoto]

Motivation

On ROCm/HIP the JIT-compiled fused metadata-copy kernel
(python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh) fails to
compile because it includes <cuda_runtime.h> unconditionally
(fatal error: 'cuda_runtime.h' file not found). The DSA backend then falls
back to a per-element copy loop whose cost scales with draft length. This
collapses EAGLE speculative-decoding throughput at draft depth >= 4
(--speculative-num-steps >= 4, --speculative-num-draft-tokens >= 5) on
MI300X/MI325X/MI355X: acceptance length keeps rising but per-step latency
craters ~100x.

GLM-5.1-MXFP4 (MI355X / gfx950, tp4), EAGLE depth-4 (4,1,5), broken path:

EAGLE config	accept len	output tok/s	median ITL
(3,1,4) depth 3	3.70	227.4	9.8 ms
(4,1,5) depth 4 — before fix	4.58	6.9	963 ms

Modifications

Guard the include with the standard USE_ROCM macro (matching
sgl_kernel/utils.cuh): use <hip/hip_runtime.h> on ROCm and keep
<cuda_runtime.h> on CUDA. The CUDA build is byte-for-byte unchanged.

#ifndef USE_ROCM
#include <cuda_runtime.h>
#else
#include <hip/hip_runtime.h>
#endif

With the kernel compiling on ROCm, depth >= 4 EAGLE no longer falls back to the
slow loop and runs at the expected per-step cost.

Speed Benchmarks

All numbers are e2e sglang.bench_serving, GLM-5.1-MXFP4, MI355X (gfx950), tp4.

Direct effect of the fix — same (4,1,5) config and workload (c=8,
2048-in/256-out), before vs after:

(4,1,5) @ c=8	accept len	output tok/s	median ITL
before fix (loop fallback)	4.58	6.9	963 ms
after fix	4.59	578.3	6.94 ms

→ ~84x output-throughput recovery, accept length unchanged.

Post-fix sweep (8192-in/1024-out), depth-4 (4,1,5) vs the depth-3 (3,1,4)
baseline, measured one server at a time:

conc	(3,1,4) TPOT / out tok/s / medITL	(4,1,5) TPOT / out tok/s / medITL
c=2	7.35 / 228 / 6.09	6.02 / 289 / 5.22
c=4	8.81 / 372 / 7.08	8.19 / 408 / 6.11
c=8	12.17 / 533 / 8.42	11.45 / 579 / 7.29
c=16	17.34 / 767 / 9.95	16.88 / 779 / 8.67
c=32	27.85 / 970 / 12.37	27.23 / 995 / 10.91
c=64	38.14 / 1058 / 14.29	38.71 / 1048 / 12.79
accept	3.56 - 3.86	4.67 - 4.71

Once depth-4 is unblocked it is the better config (higher accept length → lower
ITL at every concurrency). Both configs in the sweep were run with
--disable-custom-all-reduce (an unrelated gfx950 stability flag), applied
identically, so the comparison is fair.

Accuracy Tests

GLM-5.1 (MI355X), EAGLE (4,1,5): GSM8K 0.95, invalid 0 — lossless vs the
non-speculative baseline.

Server launched with depth-4 (4,1,5) EAGLE:

SGLANG_DSA_TRITON_PREFILL=1 sglang serve \
  --model-path amd/GLM-5.1-MXFP4 \
  --tp 4 \
  --trust-remote-code \
  --kv-cache-dtype fp8_e4m3 \
  --tool-call-parser glm47 \
  --reasoning-parser glm45 \
  --dsa-prefill-backend tilelang \
  --dsa-decode-backend tilelang \
  --chunked-prefill-size 131072 \
  --mem-fraction-static 0.85 \
  --watchdog-timeout 1200 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 4 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 5 \
  --disable-custom-all-reduce \
  --host 0.0.0.0 \
  --port 30000

Checklist

• Format with pre-commit (clang-format)
• CUDA build path unchanged

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!