[Bug] Deepseek-V2.5 capture cuda graph failed

[halexan]

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
• 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed.
• 5. Please use English, otherwise it will be closed.

Describe the bug

Bug report:

==========
== CUDA ==
==========

CUDA Version 12.4.1

Container image Copyright (c) 2016-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

This container image and its contents are governed by the NVIDIA Deep Learning Container License.
By pulling and using the container, you accept the terms and conditions of this license:
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license

A copy of this license is made available in this container at /NGC-DL-CONTAINER-LICENSE for your convenience.

[23:03:24] server_args=ServerArgs(model_path='/data/model-cache/deepseek-ai/DeepSeek-V2___5/', tokenizer_path='/data/model-cache/deepseek-ai/DeepSeek-V2___5/', tokenizer_mode='auto', skip_tokenizer_init=False, load_format='auto', dtype='auto', kv_cache_dtype='auto', trust_remote_code=True, context_length=131072, quantization=None, served_model_name='deepseek-chat', chat_template=None, is_embedding=False, host='0.0.0.0', port=50521, additional_ports=[50522, 50523, 50524, 50525], mem_fraction_static=0.75, max_running_requests=None, max_total_tokens=None, chunked_prefill_size=131072, max_prefill_tokens=131072, schedule_policy='lpm', schedule_conservativeness=0.1, tp_size=8, stream_interval=1, random_seed=1008230935, constrained_json_whitespace_pattern=None, log_level='info', log_level_http=None, log_requests=False, show_time_cost=False, api_key=None, file_storage_pth='SGLang_storage', dp_size=1, load_balance_method='round_robin', nccl_init_addr=None, nnodes=1, node_rank=None, json_model_override_args='{}', attention_backend='flashinfer', sampling_backend='flashinfer', disable_flashinfer=False, disable_flashinfer_sampling=False, disable_radix_cache=False, disable_regex_jump_forward=False, disable_cuda_graph=False, disable_cuda_graph_padding=False, disable_disk_cache=False, disable_custom_all_reduce=True, disable_mla=False, enable_mixed_chunk=False, enable_torch_compile=False, max_torch_compile_bs=32, torchao_config='', enable_p2p_check=False, triton_attention_reduce_in_fp32=False, lora_paths=None, max_loras_per_batch=8)
[23:03:24 TP1] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP2] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP3] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP4] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP5] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP6] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP0] MLA optimization is tunred on. Use triton backend.
[23:03:24 TP7] MLA optimization is tunred on. Use triton backend.
[23:03:25 TP2] Init nccl begin.
[23:03:25 TP0] Init nccl begin.
[23:03:25 TP5] Init nccl begin.
[23:03:25 TP4] Init nccl begin.
[23:03:25 TP3] Init nccl begin.
[23:03:25 TP7] Init nccl begin.
[23:03:25 TP1] Init nccl begin.
[23:03:25 TP6] Init nccl begin.
[rank4]:[W919 23:03:25.609730308 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank3]:[W919 23:03:25.609923773 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank2]:[W919 23:03:26.191450465 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank0]:[W919 23:03:26.200912535 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank5]:[W919 23:03:30.575735830 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank1]:[W919 23:03:30.617976304 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank6]:[W919 23:03:31.634186928 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[rank7]:[W919 23:03:31.636653957 ProcessGroupGloo.cpp:712] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())
[23:03:43 TP2] Load weight begin. avail mem=78.42 GB
[23:03:43 TP7] Load weight begin. avail mem=78.47 GB
[23:03:43 TP1] Load weight begin. avail mem=78.42 GB
[23:03:43 TP6] Load weight begin. avail mem=78.42 GB
[23:03:43 TP5] Load weight begin. avail mem=78.42 GB
[23:03:43 TP3] Load weight begin. avail mem=78.47 GB
[23:03:43 TP0] Load weight begin. avail mem=78.47 GB
[23:03:43 TP4] Load weight begin. avail mem=78.47 GB
Cache shape torch.Size([163840, 64])
Loading safetensors checkpoint shards:   0% Completed | 0/55 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:   2% Completed | 1/55 [00:05<05:17,  5.88s/it]
Loading safetensors checkpoint shards:   4% Completed | 2/55 [00:11<05:02,  5.71s/it]
Loading safetensors checkpoint shards:   5% Completed | 3/55 [00:18<05:31,  6.37s/it]
Loading safetensors checkpoint shards:   7% Completed | 4/55 [00:25<05:40,  6.67s/it]
Loading safetensors checkpoint shards:   9% Completed | 5/55 [00:35<06:20,  7.61s/it]
Loading safetensors checkpoint shards:  11% Completed | 6/55 [00:45<06:56,  8.50s/it]
Loading safetensors checkpoint shards:  13% Completed | 7/55 [00:55<07:09,  8.95s/it]
Loading safetensors checkpoint shards:  15% Completed | 8/55 [00:55<04:55,  6.28s/it]
Loading safetensors checkpoint shards:  16% Completed | 9/55 [00:56<03:26,  4.49s/it]
Loading safetensors checkpoint shards:  18% Completed | 10/55 [00:56<02:26,  3.27s/it]
Loading safetensors checkpoint shards:  20% Completed | 11/55 [00:57<01:47,  2.44s/it]
Loading safetensors checkpoint shards:  22% Completed | 12/55 [00:57<01:20,  1.86s/it]
Loading safetensors checkpoint shards:  24% Completed | 13/55 [00:58<01:01,  1.47s/it]
Loading safetensors checkpoint shards:  25% Completed | 14/55 [00:59<00:49,  1.22s/it]
Loading safetensors checkpoint shards:  27% Completed | 15/55 [00:59<00:41,  1.04s/it]
Loading safetensors checkpoint shards:  29% Completed | 16/55 [01:00<00:35,  1.09it/s]
Loading safetensors checkpoint shards:  31% Completed | 17/55 [01:00<00:31,  1.22it/s]
Loading safetensors checkpoint shards:  33% Completed | 18/55 [01:01<00:27,  1.33it/s]
Loading safetensors checkpoint shards:  35% Completed | 19/55 [01:02<00:25,  1.43it/s]
Loading safetensors checkpoint shards:  36% Completed | 20/55 [01:02<00:24,  1.46it/s]
Loading safetensors checkpoint shards:  38% Completed | 21/55 [01:03<00:22,  1.50it/s]
Loading safetensors checkpoint shards:  40% Completed | 22/55 [01:03<00:21,  1.53it/s]
Loading safetensors checkpoint shards:  42% Completed | 23/55 [01:04<00:19,  1.60it/s]
Loading safetensors checkpoint shards:  44% Completed | 24/55 [01:05<00:19,  1.63it/s]
Loading safetensors checkpoint shards:  45% Completed | 25/55 [01:05<00:18,  1.61it/s]
Loading safetensors checkpoint shards:  47% Completed | 26/55 [01:26<03:17,  6.79s/it]
Loading safetensors checkpoint shards:  49% Completed | 27/55 [01:49<05:24, 11.60s/it]
Loading safetensors checkpoint shards:  51% Completed | 28/55 [02:15<07:04, 15.71s/it]
Loading safetensors checkpoint shards:  53% Completed | 29/55 [02:40<08:02, 18.55s/it]
Loading safetensors checkpoint shards:  55% Completed | 30/55 [03:04<08:25, 20.24s/it]
Loading safetensors checkpoint shards:  56% Completed | 31/55 [03:30<08:44, 21.86s/it]
Loading safetensors checkpoint shards:  58% Completed | 32/55 [03:54<08:41, 22.67s/it]
Loading safetensors checkpoint shards:  60% Completed | 33/55 [04:16<08:14, 22.49s/it]
Loading safetensors checkpoint shards:  62% Completed | 34/55 [04:40<07:58, 22.81s/it]
Loading safetensors checkpoint shards:  64% Completed | 35/55 [05:04<07:44, 23.24s/it]
Loading safetensors checkpoint shards:  65% Completed | 36/55 [05:29<07:29, 23.65s/it]
Loading safetensors checkpoint shards:  67% Completed | 37/55 [05:57<07:29, 24.99s/it]
Loading safetensors checkpoint shards:  69% Completed | 38/55 [06:22<07:07, 25.14s/it]
Loading safetensors checkpoint shards:  71% Completed | 39/55 [06:47<06:39, 24.95s/it]
Loading safetensors checkpoint shards:  73% Completed | 40/55 [07:11<06:10, 24.73s/it]
Loading safetensors checkpoint shards:  75% Completed | 41/55 [07:35<05:45, 24.65s/it]
Loading safetensors checkpoint shards:  76% Completed | 42/55 [08:00<05:19, 24.58s/it]
Loading safetensors checkpoint shards:  78% Completed | 43/55 [08:24<04:53, 24.45s/it]
Loading safetensors checkpoint shards:  80% Completed | 44/55 [08:47<04:25, 24.11s/it]
Loading safetensors checkpoint shards:  82% Completed | 45/55 [09:11<03:58, 23.88s/it]
Loading safetensors checkpoint shards:  84% Completed | 46/55 [09:34<03:33, 23.72s/it]
Loading safetensors checkpoint shards:  85% Completed | 47/55 [09:59<03:13, 24.24s/it]
Loading safetensors checkpoint shards:  87% Completed | 48/55 [10:23<02:48, 24.11s/it]
Loading safetensors checkpoint shards:  89% Completed | 49/55 [10:47<02:24, 24.09s/it]
Loading safetensors checkpoint shards:  91% Completed | 50/55 [11:11<02:00, 24.06s/it]
Loading safetensors checkpoint shards:  93% Completed | 51/55 [11:36<01:36, 24.20s/it]
Loading safetensors checkpoint shards:  95% Completed | 52/55 [12:01<01:13, 24.36s/it]
Loading safetensors checkpoint shards:  96% Completed | 53/55 [12:25<00:48, 24.26s/it]
Loading safetensors checkpoint shards:  98% Completed | 54/55 [12:49<00:24, 24.17s/it]
Loading safetensors checkpoint shards: 100% Completed | 55/55 [13:15<00:00, 24.96s/it]
Loading safetensors checkpoint shards: 100% Completed | 55/55 [13:15<00:00, 14.47s/it]

