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

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
+# The single-tile local cumsum bounds variable-length extend; larger batches
-# to the unfused op sequence.
+# fall back to the unfused op sequence. Uniform-length verify needs no cumsum.
-_FUSED_EXTEND_MAX_B = 1023
+_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

@@ -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
+# Cross si tiles (BLOCK_T=128/256) and the single-tile B bound.
-BATCH_SIZES = [1, 2, 7, 64, 257, 1023]
+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,