minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix Qwen3.5 GDN multi-item scoring (#33922)

Browse Source 
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
DAI0818
2026-09-09 22:30:39 -07:00
committed by GitHub
co-authored by Po-Han Huang
parent 6f481ad0e3
commit 03e4c06589
13 changed files with 842 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_chunk_gated_delta_rule.py
+83
View File
@@ -161,6 +161,89 @@ class TestChunkGatedDeltaRule(unittest.TestCase):
def test_batch_32(self):
self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256)
@unittest.skipUnless(
torch.cuda.is_available(),
"The read-only initial-state path is not implemented by the XPU chunk kernel",
)
def test_read_only_initial_state_supports_duplicate_indices(self):
"""MIS item branches may share one query-end state without updating it."""
device = get_device()
dtype = torch.bfloat16
batch_size, tokens_per_item = 3, 65
num_heads, key_dim, value_dim = 4, 32, 32
total_tokens = batch_size * tokens_per_item
torch.manual_seed(1234)
pool_init = torch.randn(
4,
num_heads,
value_dim,
key_dim,
dtype=torch.float32,
device=device,
)
duplicate_indices = torch.tensor([2, 2, 2], dtype=torch.int32, device=device)
cu_seqlens = torch.arange(
0,
total_tokens + 1,
tokens_per_item,
dtype=torch.int32,
device=device,
)
q = torch.randn(1, total_tokens, num_heads, key_dim, dtype=dtype, device=device)
k = torch.randn_like(q)
v = torch.randn(
1, total_tokens, num_heads, value_dim, dtype=dtype, device=device
)
g = torch.nn.functional.logsigmoid(
torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device)
)
beta = torch.sigmoid(
torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device)
)
expected, _ = self._run_reference(
pool_init, duplicate_indices, q, k, v, g, beta
)
actual_pool = pool_init.clone()
actual, _, _ = chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=actual_pool,
initial_state_indices=duplicate_indices,
cu_seqlens=cu_seqlens,
head_first=False,
use_qk_l2norm_in_kernel=True,
inplace_update=False,
)
torch.testing.assert_close(
actual.float(), expected.float(), atol=self.ATOL, rtol=self.RTOL
)
self.assertTrue(torch.equal(actual_pool, pool_init))
updating_pool = pool_init.clone()
chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=updating_pool,
initial_state_indices=torch.arange(
batch_size, dtype=torch.int32, device=device
),
cu_seqlens=cu_seqlens,
head_first=False,
use_qk_l2norm_in_kernel=True,
)
self.assertFalse(
torch.equal(updating_pool[:batch_size], pool_init[:batch_size])
)
# ------------------------------------------------------------------
# Head count sweep
# ------------------------------------------------------------------
test/registered/attention/unittests/gdn/test_triton.py
+30
View File
@@ -216,6 +216,36 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
with self.subTest(case=case.name, backend=case.backend):
run_gdn_attention_case(self, case)
def test_multi_item_scoring_mixed_batch_with_empty_query(self):
case = GDNAttentionCase(
name="gdn_mis_mixed_batch_empty_query",
backend="triton",
forward_mode=ForwardMode.EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=1,
prefix_lens=(0, 0),
extend_lens=(9, 7),
mis_delimiter_indices=((0, 3, 8), (4, 6)),
conv_history_weight=0.25,
)
run_gdn_attention_case(self, case)
def test_multi_item_scoring_crosses_chunk_boundaries(self):
case = GDNAttentionCase(
name="gdn_mis_chunk_boundaries",
backend="triton",
forward_mode=ForwardMode.EXTEND,
num_k_heads=2,
num_v_heads=2,
page_size=1,
prefix_lens=(0,),
extend_lens=(198,),
mis_delimiter_indices=((5, 68, 132, 197),),
conv_history_weight=0.25,
)
run_gdn_attention_case(self, case, max_context_len=256)
# Layout-robustness. See dense/test_triton.py for the rationale.
# shuffled_pages is the default for all tests; this method opts
# into the more aggressive interleaved_pages + non_monotonic_extend.
test/registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py
+133
View File
@@ -0,0 +1,133 @@
"""End-to-end MIS coverage for the hybrid Qwen3.5 GDN architecture."""
import asyncio
import os
import unittest
import torch
from transformers import AutoTokenizer
from sglang.srt.entrypoints.engine import Engine
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST,
CustomTestCase,
)
register_cuda_ci(est_time=240, stage="extra-a", runner_config="1-gpu-large")
class TestQwen35GDNMultiItemScoring(CustomTestCase):
model = os.environ.get(
"QWEN35_GDN_TEST_MODEL", DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
)
atol = 2e-2
rtol = 2e-2
@classmethod
def setUpClass(cls):
cls.engine = Engine(
model_path=cls.model,
trust_remote_code=True,
dtype="bfloat16",
enable_mis=True,
attention_backend="flashinfer",
linear_attn_prefill_backend="triton",
disable_radix_cache=True,
chunked_prefill_size=-1,
mem_fraction_static=0.8,
log_level="error",
)
tokenizer = AutoTokenizer.from_pretrained(cls.model, trust_remote_code=True)
cls.label_token_ids = [
tokenizer.encode(label, add_special_tokens=False)[0]
for label in (" yes", " no")
]
@classmethod
def tearDownClass(cls):
if getattr(cls, "engine", None) is not None:
cls.engine.shutdown()
torch.cuda.empty_cache()
def _score(self, query, items):
return self.engine.score(
query=query,
items=items,
label_token_ids=self.label_token_ids,
apply_softmax=False,
).scores
async def _async_score(self, query, items):
result = await self.engine.async_score(
query=query,
items=items,
label_token_ids=self.label_token_ids,
apply_softmax=False,
)
return result.scores
def _pointwise(self, query, items):
return [self._score(query, [item])[0] for item in items]
def _assert_scores_close(self, actual, expected):
torch.testing.assert_close(
torch.tensor(actual),
torch.tensor(expected),
atol=self.atol,
rtol=self.rtol,
)
def test_batched_matches_pointwise_for_varied_requests(self):
cases = [
(
"Decide whether each statement is true:",
["The sky is blue.", "Two plus two is five.", "Water freezes."],
),
("", ["empty query short", "empty query with a much longer item " * 8]),
("Classify:", ["one"]),
(
"Judge each passage:",
["tiny", "medium length passage " * 5, "long passage " * 24],
),
]
for query, items in cases:
with self.subTest(query=query, item_count=len(items)):
self._assert_scores_close(
self._score(query, items), self._pointwise(query, items)
)
def test_sibling_changes_and_reordering_do_not_change_target(self):
query = "Rate each answer:"
target = "The target answer remains unchanged."
baseline = self._score(query, [target, "sibling A", "sibling B"])[0]
changed = self._score(
query, [target, "a completely different sibling " * 8, "sibling B"]
)[0]
reordered = self._score(query, ["sibling B", "sibling A", target])[2]
self._assert_scores_close(changed, baseline)
self._assert_scores_close(reordered, baseline)
def test_concurrent_requests_match_pointwise(self):
cases = [
("Is it an animal?", ["cat", "table", "blue whale"]),
("", ["alpha", "beta"]),
("Choose:", ["first"]),
("Is it a city?", ["Paris", "bread", "Tokyo", "chair"]),
]
expected = [self._pointwise(query, items) for query, items in cases]
async def gather_scores():
return await asyncio.gather(
*(self._async_score(query, items) for query, items in cases)
)
actual = self.engine.loop.run_until_complete(gather_scores())
for actual_scores, expected_scores in zip(actual, expected):
self._assert_scores_close(actual_scores, expected_scores)
if __name__ == "__main__":
unittest.main()
test/registered/layers/mamba/test_causal_conv1d.py
+57
View File
@@ -119,6 +119,63 @@ def causal_conv1d_update_ref(
return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
def test_causal_conv1d_branches_from_cloned_query_state():
device = get_device()
dtype = torch.bfloat16
torch.manual_seed(7)
dim, width = 96, 4
item_lens = [2, 5, 7]
total_tokens = sum(item_lens)
source_state = torch.randn(1, dim, width - 1, device=device, dtype=dtype)
source_state_before = source_state.clone()
branch_states = source_state.expand(len(item_lens), -1, -1).contiguous().clone()
x = torch.randn(dim, total_tokens, device=device, dtype=dtype)
weight = torch.randn(dim, width, device=device, dtype=dtype)
bias = torch.randn(dim, device=device, dtype=dtype)
query_start_loc = torch.tensor(
[0, *torch.tensor(item_lens).cumsum(0).tolist()],
dtype=torch.int32,
device=device,
)
cache_indices = torch.arange(len(item_lens), dtype=torch.int32, device=device)
actual = causal_conv1d_fn(
x,
weight,
bias=bias,
conv_states=branch_states,
query_start_loc=query_start_loc,
seq_lens_cpu=torch.tensor(item_lens),
cache_indices=cache_indices,
has_initial_state=torch.ones(len(item_lens), dtype=torch.bool, device=device),
activation="silu",
pad_slot_id=PAD_SLOT_ID,
)
expected_outputs = []
expected_states = []
for item in torch.split(x, item_lens, dim=-1):
item_output, item_state = causal_conv1d_ref(
item.unsqueeze(0),
weight,
bias,
initial_states=source_state,
return_final_states=True,
activation="silu",
)
expected_outputs.append(item_output.squeeze(0))
expected_states.append(item_state.squeeze(0))
torch.testing.assert_close(
actual, torch.cat(expected_outputs, dim=-1), atol=5e-2, rtol=1e-2
)
torch.testing.assert_close(
branch_states, torch.stack(expected_states), atol=5e-2, rtol=1e-2
)
assert torch.equal(source_state, source_state_before)
@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
@pytest.mark.parametrize("silu_activation", [True])
@pytest.mark.parametrize("has_bias", [True])
test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
+19
View File
@@ -13,6 +13,7 @@ from sglang.srt.layers.attention.linear.gdn_backend import (
GDNKernelDispatcher,
_validate_gdn_linear_attn_backends,
flashinfer_gdn_prefill_default,
validate_gdn_mis_backend,
)
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
maybe_build_flashinfer_checkpoint_plan,
@@ -59,6 +60,7 @@ def make_runner(
mamba_radix_cache_strategy="no_buffer",
enable_dynamic_chunking=False,
chunked_prefill_size=8192,
enable_mis=False,
)
fields.update(arg_overrides)
args = _publish(testcase, **fields)
@@ -79,6 +81,22 @@ def make_runner(
class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
def test_mis_requires_triton_prefill_backend(self):
runner = make_runner(self, enable_mis=True)
with self.assertRaisesRegex(ValueError, "Triton linear-attention prefill"):
validate_gdn_mis_backend(LinearAttnKernelBackend.FLASHINFER)
def test_mis_rejects_page_major_layout(self):
make_runner(self, enable_mis=True, enable_page_major_kv_layout=True)
with self.assertRaisesRegex(ValueError, "page-major"):
validate_gdn_mis_backend(LinearAttnKernelBackend.TRITON)
def test_non_gdn_linear_backend_rejects_mis(self):
with self.assertRaisesRegex(ValueError, "does not support multi-item scoring"):
MambaAttnBackendBase.validate_mis_support(SimpleNamespace(enable_mis=True))
GDNAttnBackend.validate_mis_support(SimpleNamespace(enable_mis=True))
def apply_policy(
self,
runner,
@@ -232,6 +250,7 @@ class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True)
metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None)
forward_batch = SimpleNamespace(
multi_item_delimiter_indices=None,
mamba_track_mask=torch.tensor([True]),
mamba_track_indices=torch.tensor([7]),
)
test/registered/unit/layers/test_gdn_mis_metadata.py
+99
View File
@@ -0,0 +1,99 @@
import unittest
from types import SimpleNamespace
import torch
from sglang.srt.layers.attention.linear.gdn_backend import build_gdn_mis_metadata
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestGDNMISMetadata(CustomTestCase):
def test_allows_trailing_attention_padding(self):
forward_batch = SimpleNamespace(
input_ids=torch.empty(8, dtype=torch.int64),
extend_seq_lens_cpu=[5],
extend_prefix_lens_cpu=[0],
multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)],
is_prefill_only=True,
)
metadata = build_gdn_mis_metadata(forward_batch)
torch.testing.assert_close(
torch.cat([metadata.query_token_indices, metadata.item_token_indices])
.sort()
.values,
torch.arange(5, dtype=torch.int64),
)
def test_mixed_batch_with_empty_query(self):
forward_batch = SimpleNamespace(
input_ids=torch.empty(16, dtype=torch.int64),
extend_seq_lens_cpu=[9, 7],
extend_prefix_lens_cpu=[0, 0],
multi_item_delimiter_indices=[
torch.tensor([0, 3, 8], dtype=torch.int64),
torch.tensor([4, 6], dtype=torch.int64),
],
is_prefill_only=True,
)
metadata = build_gdn_mis_metadata(forward_batch)
self.assertEqual(metadata.query_seq_lens_cpu, [4])
self.assertEqual(metadata.item_seq_lens_cpu, [3, 5, 1, 2, 1])
torch.testing.assert_close(
metadata.query_token_indices,
torch.tensor([9, 10, 11, 12], dtype=torch.int64),
)
torch.testing.assert_close(
metadata.query_cu_seqlens,
torch.tensor([0, 4], dtype=torch.int32),
)
torch.testing.assert_close(
metadata.query_request_indices,
torch.tensor([1], dtype=torch.int64),
)
torch.testing.assert_close(
metadata.item_token_indices,
torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 13, 14, 15], dtype=torch.int64),
)
torch.testing.assert_close(
metadata.item_cu_seqlens,
torch.tensor([0, 3, 8, 9, 11, 12], dtype=torch.int32),
)
torch.testing.assert_close(
metadata.item_request_indices,
torch.tensor([0, 0, 0, 1, 1], dtype=torch.int64),
)
def test_rejects_sequence_lengths_beyond_input(self):
forward_batch = SimpleNamespace(
input_ids=torch.empty(4, dtype=torch.int64),
extend_seq_lens_cpu=[5],
extend_prefix_lens_cpu=[0],
multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)],
is_prefill_only=True,
)
with self.assertRaisesRegex(ValueError, "exceed the input tokens"):
build_gdn_mis_metadata(forward_batch)
def test_rejects_missing_final_delimiter(self):
forward_batch = SimpleNamespace(
input_ids=torch.empty(5, dtype=torch.int64),
extend_seq_lens_cpu=[5],
extend_prefix_lens_cpu=[0],
multi_item_delimiter_indices=[torch.tensor([2, 3], dtype=torch.int64)],
is_prefill_only=True,
)
with self.assertRaisesRegex(ValueError, "final delimiter"):
build_gdn_mis_metadata(forward_batch)
if __name__ == "__main__":
unittest.main()