跳转至内容
  • 版块
  • 最新
  • 标签
  • 热门
  • 用户
  • 群组
皮肤
  • 浅色
  • Brite
  • Cerulean
  • Cosmo
  • Flatly
  • Journal
  • Litera
  • Lumen
  • Lux
  • Materia
  • Minty
  • Morph
  • Pulse
  • Sandstone
  • Simplex
  • Sketchy
  • Spacelab
  • United
  • Yeti
  • Zephyr
  • 深色
  • Cyborg
  • Darkly
  • Quartz
  • Slate
  • Solar
  • Superhero
  • Vapor

  • 默认(LCZ-Blue)
  • 不使用皮肤
  • LCZ-Green
  • LCZ-Blue
  • LCZ-Black
折叠
品牌标识

抡锤者

首页 版块 标签 硬件 AI 广场
  1. 主页
  2. 版块
  3. LLM讨论区
  4. # SGLang 上 MTP 投机的完整调优记录:从"draft 白占 4.7GB"到单路 89 t/s,以及一个会崩服务的 mamba 坑

# SGLang 上 MTP 投机的完整调优记录:从"draft 白占 4.7GB"到单路 89 t/s,以及一个会崩服务的 mamba 坑

已定时 已固定 已锁定 已移动 LLM讨论区
sg-langmtpqwen-27b
4 帖子 2 发布者 134 浏览
  • 从旧到新
  • 从新到旧
  • 最多赞同
回复
  • 在新帖中回复
