From a63efd9056b33a3d4a32dfba6262fac1d62b959a Mon Sep 17 00:00:00 2001 From: paulzhang-tm Date: Thu, 10 Sep 2026 18:13:17 -0400 Subject: [PATCH] [Spec] Support large MTP batches in short-convolution metadata (#38558) --- .../models/inkling_common/kernels/sconv.py | 17 ++-- .../ops/mamba/test_sconv_extend_metadata.py | 95 ++++++++++++++++++- 2 files changed, 101 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/models/inkling_common/kernels/sconv.py b/python/sglang/srt/models/inkling_common/kernels/sconv.py index 84b2f7658..b4fe7f30c 100644 --- a/python/sglang/srt/models/inkling_common/kernels/sconv.py +++ b/python/sglang/srt/models/inkling_common/kernels/sconv.py @@ -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, diff --git a/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py index 20cc96984..8a1fccfb6 100644 --- a/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py +++ b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py @@ -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,