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

[Spec] Support large MTP batches in short-convolution metadata (#38558)

Browse Source
This commit is contained in:
paulzhang-tm
2026-09-10 15:13:17 -07:00
committed by GitHub
parent fd7743e0e1
commit a63efd9056
2 changed files with 101 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/inkling_common/kernels/sconv.py
+10 -7
View File
@@ -369,9 +369,9 @@ HIS_PREFIX = 1 # extend_prefix_lens > 0
HIS_SEQ_MINUS_EXT = 2 # (seq_lens[:B] - extend_seq_lens) > 0 (draft_extend_v2 capture)
HIS_ONES = 3 # target_verify: always has initial state
# The single-tile local cumsum bounds the fused path; larger batches fall back
# to the unfused op sequence.
_FUSED_EXTEND_MAX_B = 1023
# The single-tile local cumsum bounds variable-length extend; larger batches
# fall back to the unfused op sequence. Uniform-length verify needs no cumsum.
_FUSED_EXTEND_MAX_B = 2047
@triton.jit
@@ -475,7 +475,8 @@ def fused_extend_sconv_metadata(
owning layer).
Returns ``(query_start_loc, has_initial_state, SconvExtendMetadata)`` with
tensors bit-identical to the unfused path, or None when the shape falls
outside the fused kernel's single-tile bound (caller runs unfused).
outside the variable-length kernel's single-tile bound (caller runs unfused).
Uniform-length target verification does not have that batch-size bound.
``his_mode`` selects the has_initial_state source: HIS_ZEROS (boundary-KV
draft extend), HIS_PREFIX (``his_src`` = extend_prefix_lens), HIS_SEQ_MINUS_EXT
@@ -483,10 +484,10 @@ def fused_extend_sconv_metadata(
``extend_seq_lens`` unused). Pass ``out`` to write into preallocated (e.g.
cuda-graph-static) destinations instead of fresh allocations.
"""
if B > _FUSED_EXTEND_MAX_B or not cache_indices.is_cuda:
is_verify = his_mode == HIS_ONES
if (not is_verify and B > _FUSED_EXTEND_MAX_B) or not cache_indices.is_cuda:
return None
assert cache_indices.shape[0] >= B and cache_indices.stride(0) == 1
is_verify = his_mode == HIS_ONES
if is_verify:
assert draft_token_num is not None
else:
@@ -499,7 +500,9 @@ def fused_extend_sconv_metadata(
safe_idx = dst["safe_idx"]
cu = dst["cu"]
si = dst["si"]
BLOCK_T = 256
# Keep the variable-length si comparison tile at most 256 * 1024 elements
# when the cumsum tile grows to 2048 requests (e.g. MTP batches of 1264).
BLOCK_T = 128 if not is_verify and B > 1023 else 256
dummy = cache_indices # never dereferenced thanks to masks/constexpr
_fused_extend_metadata_kernel[(1 + triton.cdiv(T, BLOCK_T),)](
cache_indices,
test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py
+91 -4
View File
@@ -16,6 +16,7 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
HIS_SEQ_MINUS_EXT,
HIS_ZEROS,
PAD_SLOT_ID,
SconvMetadataOut,
fused_extend_sconv_metadata,
precompute_helion_extend_metadata,
)
@@ -27,8 +28,8 @@ register_cuda_ci(est_time=40, stage="nightly", runner_config="1-gpu-large")
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
# cross si tiles (BLOCK_T=256) and the single-tile B bound
BATCH_SIZES = [1, 2, 7, 64, 257, 1023]
# Cross si tiles (BLOCK_T=128/256) and the single-tile B bound.
BATCH_SIZES = [1, 2, 7, 64, 257, 1023, 1024, 1264, 2047]
EXTEND_CASES = get_ci_test_range(
[
(b, his_mode, lens_dtype)
@@ -43,10 +44,17 @@ EXTEND_CASES = get_ci_test_range(
(64, HIS_ZEROS, torch.int64),
(257, HIS_PREFIX, torch.int32),
(1023, HIS_SEQ_MINUS_EXT, torch.int64),
(1024, HIS_ZEROS, torch.int32),
(1264, HIS_PREFIX, torch.int64),
(2047, HIS_SEQ_MINUS_EXT, torch.int64),
],
)
VERIFY_CASES = get_ci_test_range(
[(b, draft_token_num) for draft_token_num in (1, 9) for b in BATCH_SIZES],
[
(b, draft_token_num)
for draft_token_num in (1, 3, 9)
for b in BATCH_SIZES + [2048]
],
[
(1, 1),
(2, 9),
@@ -54,6 +62,23 @@ VERIFY_CASES = get_ci_test_range(
(64, 9),
(257, 1),
(1023, 9),
(1024, 3),
(1264, 3),
(2047, 9),
(2048, 3),
],
)
GRAPH_REPLAY_CASES = get_ci_test_range(
[
(b, his_mode)
for b in (1024, 1264, 2047)
for his_mode in (HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT, HIS_ONES)
],
[
(1024, HIS_ZEROS),
(1264, HIS_PREFIX),
(2047, HIS_SEQ_MINUS_EXT),
(1264, HIS_ONES),
],
)
@@ -162,6 +187,68 @@ def test_verify_matches_unfused(b, draft_token_num):
_assert_equal(got, ref)
@requires_cuda
@pytest.mark.parametrize("b,his_mode", GRAPH_REPLAY_CASES)
def test_large_extend_cuda_graph_replay(b, his_mode):
"""Large MTP batches refresh static metadata after lengths/PAD slots change."""
draft_token_num = 3
cache_indices = _cache_indices(b, torch.int64)
lens = torch.full((b,), draft_token_num, dtype=torch.int64, device="cuda")
his_src = lens + 1
def reference():
if his_mode == HIS_ONES:
return _ref_verify(b, draft_token_num, cache_indices)
return _ref_extend(
b, lens, his_mode, his_src, cache_indices, b * draft_token_num
)
ref = reference()
out = SconvMetadataOut(
query_start_loc=torch.empty_like(ref[0]),
has_initial_state=torch.empty_like(ref[1]),
**{key: torch.empty_like(value) for key, value in ref[2].items()},
)
def refresh():
return fused_extend_sconv_metadata(
B=b,
T=b * draft_token_num,
cache_indices=cache_indices,
his_mode=his_mode,
draft_token_num=draft_token_num,
extend_seq_lens=lens,
his_src=his_src,
out=out,
)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
assert refresh() is not None # Compile before graph capture.
torch.cuda.current_stream().wait_stream(stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=stream):
got = refresh()
assert got is not None
assert got[0].data_ptr() == out["query_start_loc"].data_ptr()
assert got[1].data_ptr() == out["has_initial_state"].data_ptr()
for key, value in got[2].items():
assert value.data_ptr() == out[key].data_ptr()
for offset in (0, 1):
cache_indices.copy_(torch.arange(b, dtype=torch.int64, device="cuda"))
cache_indices[offset::2] = PAD_SLOT_ID
lens.fill_(draft_token_num)
lens[offset::3] = 0
his_src.copy_(lens + 1)
his_src[offset::2] = 0
ref = reference()
graph.replay()
torch.cuda.synchronize()
_assert_equal(got, ref)
@requires_cuda
def test_cu_not_spanning_T():
"""Dummy capture sequences: cu stops short of T; trailing si rows clamp to
@@ -186,7 +273,7 @@ def test_cu_not_spanning_T():
@requires_cuda
def test_fallback_past_batch_bound():
b = 1024 # > _FUSED_EXTEND_MAX_B
b = 2048 # > _FUSED_EXTEND_MAX_B
lens = torch.ones(b, dtype=torch.int64, device="cuda")
got = fused_extend_sconv_metadata(
B=b,