登录后回复
此主题已被删除。只有拥有主题管理权限的用户可以查看。
  • E 离线
    E 离线
    Enigma
    德高望重
    编写于 最后由 编辑
    #1

    上一篇写了 llama.cpp + Q6_K + MTP 的成绩(工具 81.8 t/s / 代码 80.9 / 创作 52.9)。这篇是 SGLang 侧的续集:换成带 MTP 权重的 INT4 模型、把 NEXTN 投机调起来,最后不仅追平还略微超过了 llama.cpp(代码 89.3 t/s)。

    但过程比想象中曲折:MTP 打开后 KV 池从 28 万 token 掉到 1.4 万 token,连 32K 上下文都放不下。这篇把每个坑的定位过程和修法都写下来。

    1. 先给结论

    项目 结果
    权重 RedHatAI/Qwen3.8-27B-INT4(17.7 GB,自带 model_mtp.safetensors)
    引擎 SGLang 0.5.17 + --speculative-algorithm NEXTN
    代码生成 decode 41.5 t/s → 89.3 t/s(+115%,接受率 0.94)
    抽取原文 decode 40.5 t/s → 86.6 t/s(+114%,接受率 0.93)
    中文创作 decode 41.5 t/s → 59.7 t/s(+44%,接受率 0.51)
    双流聚合 82.9 t/s → 170.5 t/s
    prefill(13.7K 提示) 1846 → 1775 tok/s(−4%,基本无损)
    128K 上下文下的解码 34.4 → 37.8 t/s(接受长度掉到 2.0,增益变小但仍为正)
    代价 无(见第 7 节:原本以为要牺牲上下文,最后发现是 SGLang 的一个 bug,修完池子反而 ×4.7)

    一句话:MTP 能让解码翻倍(短上下文)到 +10%(128K 长上下文),而且第 7 节那个 bug 修掉之后,连单路 128K 都能带着 MTP 一起跑。

    2. 换权重:先验证包是不是自洽的

    上一篇(llama.cpp + Q6_K 那篇)我在 hotdogs 的 AWQ 包上栽过:它的 config.json 用 model.language_model. 前缀声明"忽略 linear_attn",但文件里真实键名是 model.layers.,而且 linear_attn 其实被量化了 —— 结果加载器按"未量化"去找 bf16 权重,找不到,输出满屏乱码。

    所以这次拿到 RedHatAI 的包先做了自洽性检查:

    # 1) ignore 列表里的名字和文件里的真实键名对得上吗
    # 2) 声明 bf16 的层,文件里真的是 bf16 吗
    model.language_model.layers.0.linear_attn.in_proj_qkv.weight_packed  I32   ← 量化,符合预期
    model.language_model.layers.0.linear_attn.in_proj_a.weight            BF16  ← 真 bf16,符合预期
    linear_attn.in_proj_a.weight (bf16) 层数: 48    ← 48 层全对
    linear_attn.in_proj_a.weight_packed 层数: 0     ← 没有残留的 packed 版本
    

    这个包是干净的:前缀和架构(Qwen3_5ForConditionalGeneration)一致,in_proj_a/b(输出维度只有 48,过不了 Marlin 的 64 整除要求)确实保留 bf16,re:^mtp.* 让 MTP 层也保持 bf16。直接就能跑,不用改权重。

    3. 下载坑:HF Xet 存储走不了镜像

    用 hf-mirror 下载时直接报错:

    RuntimeError: Task error: File reconstruction error: CAS Client Error:
    HTTP status client error (401 Unauthorized),
    domain: https://cas-server.xethub.hf.co/v2/reconstructions/...
    

    原因:这个仓库用的是 HF 的 Xet 存储,客户端会绕过 HF_ENDPOINT 直接去 cas-server.xethub.hf.co 取数据,所以镜像是失效的、而且直连会 401。

    解决:关掉 Xet,走经典 HTTP 路径(这样才会真正走 hf-mirror):

    export HF_ENDPOINT=https://hf-mirror.com
    export HF_HUB_DISABLE_XET=1      # ← 关键
    hf download RedHatAI/Qwen3.8-27B-INT4 --local-dir /mnt/sda6/download/qwen38-redhat-int4
    

    关掉之后速度稳定在 32 MB/s,19.5 GB 十分钟左右下完。

    4. 打开 NEXTN 投机

    --speculative-algorithm NEXTN \
    --speculative-draft-model-path <模型目录> \
    --speculative-eagle-topk 1 \
    --speculative-num-steps 3 \
    --speculative-num-draft-tokens 4
    

    --speculative-draft-model-path 指模型自己的目录就行:model.safetensors.index.json 里同时引用了 model.safetensors 和 model_mtp.safetensors,draft 加载器会把 mtp.* 的键挑出来用。

    效果(同一个模型,只差投机开关):

    负载 MTP 关 MTP 开 提升 接受率 接受长度
    代码生成 41.5 t/s 89.3 t/s +115% 0.941 3.76
    抽取原文 40.5 t/s 86.6 t/s +114% 0.931 3.66
    中文创作 41.5 t/s 59.7 t/s +44% 0.513 2.56
    双流聚合 82.9 t/s 170.5 t/s +109% — —

    对比之前在没有 MTP 权重时用 NGRAM 投机的成绩(接受率只有 0.06-0.23):训练过的 draft head 和 n-gram 匹配完全不是一个量级,接受率差了 4-15 倍。

    5. 最大的坑:MTP draft 白占 4.7 GB,KV 池被压到 1.4 万

    打开 MTP 后服务起来了,但一查 KV 池:

    Load weight end. type=Qwen3_5ForConditionalGeneration, mem usage=17.67 GB   ← 目标模型
    Load weight end. type=Qwen3_5ForCausalLMMTP,          mem usage=5.53 GB    ← MTP draft?!
    Mamba Cache is allocated. ssm_state size: 0.98GB, intermediate_ssm_state_cache size: 1.12GB
    KV Cache is allocated. #tokens: 14161        ← 只有 0.44 GB!
    Memory pool end. avail mem=4.17 GB           ← 还剩 4GB 没被用
    

    draft 的权重文件只有 0.85 GB,加载却占 5.53 GB。翻代码发现 Qwen3_5ForCausalLMMTP.__init__ 会给 draft 建自己的 embed_tokens 和 lm_head(各 248320×5120×2 byte = 2.5 GB),但运行时 eagle_worker_v2.init_lm_head() 会调用 set_embed_and_head():

    def set_embed_and_head(self, embed, head):
        del self.model.embed_tokens.weight      # 删掉自己的
        if not self.config.tie_word_embeddings:
            del self.lm_head.weight
        self.model.embed_tokens.weight = embed  # 换成目标模型的张量引用
        self.lm_head.weight = head
        torch.cuda.empty_cache()
    

    问题是:KV 池是在这之前就按"draft 还占着 5.53GB"算好尺寸的,之后 empty_cache() 释放出来的 5GB 就白空着了(所以才有 avail mem=4.17 GB 却只给 1.4 万 token 的怪现象)。

    修法:构造时直接把这两个权重换成 0 尺寸占位(del 仍能成功,之后会被目标张量替换):

    # Qwen3_5ForCausalLMMTP.__init__ 末尾插入
    self.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(0, device=dev))
    if not config.tie_word_embeddings:
        self.lm_head.weight = torch.nn.Parameter(torch.empty(0, device=dev))
    

    效果立竿见影:

    指标 打补丁前 打补丁后
    draft 加载显存 5.53 GB 0.79 GB
    KV 池 14161 token 45423 token(3.2 倍)

    提醒:weight 必须是 nn.Parameter,直接赋 torch.empty(0) 会报 TypeError: cannot assign ... as parameter 'weight'。

    6. 并发被 mamba 槽位卡住

    补丁之后最大上下文到 45K,但日志还有一条:

    max_running_requests is capped to 1 by the mamba state cache
      (max_mamba_cache_size=6, 5 state slots per request)
    

    Qwen3.8 是混合架构(48 层 GDN 线性注意力 + 16 层全注意力),投机解码下每个请求要占用 5 个 mamba 状态槽(要在验证失败时回滚)。默认/我们之前用的 6 槽只够 1 个请求,双路并发直接变成排队。

    解决:--max-mamba-cache-size 16 + --mamba-ssm-dtype bfloat16(SSM 状态用 bf16 存储,槽位成本减半)→ 并发恢复成 2,而且 mamba 显存从 0.98 GB 变成 1.20 GB(16 槽 bf16),中间态缓存反而更省。

    7. 真正的元凶:MTP 层数识别错误,把 KV 池砍到了 1/5

    修完上面的坑,KV 池到 45423,但 128K 需要 ≥131072,于是开始找显存。我先试了各种参数(mem-fraction 0.92→0.99、并发降到 1、prefill 压到 2048、hicache 缩到 4、关 CUDA graph、减投机步数),池子最多只到 60973,而且每次分配完都还空着 5-6 GB 显存 —— 说明不是显存不够,是算错了。

    于是往 sizing 代码里插了调试日志,一次就露馅了:

    无投机: cell_size=32768B (32.0KiB/token)  num_layers=16  eagle_draft_layers=None   → KV 池 279214
    有投机: cell_size=32768B → EAGLE 缩放生效:
            eagle_draft_layers=64  num_layers=16  →  cell_size=163840B  (×5)          → KV 池 44913
    

    pool_configurator.py 里有一段为投机解码准备的缩放:

    # EAGLE/STANDALONE: scale cell_size to account for draft model KV cache
    self._cell_size = int(
        self._cell_size * (1 + int(eagle_draft_num_layers) / int(num_layers))
    )
    

    而 eagle_draft_num_layers 的来源在 spec_aux_hidden_state.py:

    num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
    if num_nextn_predict_layers is not None:
        config.eagle_draft_num_layers = int(num_nextn_predict_layers)
    else:
        config.eagle_draft_num_layers = int(max(
            draft_model_config.num_hidden_layers,      # ← 因为 draft 路径就是模型自己,这里拿到 64
            draft_model_config.num_attention_layers,
        ))
    

    Qwen3.5/3.8 的 MTP 层数写在 text_config.mtp_num_hidden_layers 里(=1),不是 DeepSeek/GLM 用的 num_nextn_predict_layers。字段认不出来 → 走了 fallback → max(num_hidden_layers, ...) = 64(因为 --speculative-draft-model-path 指向模型自己的目录,读到的是整模型的 64 层)。

    而真正带 KV 的层只有 16 层(48 层是 GDN 线性注意力,不占 KV),于是缩放系数变成 1 + 64/16 = 5,每 token 成本被高估 5 倍,KV 池被砍到 1/5。

    修法(一行级的补丁)

    fix_mtp_draft_layers.py:在 fallback 之前先认 mtp_num_hidden_layers:

    elif _mtp_layers is not None and int(_mtp_layers) > 0:      # mtp_num_hidden_layers
        config.eagle_draft_num_layers = int(_mtp_layers)
    

    效果

    mem-fraction 修复前 KV 池 修复后 KV 池
    0.92 44913 211357(4.7 倍)
    0.97 59150 257367

    修好之后,单路 128K + MTP 直接就跑通了:

    配置 prompt prefill 解码 @128K 接受长度
    MTP 126863 tok 111.4 s(1139 tok/s) 37.8 t/s 2.00
    无 MTP(同 mamba 配置) 126863 tok 106.9 s(1187 tok/s) 34.4 t/s —

    也就是说 128K 上下文下 MTP 仍有 +10%(接受长度从短上下文的 3.76 掉到 2.0,所以增益变小,但没变负)。

    另外补充一句:社区那个 PR #37155 是真的,它把 draft 的 vocab 层(embed/lm_head)释放时机提前到 KV 池划分之前 —— 但那解决的是另一个问题(2.4 GiB 的 vocab 层白占),和这里的 5 倍 cell_size 是两码事。两处都值得修:vocab 层那个值 2.4 GB,层数这个值池子的 4.7 倍。

    一个还没解决的并发崩溃

    修好池子后,双路 2×60K 并发(两个请求都成功完成,51.4s)之后服务会崩在同一个 mamba 槽位断言上:

    AssertionError: Can not alloc mamba cache
      mem_cache/unified_cache/components/mamba_component.py:479 _alloc_mamba_slot
      ← cache_unfinished_req ← stash_chunked_request
    

    而且这个崩溃 --mamba-ssm-dtype bfloat16 + 16 槽 才会出现;换回 --max-mamba-cache-size 6 + 默认 float32,同样 2×64K 并发是稳的(只是投机下并发会被压成 1)。所以:

    • MTP + 单路 128K:--max-mamba-cache-size 6(默认 dtype),稳 ✅
    • MTP + 双路(短/中上下文):bf16 + 16 槽,但别跑 2×60K 这种长并发

    8. 顺带抓到一个会崩服务的 bug(强烈建议避开)

    为了让长上下文档多留并发槽位,我把 mamba 调到 16 槽 + bf16 之后,跑 2×64K 并发,服务整个崩了:

    AssertionError: Can not alloc mamba cache
      sglang/srt/mem_cache/unified_cache/components/mamba_component.py:479 _alloc_mamba_slot
      ← cache_unfinished_req
      ← stash_chunked_request        (chunked prefill 每分一块都要暂存 mamba 状态)
    → SIGQUIT → scheduler 异常 → 整个服务退出
    

    根因:chunked prefill 是分块的,每分一块都要暂存一次当前 mamba 状态;长提示分块多、并发两个就更容易把槽位耗光,而这里只有一句 assert,没有任何保护或降级,直接崩进程。

    换成验证过的 --max-mamba-cache-size 6 之后单路 128K 和 2×64K 都稳定(服务存活),但并发会被压成 1:2×64K 墙钟 84.4s ≈ 2×42s,是排队串行而不是并行。

    9. 顺便把不带 MTP 的单路 128K 验了

    测试 prompt prefill 解码 @128K
    单路 ~127K 123462 tok 103.0 s(1199 tok/s) 34.5 t/s
    同上重放 123286 tok 102.5 s(1202 tok/s) 37.1 t/s

    KV 池 289577 token,单路 128K 只占 43%。注意 fp8 KV 下长上下文的解码会从 41.5 掉到 34-37 t/s。

    10. 最终启动脚本(三档切换)

    跑 MTP 之前请先打两个补丁:patch_mtp_lmhead.py(draft vocab 层占位,省 4.7 GB)和
    fix_mtp_draft_layers.py(修 MTP 层数识别,KV 池 ×4.7)。两个都是幂等的。

    # ===== 档 A:MTP + 单路 128K(修复后新增,decode 37.8 t/s @128K)=====
    python -m sglang.launch_server \
      --model-path /mnt/sda6/download/qwen38-redhat-int4 \
      --language-only \
      --context-length 131072 --max-running-requests 1 \
      --mem-fraction-static 0.92 \
      --kv-cache-dtype fp8_e4m3 --page-size 1 \
      --max-mamba-cache-size 6 \
      --enable-hierarchical-cache --hicache-size 12 \
      --hicache-write-policy write_through \
      --speculative-algorithm NEXTN \
      --speculative-draft-model-path /mnt/sda6/download/qwen38-redhat-int4 \
      --speculative-eagle-topk 1 \
      --speculative-num-steps 3 --speculative-num-draft-tokens 4 \
      --trust-remote-code
    
    # ===== 档 B:MTP 速度优先(短/中上下文,code 89 t/s、双流聚合 170 t/s)=====
    python -m sglang.launch_server \
      --model-path /mnt/sda6/download/qwen38-redhat-int4 \
      --language-only \
      --context-length 32768 --max-running-requests 2 \
      --mem-fraction-static 0.92 \
      --kv-cache-dtype fp8_e4m3 --page-size 1 \
      --max-mamba-cache-size 16 --mamba-ssm-dtype bfloat16 \
      --enable-hierarchical-cache --hicache-size 12 \
      --hicache-write-policy write_through \
      --speculative-algorithm NEXTN \
      --speculative-draft-model-path /mnt/sda6/download/qwen38-redhat-int4 \
      --speculative-eagle-topk 1 \
      --speculative-num-steps 3 --speculative-num-draft-tokens 4 \
      --trust-remote-code
    
    # ===== 档 C:长上下文/双路(无 MTP,最稳,41.5 t/s)=====
    python -m sglang.launch_server \
      --model-path /mnt/sda6/download/qwen38-redhat-int4 \
      --language-only \
      --context-length 131072 --max-running-requests 2 \
      --mem-fraction-static 0.92 \
      --kv-cache-dtype fp8_e4m3 --page-size 1 \
      --max-mamba-cache-size 6 \
      --enable-hierarchical-cache --hicache-size 12 \
      --hicache-write-policy write_through \
      --trust-remote-code
    

    11. 给同是小白的提醒

    1. 换带 MTP 的权重前,先确认它带 model_mtp.safetensors,并检查 model.safetensors.index.json 里有没有引用它。
    2. Xet 存储的仓库必须 HF_HUB_DISABLE_XET=1,否则镜像站等于没用,还会 401。
    3. MTP 打开后第一件事是看 KV 池(日志里的 max_total_num_tokens),别以为服务起来了就没事 —— 池子被挤到 1.4 万 token 时,长提示会被 400 拒掉。
    4. 池子异常小的时候,别急着调参数,先去读 sizing 代码。我这次就是靠往 _profile_available_bytes / _compute_cell_size 里插几行日志,一眼看出 cell_size 被放大了 5 倍 —— 参数调到天亮也没用,因为那是算错而不是不够。
    5. max_running_requests is capped to 1 这条日志一定要看,混合 mamba 模型 + 投机解码下每请求要 5 个 mamba 槽,槽不够就变排队。
    6. --max-mamba-cache-size 别乱调大:它和 chunked prefill 的交互有个会崩整个服务的断言(Can not alloc mamba cache),调完一定要跑 2×64K 压测。
    7. 投机参数要自己扫(和扫 n-max 一样),steps 越大代码/回声越快、但创作类接受率反而下降:1/2 → 接受率 0.97/0.95/0.74;3/4 → 0.94/0.93/0.51;5/6 → 0.91/—/0.42。而且上下文越长接受长度越低:短上下文 3.76,128K 时只有 2.0(增益从 +115% 降到 +10%)。
    8. 贪婪解码下投机不改变结果:3 个测试题里 2 个输出和关闭 MTP 时逐字节一致,1 个措辞不同(bf16 下 batch-verify 与单 token 路径的数值抖动),不是投机算错。

    感谢论坛里关于 MTP/投机解码的讨论,让我这次少走了不少弯路。

    1 条回复 最后回复
    1
    • terryT 离线
      terryT 离线
      terry
      超级版主
      编写于 最后由 编辑
      #2

      第七节是精髓,就是不知道比起DFlash如何,DFlash2的版本,显卡运算量明显减少,除了一开始任务时,后期风扇都不转。

      油管:https://www.youtube.com/@抡锤者

      E 2 条回复 最后回复
      0
      • terryT terry

        第七节是精髓,就是不知道比起DFlash如何,DFlash2的版本,显卡运算量明显减少,除了一开始任务时,后期风扇都不转。

        E 离线
        E 离线
        Enigma
        德高望重
        编写于 最后由 编辑
        #3

        @terry 说:

        比起DFlash如何,DFlash2的版本

        谢谢坛主鼓励,我再去调试

        1 条回复 最后回复
        0
        • terryT terry

          第七节是精髓,就是不知道比起DFlash如何,DFlash2的版本,显卡运算量明显减少,除了一开始任务时,后期风扇都不转。

          E 离线
          E 离线
          Enigma
          德高望重
          编写于 最后由 编辑
          #4

          @terry 说:

          第七节是精髓,就是不知道比起DFlash如何,DFlash2的版本,显卡运算量明显减少,除了一开始任务时,后期风扇都不转。

          成功了,提升很大,准备发帖子

          1 条回复 最后回复
          0

          你好!看起来您对这段对话很感兴趣,但您还没有一个账号。

          厌倦了每次访问都刷到同样的帖子?您注册账号后,您每次返回时都能精准定位到您上次浏览的位置,并可选择接收新回复通知(通过邮件或推送通知)。您还能收藏书签、为帖子顶,向社区成员表达您的欣赏。

          有了你的建议,这篇帖子会更精彩哦 💗

          注册 登录
          回复
          • 在新帖中回复
          登录后回复
          • 从旧到新
          • 从新到旧
          • 最多赞同


          • 登录

          • 登录或注册以进行搜索。
          • 第一个帖子
            最后一个帖子
          0
          • 版块
          • 最新
          • 标签
          • 热门
          • 用户
          • 群组