Fuse GLM-5.3-Flash KDA projections and prefill metadata (#39688)
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
Yuxuan Zhang
co-authored by
Xinyuan Tong
Xinyuan Tong
parent
c1a1eb5f66
commit
c8eb54c41d
@@ -25,11 +25,19 @@ def _backend():
|
||||
|
||||
def _forward_batch(extend_lens, prefix_lens, track_seqlens, track_mask):
|
||||
return SimpleNamespace(
|
||||
forward_mode=SimpleNamespace(
|
||||
is_extend=lambda: True, is_target_verify=lambda: False
|
||||
),
|
||||
extend_seq_lens=torch.tensor(extend_lens),
|
||||
extend_prefix_lens=torch.tensor(prefix_lens),
|
||||
mamba_track_seqlens=torch.tensor(track_seqlens),
|
||||
mamba_track_mask=torch.tensor(track_mask),
|
||||
mamba_track_indices=torch.arange(100, 100 + len(extend_lens)),
|
||||
# Exercise the legacy GPU planner, not the CPU-metadata fast path.
|
||||
mamba_prefill_track_mask_cpu=None,
|
||||
mamba_track_seqlens_cpu=None,
|
||||
extend_seq_lens_cpu=None,
|
||||
extend_prefix_lens_cpu=None,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
import random
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
|
||||
Mamba2AttnBackend,
|
||||
MambaAttnBackendBase,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
||||
from sglang.srt.model_executor.forward_batch_info import (
|
||||
CaptureHiddenMode,
|
||||
ForwardBatch,
|
||||
ForwardMode,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.srt.speculative import spec_utils
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class NoHostRead(torch.Tensor):
|
||||
def cpu(self, *args, **kwargs):
|
||||
raise AssertionError("CPU tracking must not copy device metadata to the host")
|
||||
|
||||
|
||||
def make_batch(lengths, prefix, track_lens, mask, mirrored):
|
||||
def tensor(values):
|
||||
return torch.tensor(values).as_subclass(
|
||||
NoHostRead if mirrored else torch.Tensor
|
||||
)
|
||||
|
||||
return SimpleNamespace(
|
||||
batch_size=len(lengths),
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
extend_seq_lens=tensor(lengths),
|
||||
extend_prefix_lens=tensor(prefix),
|
||||
mamba_track_seqlens=tensor(track_lens),
|
||||
mamba_track_mask=tensor(mask),
|
||||
mamba_track_indices=tensor([37 + i * 17 for i in range(len(lengths))]),
|
||||
extend_seq_lens_cpu=lengths,
|
||||
extend_prefix_lens_cpu=prefix,
|
||||
mamba_track_seqlens_cpu=track_lens if mirrored else None,
|
||||
mamba_prefill_track_mask_cpu=mask if mirrored else None,
|
||||
)
|
||||
|
||||
|
||||
def make_forward_batch(lengths, starts, cpu_lengths, mode=ForwardMode.EXTEND):
|
||||
return ForwardBatch(
|
||||
forward_mode=mode,
|
||||
batch_size=len(lengths),
|
||||
input_ids=torch.zeros(sum(lengths), dtype=torch.int64),
|
||||
req_pool_indices=torch.arange(len(lengths)),
|
||||
seq_lens=torch.tensor(lengths),
|
||||
seq_lens_sum=sum(lengths),
|
||||
out_cache_loc=torch.zeros(sum(lengths), dtype=torch.int64),
|
||||
extend_start_loc=torch.tensor(starts, dtype=torch.int32),
|
||||
extend_seq_lens=torch.tensor(lengths, dtype=torch.int32),
|
||||
extend_seq_lens_cpu=cpu_lengths,
|
||||
)
|
||||
|
||||
|
||||
def make_metadata_backend():
|
||||
backend = object.__new__(MambaAttnBackendBase)
|
||||
backend.device = "cpu"
|
||||
backend.topk = 1
|
||||
backend.req_to_token_pool = SimpleNamespace(
|
||||
get_mamba_indices=lambda rows: rows,
|
||||
translate_mamba_indices=lambda slots: slots,
|
||||
)
|
||||
return backend
|
||||
|
||||
|
||||
class TestMambaPrefillTrackMetadata(unittest.TestCase):
|
||||
def test_cpu_plan_matches_existing_tensor_planner(self):
|
||||
rng = random.Random(2026)
|
||||
for backend_type in (MambaAttnBackendBase, Mamba2AttnBackend):
|
||||
for chunk in (16, 64, 128):
|
||||
backend = object.__new__(backend_type)
|
||||
backend.device = "cpu"
|
||||
backend._mamba_chunk_size = chunk
|
||||
cases = [
|
||||
(
|
||||
[chunk + 6, 2 * chunk + 1, chunk],
|
||||
[0, 2 * chunk, 0],
|
||||
[chunk + 1, 3 * chunk + 1, chunk],
|
||||
[True, True, True],
|
||||
),
|
||||
(
|
||||
[2 * chunk + 1, 1, 1],
|
||||
[chunk, 0, 0],
|
||||
[2 * chunk + 1, 0, 0],
|
||||
[True, False, False],
|
||||
),
|
||||
([1, chunk], [0, 0], [0, 0], [False, False]),
|
||||
]
|
||||
for _ in range(10):
|
||||
lengths = [rng.randrange(1, chunk * 6) for _ in range(5)]
|
||||
prefix = [rng.randrange(4) * chunk for _ in lengths]
|
||||
cases.append(
|
||||
(
|
||||
lengths,
|
||||
prefix,
|
||||
[
|
||||
p + rng.randrange(1, n + 1)
|
||||
for p, n in zip(prefix, lengths)
|
||||
],
|
||||
[bool(rng.randrange(2)) for _ in lengths],
|
||||
)
|
||||
)
|
||||
for lengths, prefix, track, mask in cases:
|
||||
slots = torch.tensor([111 - i * 5 for i in range(len(lengths))])
|
||||
with self.subTest(
|
||||
backend=backend_type.__name__, chunk=chunk, mask=mask
|
||||
):
|
||||
expected = backend._init_track_ssm_indices(
|
||||
slots, make_batch(lengths, prefix, track, mask, False)
|
||||
)
|
||||
actual = backend._init_track_ssm_indices(
|
||||
slots.as_subclass(NoHostRead),
|
||||
make_batch(lengths, prefix, track, mask, True),
|
||||
)
|
||||
for result, reference in zip(actual, expected):
|
||||
if reference is None:
|
||||
self.assertIsNone(result)
|
||||
else:
|
||||
torch.testing.assert_close(result, reference)
|
||||
|
||||
def test_verify_and_incomplete_mirrors_use_existing_planner(self):
|
||||
batch = make_batch([64], [0], [64], [True], True)
|
||||
eligible = MambaAttnBackendBase._has_cpu_prefill_track_metadata
|
||||
self.assertTrue(eligible(batch))
|
||||
batch.forward_mode = ForwardMode.TARGET_VERIFY
|
||||
self.assertFalse(eligible(batch))
|
||||
batch.forward_mode = ForwardMode.EXTEND
|
||||
batch.mamba_track_seqlens_cpu = None
|
||||
self.assertFalse(eligible(batch))
|
||||
batch.mamba_track_seqlens_cpu = [64, 0]
|
||||
self.assertFalse(eligible(batch))
|
||||
|
||||
def test_logical_token_extent_avoids_scalar_reads_with_valid_cpu_lengths(self):
|
||||
backend = make_metadata_backend()
|
||||
original_int = torch.Tensor.__int__
|
||||
for lengths, starts, cpu_lengths, tbo_range, expected, scalar_reads in (
|
||||
([3, 5], [0, 3], [3, 5], None, 8, 0),
|
||||
([3, 5, 0], [0, 3, 8], [3, 5, 0], None, 8, 0),
|
||||
([3, 5], [4, 7], None, None, 12, 1),
|
||||
([3, 5], [4, 7], [8], None, 12, 1),
|
||||
([3, 5], [4, 7], [3, 5], (4, 12), 12, 1),
|
||||
):
|
||||
with self.subTest(cpu_lengths=cpu_lengths, tbo_range=tbo_range):
|
||||
batch = make_forward_batch(lengths, starts, cpu_lengths)
|
||||
batch.tbo_parent_token_range = tbo_range
|
||||
reads = []
|
||||
|
||||
def read_scalar(tensor):
|
||||
if scalar_reads == 0:
|
||||
raise AssertionError(
|
||||
"Valid CPU lengths must avoid scalar reads"
|
||||
)
|
||||
reads.append(tensor.clone())
|
||||
return original_int(tensor)
|
||||
|
||||
with patch.object(torch.Tensor, "__int__", read_scalar):
|
||||
metadata = backend._forward_metadata(batch)
|
||||
self.assertEqual(metadata.logical_num_tokens, expected)
|
||||
self.assertEqual(len(reads), scalar_reads)
|
||||
self.assertEqual(metadata.query_start_loc[-1].item(), expected)
|
||||
|
||||
def test_verify_decode_and_idle_ignore_stale_cpu_token_lengths(self):
|
||||
backend = make_metadata_backend()
|
||||
for mode, lengths, expected_starts in (
|
||||
(ForwardMode.TARGET_VERIFY, [3, 3], [0, 3, 6]),
|
||||
(ForwardMode.DECODE, [1, 1], [0, 1, 2]),
|
||||
(ForwardMode.IDLE, [], [0]),
|
||||
):
|
||||
with self.subTest(mode=mode):
|
||||
batch = make_forward_batch(
|
||||
lengths, [0, 3][: len(lengths)], [100, 200], mode
|
||||
)
|
||||
if mode == ForwardMode.TARGET_VERIFY:
|
||||
batch.spec_info = SimpleNamespace(
|
||||
ragged_verify_layout=None, draft_token_num=3
|
||||
)
|
||||
with patch.object(
|
||||
torch.Tensor,
|
||||
"__int__",
|
||||
side_effect=AssertionError("This mode must not read token scalars"),
|
||||
):
|
||||
metadata = backend._forward_metadata(batch)
|
||||
self.assertIsNone(metadata.logical_num_tokens)
|
||||
torch.testing.assert_close(
|
||||
metadata.query_start_loc,
|
||||
torch.tensor(expected_starts, dtype=torch.int32),
|
||||
)
|
||||
|
||||
def test_forward_snapshot_and_padding_do_not_mutate_scheduler_lists(self):
|
||||
override = get_context().override_server_args(device="cpu")
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
batch = ScheduleBatch(
|
||||
reqs=[SimpleNamespace(rid="one", lora_id=None, token_type_ids=None)],
|
||||
device="cpu",
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
input_ids=torch.tensor([3]),
|
||||
req_pool_indices=torch.tensor([2]),
|
||||
seq_lens=torch.tensor([65]),
|
||||
seq_lens_cpu=torch.tensor([65]),
|
||||
seq_lens_sum=65,
|
||||
out_cache_loc=torch.tensor([1]),
|
||||
extend_lens=[1],
|
||||
prefix_lens=[64],
|
||||
extend_num_tokens=1,
|
||||
mamba_track_mask=torch.tensor([True]),
|
||||
mamba_track_seqlens=torch.tensor([65]),
|
||||
mamba_prefill_track_mask_cpu=[True],
|
||||
mamba_track_seqlens_cpu=[65],
|
||||
)
|
||||
runner = SimpleNamespace(
|
||||
device="cpu",
|
||||
model_config=SimpleNamespace(
|
||||
requires_mm_token_modalities=False, model_is_mrope=False
|
||||
),
|
||||
kv_index_translator=SimpleNamespace(rebind_write_loc=lambda forward: None),
|
||||
prefill_attention_backend_str="torch_native",
|
||||
ngram_embedding_manager=SimpleNamespace(enabled=False),
|
||||
lora_manager=None,
|
||||
ps=SimpleNamespace(attn_dcp_size=1),
|
||||
attn_backend=SimpleNamespace(
|
||||
get_cpu_graph_seq_len_fill_value=lambda: 1,
|
||||
get_cuda_graph_seq_len_fill_value=lambda: 1,
|
||||
),
|
||||
)
|
||||
forward = ForwardBatch.init_new(
|
||||
batch,
|
||||
runner,
|
||||
capture_hidden_mode=CaptureHiddenMode.NULL,
|
||||
return_hidden_states_before_norm=False,
|
||||
)
|
||||
for target, source in (
|
||||
("mamba_prefill_track_mask_cpu", "mamba_prefill_track_mask_cpu"),
|
||||
("mamba_track_seqlens_cpu", "mamba_track_seqlens_cpu"),
|
||||
("extend_seq_lens_cpu", "extend_lens"),
|
||||
("extend_prefix_lens_cpu", "prefix_lens"),
|
||||
):
|
||||
self.assertEqual(getattr(forward, target), getattr(batch, source))
|
||||
self.assertIsNot(getattr(forward, target), getattr(batch, source))
|
||||
forward._pad_inputs_to_size(runner, num_tokens=3, bs=3)
|
||||
self.assertEqual(batch.mamba_prefill_track_mask_cpu, [True])
|
||||
self.assertEqual(batch.mamba_track_seqlens_cpu, [65])
|
||||
self.assertEqual(batch.extend_lens, [1])
|
||||
self.assertEqual(batch.prefix_lens, [64])
|
||||
for host, device, expected in (
|
||||
("mamba_prefill_track_mask_cpu", "mamba_track_mask", [True, False, False]),
|
||||
("mamba_track_seqlens_cpu", "mamba_track_seqlens", [65, 0, 0]),
|
||||
("extend_seq_lens_cpu", "extend_seq_lens", [1, 0, 0]),
|
||||
("extend_prefix_lens_cpu", "extend_prefix_lens", [64, 0, 0]),
|
||||
):
|
||||
self.assertEqual(getattr(forward, host), expected)
|
||||
self.assertEqual(getattr(forward, device).tolist(), expected)
|
||||
|
||||
def test_decode_and_verify_clear_prefill_lists_without_losing_snapshot(self):
|
||||
for verify in (False, True):
|
||||
with self.subTest(verify=verify):
|
||||
batch = ScheduleBatch(
|
||||
reqs=[],
|
||||
spec_algorithm=SimpleNamespace(is_none=lambda: False),
|
||||
mamba_track_mask=torch.tensor([True]),
|
||||
mamba_track_seqlens=torch.tensor([65]),
|
||||
mamba_prefill_track_mask_cpu=[True],
|
||||
mamba_track_seqlens_cpu=[65],
|
||||
)
|
||||
snapshot = batch.copy()
|
||||
if verify:
|
||||
settings = SimpleNamespace(
|
||||
mamba=SimpleNamespace(
|
||||
enable_mamba_extra_buffer=True,
|
||||
enable_mamba_extra_buffer_lazy=False,
|
||||
)
|
||||
)
|
||||
with (
|
||||
patch.object(spec_utils, "get_exec", return_value=settings),
|
||||
patch.object(spec_utils, "set_mamba_track_indices_from_reqs"),
|
||||
):
|
||||
spec_utils.prepare_mamba_track_for_verify(batch)
|
||||
self.assertIsNone(batch.mamba_track_mask)
|
||||
self.assertIsNone(batch.mamba_track_seqlens)
|
||||
else:
|
||||
with patch.object(spec_utils, "spec_prepare_for_decode"):
|
||||
batch.prepare_for_decode()
|
||||
self.assertIsNone(batch.mamba_prefill_track_mask_cpu)
|
||||
self.assertIsNone(batch.mamba_track_seqlens_cpu)
|
||||
self.assertEqual(snapshot.mamba_prefill_track_mask_cpu, [True])
|
||||
self.assertEqual(snapshot.mamba_track_seqlens_cpu, [65])
|
||||
self.assertIsNone(snapshot.mamba_track_mask_cpu)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,182 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||
from sglang.srt.models import glm5_next
|
||||
from sglang.srt.runtime_context import get_context, get_parallel
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
PREFIX = "model.layers.0.self_attn"
|
||||
QKV = ("q_proj", "k_proj", "v_proj")
|
||||
BFG = ("b_proj", "f_a_proj", "g_a_proj", "f_b_proj", "g_b_proj")
|
||||
|
||||
|
||||
class MockQuantizedLinearMethod:
|
||||
"""Keep dense storage so the test isolates routing and checkpoint loading."""
|
||||
|
||||
create_weights = UnquantizedLinearMethod.create_weights
|
||||
|
||||
def apply(self, layer, x, bias=None):
|
||||
return F.linear(x, layer.weight, bias)
|
||||
|
||||
|
||||
class MockFp8Config:
|
||||
def __init__(self, ignored):
|
||||
self.ignored_layers = {f"{PREFIX}.{name}" for name in ignored}
|
||||
|
||||
def get_name(self):
|
||||
return "fp8"
|
||||
|
||||
def get_quant_method(self, layer, prefix):
|
||||
names = (
|
||||
[prefix.replace("qkv_proj", name) for name in QKV]
|
||||
if prefix.endswith(".qkv_proj")
|
||||
else [prefix]
|
||||
)
|
||||
if all(name in self.ignored_layers for name in names):
|
||||
return UnquantizedLinearMethod()
|
||||
return MockQuantizedLinearMethod()
|
||||
|
||||
|
||||
class TestGlm5NextBfgFusion(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.addCleanup(torch.set_default_dtype, torch.get_default_dtype())
|
||||
torch.set_default_dtype(torch.float32)
|
||||
override = get_context().override_server_args(
|
||||
device="cpu", enable_lora=False, lora_paths=None
|
||||
)
|
||||
override.install()
|
||||
self.addCleanup(override.restore)
|
||||
patcher = patch.object(
|
||||
UnquantizedLinearMethod,
|
||||
"apply",
|
||||
MockQuantizedLinearMethod.apply,
|
||||
)
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
@torch.no_grad()
|
||||
def test_projection_loading_matches_unfused_reference(self):
|
||||
torch.manual_seed(42)
|
||||
hidden, heads, dim = 16, 4, 8
|
||||
shapes = {name: (heads * dim, hidden) for name in QKV}
|
||||
shapes.update(
|
||||
b_proj=(heads, hidden),
|
||||
f_a_proj=(dim, hidden),
|
||||
g_a_proj=(dim, hidden),
|
||||
f_b_proj=(heads * dim, dim),
|
||||
g_b_proj=(heads * dim, dim),
|
||||
)
|
||||
weights = {name: torch.randn(shape) for name, shape in shapes.items()}
|
||||
x = torch.randn(7, hidden)
|
||||
for ignored, expected_route in (
|
||||
(QKV + BFG, (True, False)),
|
||||
(BFG, (False, True)),
|
||||
((), (False, False)),
|
||||
):
|
||||
for attn_tp, rank in ((1, 0), (2, 0), (2, 1)):
|
||||
with (
|
||||
self.subTest(route=expected_route, attn_tp=attn_tp, rank=rank),
|
||||
get_parallel().override(
|
||||
tp_size=4, tp_rank=3, attn_tp_size=attn_tp, attn_tp_rank=rank
|
||||
),
|
||||
):
|
||||
quant = MockFp8Config(ignored)
|
||||
attention = glm5_next.Glm5NextLinearAttention(
|
||||
layer_idx=0,
|
||||
hidden_size=hidden,
|
||||
config=SimpleNamespace(
|
||||
linear_attn_config={
|
||||
"head_dim": dim,
|
||||
"num_heads": heads,
|
||||
"short_conv_kernel_size": 4,
|
||||
}
|
||||
),
|
||||
quant_config=quant,
|
||||
prefix=PREFIX,
|
||||
)
|
||||
self.assertEqual(
|
||||
(attention.do_fuse_qkvbfg, attention.fuse_bfg), expected_route
|
||||
)
|
||||
for parameter in attention.parameters():
|
||||
parameter.fill_(torch.nan)
|
||||
model = SimpleNamespace(
|
||||
config=SimpleNamespace(n_routed_experts=0),
|
||||
num_fused_shared_experts=0,
|
||||
quant_config=quant,
|
||||
named_parameters=lambda: (
|
||||
(f"{PREFIX}.{name}", param)
|
||||
for name, param in attention.named_parameters()
|
||||
),
|
||||
)
|
||||
with patch.object(
|
||||
glm5_next.DeepseekV2WeightLoaderMixin, "post_load_weights"
|
||||
):
|
||||
glm5_next.Glm5NextForConditionalGeneration.load_weights(
|
||||
model,
|
||||
[
|
||||
(f"{PREFIX}.{name}.weight", w)
|
||||
for name, w in weights.items()
|
||||
],
|
||||
)
|
||||
|
||||
def linear(value, name):
|
||||
weight = weights[name]
|
||||
if name not in ("f_a_proj", "g_a_proj"):
|
||||
weight = weight.chunk(attn_tp, dim=0)[rank]
|
||||
return F.linear(value, weight)
|
||||
|
||||
expected = (
|
||||
torch.cat([linear(x, name) for name in QKV], dim=-1),
|
||||
linear(x, "b_proj"),
|
||||
linear(linear(x, "f_a_proj"), "f_b_proj"),
|
||||
linear(linear(x, "g_a_proj"), "g_b_proj"),
|
||||
)
|
||||
forward = (
|
||||
attention.forward_qkvbfg_fused
|
||||
if attention.do_fuse_qkvbfg
|
||||
else attention.forward_qkvbfg
|
||||
)
|
||||
for actual, reference in zip(forward(x, None), expected):
|
||||
torch.testing.assert_close(
|
||||
actual, reference, atol=1e-5, rtol=1e-5
|
||||
)
|
||||
|
||||
def test_each_quantized_gate_projection_disables_fusion(self):
|
||||
for quantized in BFG:
|
||||
quant = MockFp8Config(name for name in QKV + BFG if name != quantized)
|
||||
for packed in ("fused_qkvbfg_a_proj", "fused_bfg_a_proj"):
|
||||
with self.subTest(quantized=quantized, packed=packed):
|
||||
self.assertFalse(
|
||||
glm5_next.Glm5NextLinearAttention._can_fuse_proj(
|
||||
quant, PREFIX, packed, "fused_fg_b_proj"
|
||||
)
|
||||
)
|
||||
|
||||
def test_lora_disables_full_and_bfg_fusion(self):
|
||||
for enable_lora, paths in ((True, None), (False, ["adapter"])):
|
||||
with patch.object(
|
||||
glm5_next,
|
||||
"get_lora",
|
||||
return_value=SimpleNamespace(enable_lora=enable_lora, lora_paths=paths),
|
||||
):
|
||||
for quant in (None, MockFp8Config(QKV + BFG)):
|
||||
for packed in ("fused_qkvbfg_a_proj", "fused_bfg_a_proj"):
|
||||
with self.subTest(
|
||||
enabled=enable_lora, paths=paths, packed=packed
|
||||
):
|
||||
self.assertFalse(
|
||||
glm5_next.Glm5NextLinearAttention._can_fuse_proj(
|
||||
quant, PREFIX, packed, "fused_fg_b_proj"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user