[HotFix][Ling 2.6] Fix HybridLinearAttn dispatcher for Ling-2.6 (#26474)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -811,9 +811,18 @@ class HybridLinearAttnBackend(AttentionBackend):
|
|||||||
def _is_full_attn(
|
def _is_full_attn(
|
||||||
self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
|
self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
|
||||||
) -> bool:
|
) -> bool:
|
||||||
# Dispatch by the layer's runtime type
|
# Explicit linear-attention subclass → strong linear signal (KDA, GDN,
|
||||||
|
# Qwen3-Next, Qwen3.5 main linear layers).
|
||||||
if isinstance(layer, RadixLinearAttention):
|
if isinstance(layer, RadixLinearAttention):
|
||||||
return False
|
return False
|
||||||
|
# Some hybrid models (Ling-2.5/2.6) wrap their linear layers in plain
|
||||||
|
# `RadixAttention` rather than `RadixLinearAttention`. Those wrappers
|
||||||
|
# set `_is_linear_attention=True` on the attn module so we can
|
||||||
|
# distinguish them from full-attention RadixAttention instances —
|
||||||
|
# including MTP/NEXTN draft layers, which are full and must default to
|
||||||
|
# the full-attn path.
|
||||||
|
if layer is not None and getattr(layer, "_is_linear_attention", False):
|
||||||
|
return False
|
||||||
if isinstance(layer, RadixAttention):
|
if isinstance(layer, RadixAttention):
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|||||||
@@ -508,6 +508,12 @@ class BailingMoELinearAttention(nn.Module):
|
|||||||
quant_config=quant_config,
|
quant_config=quant_config,
|
||||||
prefix=f"{prefix}.attn",
|
prefix=f"{prefix}.attn",
|
||||||
)
|
)
|
||||||
|
# Marker for HybridLinearAttnBackend._is_full_attn: Bailing wraps
|
||||||
|
# linear-attention layers in a plain RadixAttention, so the
|
||||||
|
# dispatcher can't tell from the type alone that this is a linear
|
||||||
|
# layer (would otherwise default to the full-attn backend, e.g. the
|
||||||
|
# same way MTP/NEXTN draft layers are routed).
|
||||||
|
self.attn._is_linear_attention = True
|
||||||
|
|
||||||
self.group_norm_size = getattr(config, "group_norm_size", 1)
|
self.group_norm_size = getattr(config, "group_norm_size", 1)
|
||||||
self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))
|
self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))
|
||||||
|
|||||||
@@ -0,0 +1,54 @@
|
|||||||
|
"""GSM8K accuracy test for Ling-2.6-flash (BailingMoELinearForCausalLM).
|
||||||
|
|
||||||
|
Guards the hybrid linear / full attention dispatcher: Ling-2.5/2.6
|
||||||
|
has 32 layers with `layer_group_size=8`, so layers {7, 15, 23, 31}
|
||||||
|
are full attention (MLA) and the rest are linear (Lightning seg_la).
|
||||||
|
Runs on the 8-GPU H200 runner with TP=4.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||||
|
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||||
|
|
||||||
|
register_cuda_ci(est_time=600, stage="base-c", runner_config="8-gpu-h200")
|
||||||
|
|
||||||
|
|
||||||
|
class TestLing26Flash(GSM8KMixin, DefaultServerBase):
|
||||||
|
model = "inclusionAI/Ling-2.6-flash"
|
||||||
|
|
||||||
|
# Native 128K context (no YaRN) — avoids the
|
||||||
|
# SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN env-var dance and keeps the
|
||||||
|
# smoke test focused on the dispatcher / hybrid-attention path.
|
||||||
|
other_args = [
|
||||||
|
"--tp-size",
|
||||||
|
"4",
|
||||||
|
"--trust-remote-code",
|
||||||
|
"--mamba-scheduler-strategy",
|
||||||
|
"extra_buffer",
|
||||||
|
"--mem-fraction-static",
|
||||||
|
"0.75",
|
||||||
|
"--max-running-requests",
|
||||||
|
"64",
|
||||||
|
"--max-mamba-cache-size",
|
||||||
|
"256",
|
||||||
|
# MTP path also exercises the dispatcher (draft + target verify),
|
||||||
|
# so keep it on to maximize coverage.
|
||||||
|
"--speculative-algorithm",
|
||||||
|
"NEXTN",
|
||||||
|
"--speculative-num-steps",
|
||||||
|
"3",
|
||||||
|
"--speculative-eagle-topk",
|
||||||
|
"1",
|
||||||
|
"--speculative-num-draft-tokens",
|
||||||
|
"4",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Observed 0.825 on H200 TP=4 + NEXTN MTP with default 200-question GSM8K
|
||||||
|
# (the model card's 0.96 is from full 1319-question runs of the 1T model).
|
||||||
|
gsm8k_accuracy_thres = 0.825
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main(verbosity=3)
|
||||||
Reference in New Issue
Block a user