amd/deepseek_v4 integration 2/N - cuda graph 0426

[kkHuang-amd]

Motivation

Update amd/deepseek_v4 integration branch

Following PRs have large set of conflict, we use this PR and upstream amd/deepseek_v4 branch to integrate in parallel.
#23600
#23608

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request implements comprehensive support for DeepSeek-V4, introducing optimized CUDA kernels for compressed attention (C4/C128), expert-filtered activations, and a specialized memory pool architecture. It also updates diffusion pipelines to support request-local schedulers and adds extensive configuration for deployment. Review feedback identifies several critical issues in the attention backend, specifically regarding performance bottlenecks from host-device synchronizations, incorrect metadata copying during CUDA graph replay for variable batch sizes, and potential indexing errors in multi-token prediction (MTP). Additionally, improvements were suggested to enable expert filtering on non-CUDA platforms and to enhance thread safety by avoiding stateful caching within memory pools.

I am having trouble creating individual review comments. Click here to see my feedback.

python/sglang/srt/layers/attention/deepseek_v4_backend_radix.py (1132)

Calling .item() on a GPU tensor causes a host-device synchronization. This is a significant performance bottleneck in the hot path and will break CUDA graph capture if this method is called during the capture phase. Consider using torch.cat with a sliced and expanded tensor to avoid the sync: torch.cat([req_pool_indices_repeated, req_pool_indices_repeated[-1:].expand(pad_size)]).

python/sglang/srt/layers/attention/deepseek_v4_backend_radix.py (190-237)

The copy_ implementation for DSV4AttnMetadataRadix (and its constituent metadata classes) uses dst_val.copy_(src_val) which requires exact shape matching. During CUDA graph replay, the destination tensors have the captured max_bs size, while the source tensors from temp_metadata have the current bs size. This will cause a runtime error when bs < max_bs. You should use slicing to copy only the active portion: dst_val[:src_val.shape[0]].copy_(src_val).

python/sglang/srt/layers/attention/deepseek_v4_backend_radix.py (1044)

The use of _pad_tensor_to_size (which internally uses torch.cat) is problematic for MTP if topk > 1. If q has bs * topk tokens but swa_page_indices only has bs rows, padding with 0 will result in incorrect indices for all but the first token of each request. It should likely be repeat_interleave(topk, dim=0) instead. Furthermore, torch.cat should be avoided in the forward path to maintain CUDA graph compatibility and performance.

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py (532)

The expert filtering logic is restricted to _is_cuda. While the comment mentions that HIP/XPU fall through to the unfiltered path where the down kernel handles zeros, this results in redundant computation in the activation kernel for filtered tokens on those platforms. If the silu_and_mul JIT kernel supports filtering on ROCm, it should be enabled here as well to improve efficiency.

python/sglang/srt/mem_cache/deepseekv4_memory_pool.py (830)

Caching swa_loc in self.cached_loc assumes that the model forward pass is strictly sequential across layers and that only one forward pass happens at a time. While this is currently true for the SGLang model runner, it makes the code non-thread-safe and could lead to subtle bugs if the execution model changes (e.g., multi-threaded speculative decoding steps). Consider passing this metadata through the forward_batch or metadata objects instead of storing it as state in the pool.