[23:17:00 TP2] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=21.96 GB
[23:17:00 TP3] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=22.01 GB
[23:17:00 TP0] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=22.01 GB
[23:17:01 TP5] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=21.96 GB
[23:17:01 TP6] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=21.96 GB
[23:17:01 TP1] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=21.96 GB
[23:17:01 TP4] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=22.01 GB
[23:17:01 TP7] Load weight end. type=DeepseekV2ForCausalLM, dtype=torch.bfloat16, avail mem=22.01 GB
[23:17:01 TP6] Memory pool end. avail mem=18.61 GB
[23:17:01 TP5] Memory pool end. avail mem=18.61 GB
[23:17:01 TP1] Memory pool end. avail mem=18.61 GB
[23:17:01 TP4] Memory pool end. avail mem=18.66 GB
[23:17:01 TP3] Memory pool end. avail mem=18.66 GB
[23:17:01 TP7] Memory pool end. avail mem=18.66 GB
[23:17:01 TP0] Memory pool end. avail mem=18.66 GB
[23:17:01 TP2] Memory pool end. avail mem=18.61 GB
[23:17:01 TP3] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP1] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP6] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP4] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP5] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP2] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP7] Capture cuda graph begin. This can take up to several minutes.
[23:17:01 TP0] Capture cuda graph begin. This can take up to several minutes.
[23:17:08 TP4] Exception in run_tp_server:
Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 147, in __init__
    self.capture()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 176, in capture
    ) = self.capture_one_batch_size(bs, forward)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 217, in capture_one_batch_size
    run_once()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 211, in run_once
    return forward(input_ids, input_metadata.positions, input_metadata)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py", line 663, in forward
    hidden_states = self.model(input_ids, positions, input_metadata)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py", line 632, in forward
    hidden_states, residual = layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py", line 587, in forward
    hidden_states = self.mlp(hidden_states)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py", line 85, in forward
    gate_up, _ = self.gate_up_proj(x)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/linear.py", line 365, in forward
    output_parallel = self.quant_method.apply(self, input_, bias)
  File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/linear.py", line 131, in apply
    return F.linear(x, layer.weight, bias)
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16BF, lda, b, CUDA_R_16BF, ldb, &fbeta, c, CUDA_R_16BF, ldc, compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)`

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/managers/tp_worker.py", line 959, in run_tp_server
    model_server = ModelTpServer(
  File "/sgl-workspace/sglang/python/sglang/srt/managers/tp_worker.py", line 100, in __init__
    self.model_runner = ModelRunner(
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py", line 128, in __init__
    self.init_cuda_graphs()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py", line 468, in init_cuda_graphs
    self.cuda_graph_runner = CudaGraphRunner(self)
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 149, in __init__
    raise Exception(
Exception: Capture cuda graph failed: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16BF, lda, b, CUDA_R_16BF, ldb, &fbeta, c, CUDA_R_16BF, ldc, compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)`
Possible solutions:
1. disable cuda graph by --disable-cuda-graph
2. set --mem-fraction-static to a smaller value
3. disable torch compile by not using --enable-torch-compile
Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose

Reproduction

python3 -m sglang.launch_server \
    --model-path /data/model-cache/deepseek-ai/DeepSeek-V2___5/ \
    --served-model-name deepseek-chat \
    --tp 8 \
    --mem-fraction-static 0.75 \
    --schedule-conservativeness 0.1 \
    --disable-custom-all-reduce \
    --trust-remote-code \
    --host 0.0.0.0 \
    --port 50521

Environment

docker image lmsysorg/sglang:v0.3.1.post2-cu124