update sglang && support qwen3 next (#355)

* support qwen3 next

* fix bug

* fix bug

* update sglang

Author: sleepcoo

configs/qwen3-next-80b-a3b-eagle3.json

@@ -2,13 +2,6 @@
     "architectures": [
       "LlamaForCausalLMEagle3"
     ],
-    "eagle_config": {
-        "eagle_aux_hidden_state_layer_ids": [
-            1,
-            23,
-            45
-        ]
-    },
     "attention_dropout": 0.0,
     "bos_token_id": 151643,
     "decoder_sparse_step": 1,

examples/run_qwen3_next_80b_eagle3_online.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname $SCRIPT_DIR)
+export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
+
+NUM_GPUS=${1:-8}
+TP_SIZE=4
+BUILD_DATASET_NUM_PROC=${BUILD_DATASET_NUM_PROC:-64}
+
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
+    $ROOT_DIR/scripts/train_eagle3.py \
+    --target-model-path $ROOT_DIR//Qwen/Qwen3-Next-80B-A3B-Instruct-FP8/\
+    --draft-model-config $ROOT_DIR/configs/qwen3-next-80b-a3b-eagle3.json \
+    --train-data-path $ROOT_DIR/data_qwen80b/qwen3_80b_perfectblend_train_regen.jsonl \
+    --output-dir $ROOT_DIR/qwen3-80b-regen-blend \
+    --num-epochs 2 \
+    --batch-size 2 \
+    --learning-rate 1e-4 \
+    --max-length 4096 \
+    --chat-template qwen \
+    --cache-dir $ROOT_DIR/cache \
+    --embedding-key model.embed_tokens.weight \
+    --tp-size $TP_SIZE \
+    --sglang-mem-fraction-static 0.5 \
+    --build-dataset-num-proc $BUILD_DATASET_NUM_PROC \
+    --target-model-backend sglang

requirements.txt

@@ -1,7 +1,7 @@
 pre-commit
-torch==2.8.0
-torchaudio==2.8.0
-torchvision==0.23.0
+torch==2.9.1
+torchaudio==2.9.1
+torchvision==0.24.1
 transformers==4.57.1
 qwen-vl-utils==0.0.11
 datasets
@@ -12,5 +12,5 @@ psutil
 numpy
 accelerate
 pydantic
-sglang[all]==0.5.5
+sglang[all]==0.5.6
 openai-harmony

specforge/modeling/target/eagle3_target_model.py

@@ -8,6 +8,7 @@
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch
 from sglang.srt.sampling.sampling_params import SamplingParams
@@ -306,11 +307,13 @@ def _extend(
                 module.return_last_hidden_states = return_last_hidden_states
                 module.return_logits = return_logits
 
-        tree_cache = RadixCache(
-            None,
+        cache_params = CacheInitParams(
+            disable=False,
+            req_to_token_pool=self.model_runner.req_to_token_pool,
             token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
             page_size=self.model_runner.server_args.page_size,
         )
+        tree_cache = RadixCache(cache_params)
 
         batch = ScheduleBatch.init_new(
             reqs=reqs,