Refactor eagle speculative decoding

[Ying1123]

No description provided.

[zhyncs]

@Ying1123 May you help resolve the conflicts? Thanks!

[xiezhq-hermann]

I think style-wise this is a bit confusing, if we define the allocator to be in charge of the agnostic memory operations and define another memory pool class for the underlying layouts, we should be using allocators consistently in scheduler and only use memory pool at lower level codes.

[mpjlu]

commit a574770: there is illegal bug when run DeepSeek:
409 File "/data/peng/sglang/python/sglang/srt/managers/scheduler.py", line 1218, in run_batch
410 ) = self.draft_worker.forward_batch_speculative_generation(batch)
411 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 189, in forward_batch_speculative_generation
412 spec_info, to_free_cache_loc = self.draft(batch)
413 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 244, in draft
414 assign_draft_cache_locs[(num_seqs,)](
415 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 345, in
416 return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
417 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 691, in run
418 kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata,
419 File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/driver.py", line 365, in call
420 self.launch(*args, **kwargs)
421 RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

[Ying1123]

commit a574770: there is illegal bug when run DeepSeek: 409 File "/data/peng/sglang/python/sglang/srt/managers/scheduler.py", line 1218, in run_batch 410 ) = self.draft_worker.forward_batch_speculative_generation(batch) 411 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 189, in forward_batch_speculative_generation 412 spec_info, to_free_cache_loc = self.draft(batch) 413 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 244, in draft 414 assign_draft_cache_locs[(num_seqs,)]( 415 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 345, in 416 return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs) 417 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 691, in run 418 kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, 419 File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/driver.py", line 365, in call 420 self.launch(*args, **kwargs) 421 RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

@mpjlu Thanks for reporting. Could you provide the reproducible command?
@ispobock @zhyncs Could you also help take a look?

[mpjlu]

commit a574770: there is illegal bug when run DeepSeek: 409 File "/data/peng/sglang/python/sglang/srt/managers/scheduler.py", line 1218, in run_batch 410 ) = self.draft_worker.forward_batch_speculative_generation(batch) 411 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 189, in forward_batch_speculative_generation 412 spec_info, to_free_cache_loc = self.draft(batch) 413 File "/data/peng/sglang/python/sglang/srt/speculative/eagle_worker.py", line 244, in draft 414 assign_draft_cache_locs[(num_seqs,)]( 415 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 345, in 416 return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs) 417 File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 691, in run 418 kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, 419 File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/driver.py", line 365, in call 420 self.launch(*args, **kwargs) 421 RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

@mpjlu Thanks for reporting. Could you provide the reproducible command? @ispobock @zhyncs Could you also help take a look?
The following command can reproduce:
python3 -m sglang.launch_server
--model-path $model_path
--tp $tp_size
--dist-init-addr 29.224.56.106:5000
--nnodes 2
--node-rank 0
--trust-remote-code
--mem-fraction-static 0.6
--max-running-requests 64
--speculative-draft-model-path $draft_path
--speculative-algorithm NEXTN
--speculative-num-steps 2 \
--speculative-eagle-topk 2 \
--speculative-num-draft-tokens 4 \
--disable-cuda-graph \

test.py

import openai
client = openai.Client(
base_url="http://127.0.0.1:30000/v1", api_key="EMPTY")
import time

Chat completion

start = time.time()
for i in range(100):
response = client.chat.completions.create(
model="default",
messages=[
{"role": "user", "content": "请以诚信为主题写一篇1000字作文？"},
], temperature=0.6, max_tokens=1000,
extra_body={"top_p": 0.6, "top_k": 50}
)
print(response)
print("dur=", time.time() - start)

[ispobock]

I verified e19e733 on 8*H200 for DeepSeek-V3 model with nextn enabled and it works fine.

@mpjlu I cannot reproduce your error in my environment. I am not sure if it's an error for multi-node setting since I tried the same args as your command but on one node.