fix(dsa): guard cuda_runtime.h include so fused_metadata_copy builds on ROCm

[andyluo7]

Summary

The DSA fused metadata-copy JIT kernel (python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh) includes <cuda_runtime.h> unconditionally. The tvm-ffi JIT path compiles the .cuh directly with hipcc -x hip without a hipify pass, so on ROCm there is no cuda_runtime.h on the include path and compilation fails:

fused_metadata_copy.cuh:38:10: fatal error: 'cuda_runtime.h' file not found
   38 | #include <cuda_runtime.h>

This kernel is reached on AMD whenever speculative_num_steps > 3, which selects the multi-backend fused-copy path in dsa_backend.py. The single-backend path is already HIP-guarded (_USE_FUSED_METADATA_COPY = ... and not _is_hip), but the multi path is not. And because cache_once (jit_kernel/utils.py) only memoizes successful compiles, the failed build is retried on every CUDA-graph replay — the call site catches the error and falls back to the per-backend loop, but pays a full failed hipcc invocation per decode step, collapsing throughput to ~0.13 tok/s on gfx950.

The fix guards just the include with #ifndef USE_ROCM. The rest of the file is already HIP-compatible: __grid_constant__ is shimmed in include/sgl_kernel/utils.cuh, and host::LaunchKernel has a hipLaunchKernelGGL branch.

Test plan

Verified on MI350X (gfx950) / ROCm 7.2.4:

• Before: _jit_fused_metadata_copy_multi_module(...) fails with the cuda_runtime.h fatal error above.
• After: the multi kernel compiles cleanly (no further HIP errors) for has_real_page_table ∈ {false, true} and has_flashmla ∈ {false, true}.
• Functional test: launched the compiled kernel on synthetic int32 tensors and compared against a torch reference — cache_seqlens, cu_seqlens_k[1:], page table (col < max_len), dsa_cache_seqlens, dsa_cu_seqlens_k[1:], and real_page_table all copy exactly to all three destination backends; index-0 and the page-table tail are correctly left untouched. PASS for both has_real_page_table false and true.

Note: this restores num_steps > 3 throughput on ROCm; it is independent of the gfx950 block-FP8 accuracy issue tracked in #28685.

---

CI States

Latest PR Test (Base): ❌ Run #27849498701
Latest PR Test (Extra): ❌ Run #27849498600

[gemini-code-assist]

Code Review

This pull request introduces a conditional preprocessor check in fused_metadata_copy.cuh to exclude the <cuda_runtime.h> header when compiling for ROCm (USE_ROCM is defined). This improves compatibility with AMD GPU environments. There are no review comments, and I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[andyluo7]

Could a maintainer please add the run-ci label and take a look? The change is a one-line #ifndef USE_ROCM include guard (lint passed; gemini-code-assist had no comments).

Why this matters (end-to-end validation on MI350X / gfx950, ROCm 7.2.4)

Without this guard, any EAGLE/MTP run with --speculative-num-steps > 3 selects the multi-backend fused metadata-copy path, whose JIT kernel fails to compile on ROCm:

fused_metadata_copy.cuh:38:10: fatal error: 'cuda_runtime.h' file not found

Because cache_once memoizes only successful compiles, the build is retried on every CUDA-graph replay, so the DSA spec-decode loop falls back step-by-step and decode collapses to ~0.13 tok/s.

With this PR (and, separately, the gfx950 block-FP8 accuracy fix tracked in #28685), a full --speculative-num-steps 5 GLM-5.2-FP8 serve works end-to-end (TP4, SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA=1, cuda graph on):

Check	Result
Draft cuda-graph capture	clean (53.7s); zero "falling back to loop" in the entire log
Accept length	2.92 – 3.92
Decode throughput	130 – 175 tok/s (vs ~0.13 tok/s before)

This PR only addresses the throughput cliff (the compile failure); it is independent of the accuracy issue in #28685. Verified the kernel both compiles and produces correct copies on gfx950 (functional test vs a torch reference, all metadata fields, both has_real_page_table and has_flashmla variants).

[Fridge003]

@andyluo7 Please fix conflict, thanks

[andyluo7]

Superseded by #29373 (merged 2026-06-27 by @HaiShaw), which landed the same #ifndef USE_ROCM guard around cuda_runtime.h in fused_metadata_copy.cuh plus the #else #include <hip/hip_runtime.h> branch — a strict superset of this PR. The guard is now on main verbatim, so this change is a no-op against main. Closing as already-fixed; thanks all.