Fix Engine error when enabling DP attention

[fzyzcjy]

Motivation

Before fix

➜  misc CUDA_VISIBLE_DEVICES=4,5,6,7 python3 adhoc_offline_generate.py '{"model_path": "deepseek-ai/DeepSeek-V2-Lite", "tp_size": 2, "dp_size": 2, "trust_remote_code": true, "disable_cuda_graph": true, "port": 13579, "enable_dp_attention": true}'                           
options={'model_path': 'deepseek-ai/DeepSeek-V2-Lite', 'tp_size': 2, 'dp_size': 2, 'trust_remote_code': True, 'disable_cuda_graph': True, 'port': 13579, 'enable_dp_attention': True}
INFO 03-21 07:09:04 __init__.py:190] Automatically detected platform cuda.
The following error message 'operation scheduled before its operands' can be ignored.
DP attention is enabled. The chunked prefill size is adjusted to 4096 to avoid MoE kernel issues. 
INFO 03-21 07:09:15 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:09:16 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:09:23 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:09:23 __init__.py:190] Automatically detected platform cuda.
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:01<00:03,  1.13s/it]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:02<00:02,  1.23s/it]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:03<00:01,  1.09s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:04<00:00,  1.13s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:04<00:00,  1.14s/it]

Traceback (most recent call last):
  File "/host_home/primary_synced/tom_sglang_server/misc/adhoc_offline_generate.py", line 10, in <module>
    engine = sgl.Engine(**options)
  File "/host_home/primary_synced/sglang/python/sglang/api.py", line 45, in Engine
    return Engine(*args, **kwargs)
  File "/host_home/primary_synced/sglang/python/sglang/srt/entrypoints/engine.py", line 127, in __init__
    self.send_to_rpc = get_zmq_socket(
  File "/host_home/primary_synced/sglang/python/sglang/srt/utils.py", line 883, in get_zmq_socket
    socket.bind(endpoint)
  File "/usr/local/lib/python3.10/dist-packages/zmq/sugar/socket.py", line 317, in bind
    super().bind(addr)
  File "_zmq.py", line 917, in zmq.backend.cython._zmq.Socket.bind
  File "_zmq.py", line 179, in zmq.backend.cython._zmq._check_rc

After fix

➜  misc CUDA_VISIBLE_DEVICES=4,5,6,7 python3 adhoc_offline_generate.py '{"model_path": "deepseek-ai/DeepSeek-V2-Lite", "tp_size": 2, "dp_size": 2, "trust_remote_code": true, "disable_cuda_graph": true, "port": 13579, "enable_dp_attention": true}'
options={'model_path': 'deepseek-ai/DeepSeek-V2-Lite', 'tp_size': 2, 'dp_size': 2, 'trust_remote_code': True, 'disable_cuda_graph': True, 'port': 13579, 'enable_dp_attention': True}
INFO 03-21 07:13:08 __init__.py:190] Automatically detected platform cuda.
The following error message 'operation scheduled before its operands' can be ignored.
DP attention is enabled. The chunked prefill size is adjusted to 4096 to avoid MoE kernel issues. 
INFO 03-21 07:13:19 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:13:20 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:13:28 __init__.py:190] Automatically detected platform cuda.
INFO 03-21 07:13:28 __init__.py:190] Automatically detected platform cuda.
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:01<00:03,  1.12s/it]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:02<00:02,  1.25s/it]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:03<00:01,  1.11s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:04<00:00,  1.14s/it]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:04<00:00,  1.15s/it]

prompts=['The capital of France is', 'The capital of the United Kindom is', 'Today is a sunny day and I like', '1+1==', '9-1=8, 8-1==', '1*1=1, 1*2=3, 1*3=', '8/1=8, 7/1=7, 6/1=6, 5/1='] sampling_params={'temperature': 0, 'max_new_tokens': 2}
loc("/host_home/primary_synced/sglang/python/sglang/srt/layers/attention/triton_ops/decode_attention.py":327:16): error: operation scheduled before its operands
loc("/host_home/primary_synced/sglang/python/sglang/srt/layers/attention/triton_ops/decode_attention.py":327:16): error: operation scheduled before its operands
out=[{'text': ' Paris,', 'meta_info': {'id': 'd8f910a1f6374b2583ac201b5feec71d', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 6, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.91384482383728}}, {'text': ' a city', 'meta_info': {'id': 'fd2bfd9b2bfb4ede91e9ab7050d5c63e', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 9, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.914016723632812}}, {'text': ' it.', 'meta_info': {'id': '178c2b9319cb4b67bb40d9f0c975920d', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 9, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.913856029510498}}, {'text': '2\n', 'meta_info': {'id': '93569cfe03d7483695c55ce970b695db', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 5, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.914022207260132}}, {'text': '7,', 'meta_info': {'id': 'da65eab1711e4b508f40f367e34655e6', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 12, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.913860082626343}}, {'text': '6,', 'meta_info': {'id': '8009271a07044d89881fa99603282793', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 19, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.91402578353882}}, {'text': '5,', 'meta_info': {'id': '7cf04c57f29e4b9db64bf8f842876c1e', 'finish_reason': {'type': 'length', 'length': 2}, 'prompt_tokens': 26, 'completion_tokens': 2, 'cached_tokens': 0, 'e2e_latency': 19.913864850997925}}]
prompt='The capital of France is' out_item["text"]=' Paris,'
prompt='The capital of the United Kindom is' out_item["text"]=' a city'
prompt='Today is a sunny day and I like' out_item["text"]=' it.'
prompt='1+1==' out_item["text"]='2\n'
prompt='9-1=8, 8-1==' out_item["text"]='7,'
prompt='1*1=1, 1*2=3, 1*3=' out_item["text"]='6,'
prompt='8/1=8, 7/1=7, 6/1=6, 5/1=' out_item["text"]='5,'

Test script

import json
import sys

import sglang as sgl

if __name__ == '__main__':
    options = json.loads(sys.argv[1])
    print(f'{options=}')

    engine = sgl.Engine(**options)

    prompts = [
        "The capital of France is",
        "The capital of the United Kindom is",
        "Today is a sunny day and I like",
        "1+1==",
        "9-1=8, 8-1==",
        "1*1=1, 1*2=3, 1*3=",
        "8/1=8, 7/1=7, 6/1=6, 5/1=",
    ]
    sampling_params = dict(
        temperature=0,
        max_new_tokens=2,
    )
    print(f'{prompts=} {sampling_params=}')

    out = engine.generate(prompt=prompts, sampling_params=sampling_params)
    print(f'{out=}')

    for prompt, out_item in zip(prompts, out):
        print(f'{prompt=} {out_item["text"]=}')

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.