Add Inkling model support (#31681)

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Qiaolin Yu <qiaolin.yu@radixark.ai>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Aurick Qiao <aurick@thinkingmachines.ai>
Co-authored-by: Joseph <jk@thinkingmachines.ai>

python/sglang/jit_kernel/tests/test_moe_preprocess.py

@@ -0,0 +1,93 @@
+"""fused_moe_preprocess must be bit-identical to the torch.sort-based path,
+and the grouped GEMM must produce identical results under both block_size_m
+configs (the block schedule and kernel config are chosen together).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import (
+    SMALL_M_BLOCK_SIZE_M,
+    compute_grouped_gemm_metadata,
+    fused_moe_preprocess,
+    get_src2dst,
+    grouped_gemm_triton,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+E = 256
+TOPK = 6
+
+
+def _reference(topk_ids_flat: torch.Tensor):
+    reorder_topk_ids, reorder_ids = torch.sort(
+        topk_ids_flat.to(torch.int16), stable=True
+    )
+    src2dst = get_src2dst(reorder_ids)
+    meta = compute_grouped_gemm_metadata(
+        reorder_topk_ids, E, block_size_m=SMALL_M_BLOCK_SIZE_M
+    )
+    return (src2dst, *meta, reorder_topk_ids)
+
+
+def _ids(tokens: int, seed: int, skew: bool = False) -> torch.Tensor:
+    torch.manual_seed(seed)
+    if skew:  # all tokens on few experts (stresses multi-block experts)
+        return torch.randint(0, 3, (tokens * TOPK,), dtype=torch.int32, device="cuda")
+    return (
+        torch.stack([torch.randperm(E, device="cuda")[:TOPK] for _ in range(tokens)])
+        .view(-1)
+        .to(torch.int32)
+    )
+
+
+@requires_cuda
+@pytest.mark.parametrize("tokens", [1, 2, 7, 32, 64, 170, 341])  # n = 6*T <= 2048
+@pytest.mark.parametrize("skew", [False, True])
+def test_matches_sort_path(tokens: int, skew: bool):
+    ids = _ids(tokens, seed=tokens, skew=skew)
+    ref = _reference(ids)
+    got = fused_moe_preprocess(ids, E)
+    names = [
+        "src2dst",
+        "num_tokens_per_expert",
+        "expert_token_offs",
+        "expert_block_offs",
+        "expert_block_schedule",
+        "reorder_topk_ids",
+    ]
+    for tag, g, r in zip(names, got, ref):
+        assert g.shape == r.shape, (tag, g.shape, r.shape)
+        assert torch.equal(g.long(), r.long()), (
+            tag,
+            g[: min(16, g.numel())],
+            r[: min(16, r.numel())],
+        )
+
+
+@requires_cuda
+@pytest.mark.parametrize("tokens", [1, 16, 64])
+def test_grouped_gemm_small_config_matches(tokens: int):
+    """GEMM output must be identical whichever (block_size_m, config) runs."""
+    torch.manual_seed(tokens)
+    ids = _ids(tokens, seed=tokens)
+    m, k, n = tokens * TOPK, 768, 1024
+    a = (torch.randn(m, k, device="cuda") * 0.05).to(torch.bfloat16)
+    b = (torch.randn(E, n, k, device="cuda") * 0.02).to(torch.bfloat16)
+
+    sorted_ids, _ = torch.sort(ids.to(torch.int16), stable=True)
+    meta128 = compute_grouped_gemm_metadata(sorted_ids, E)
+    out128 = grouped_gemm_triton(a, b, E, *meta128)
+
+    pre = fused_moe_preprocess(ids, E)
+    out16 = grouped_gemm_triton(a, b, E, *pre[1:5], block_size_m=SMALL_M_BLOCK_SIZE_M)
+    # both are fp32-accumulated bf16 tensor-core dots; BLOCK_K differs so
+    # accumulation grouping may differ by a few ulp
+    torch.testing.assert_close(out16.float(), out128.float(), atol=1e-3, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))

python/sglang/jit_kernel/tests/test_sconv_decode_metadata.py

@@ -0,0 +1,70 @@
+"""fused_decode_sconv_metadata must be bit-identical to the unfused prep.
+
+The unfused reference is the exact op sequence `_prepare_decode_sconv_metadata`
+used to launch: two arange calls + ones + precompute_helion_decode_metadata
+(!= PAD, &, clamp, long, arange x2).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.models.inkling_common.kernels.sconv import (
+    PAD_SLOT_ID,
+    fused_decode_sconv_metadata,
+    precompute_helion_decode_metadata,
+)
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
+
+# cross the BLOCK=1024 grid boundary and hit odd sizes
+BATCH_SIZES = [1, 2, 3, 17, 64, 160, 257, 1023, 1024, 1025]
+
+
+def _reference(B: int, cache_indices: torch.Tensor):
+    device = cache_indices.device
+    query_start_loc = torch.arange(B + 1, dtype=torch.int32, device=device)
+    has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
+    precomputed = precompute_helion_decode_metadata(
+        B=B, W=4, cache_indices=cache_indices, has_initial_state=has_initial_state
+    )
+    return query_start_loc, has_initial_state, precomputed
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64])
+def test_matches_unfused(b: int, idx_dtype: torch.dtype):
+    torch.manual_seed(b)
+    cache_indices = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
+    # sprinkle PAD slots (cudagraph padding lanes)
+    pad = torch.rand(b, device="cuda") < 0.25
+    cache_indices[pad] = PAD_SLOT_ID
+
+    ref_qsl, ref_his, ref_meta = _reference(b, cache_indices)
+    qsl, his, meta = fused_decode_sconv_metadata(B=b, cache_indices=cache_indices)
+
+    for tag, got, ref in (
+        ("query_start_loc", qsl, ref_qsl),
+        ("has_initial_state", his, ref_his),
+        ("cache_mask", meta["cache_mask"], ref_meta["cache_mask"]),
+        ("safe_idx", meta["safe_idx"], ref_meta["safe_idx"]),
+        ("cu", meta["cu"], ref_meta["cu"]),
+        ("si", meta["si"], ref_meta["si"]),
+    ):
+        assert got.dtype == ref.dtype, (tag, got.dtype, ref.dtype)
+        assert got.shape == ref.shape, (tag, got.shape, ref.shape)
+        assert torch.equal(got, ref), tag
+
+
+@requires_cuda
+def test_all_pad():
+    cache_indices = torch.full((8,), PAD_SLOT_ID, dtype=torch.int32, device="cuda")
+    _, _, meta = fused_decode_sconv_metadata(B=8, cache_indices=cache_indices)
+    assert not meta["cache_mask"].any()
+    assert (meta["safe_idx"] == 0).all()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))

python/sglang/jit_kernel/tests/test_sconv_extend_metadata.py

@@ -0,0 +1,174 @@
+"""fused_extend_sconv_metadata must be bit-identical to the unfused prep.
+
+The unfused reference is the exact op sequence _prepare_extend_common_metadata
++ precompute_helion_extend_metadata used to launch: zeros + cumsum + slice-copy
+(or arange + ones for verify) + the has_initial_state compare, then != PAD, &,
+clamp, long, to(int64), arange, searchsorted, clamp, to(int32).
+"""
+
+import pytest
+import torch
+
+from sglang.srt.models.inkling_common.kernels.sconv import (
+    HIS_ONES,
+    HIS_PREFIX,
+    HIS_SEQ_MINUS_EXT,
+    HIS_ZEROS,
+    PAD_SLOT_ID,
+    fused_extend_sconv_metadata,
+    precompute_helion_extend_metadata,
+)
+
+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]
+
+
+def _ref_extend(B, extend_seq_lens, his_mode, his_src, cache_indices, T):
+    device = cache_indices.device
+    query_start_loc = torch.zeros(B + 1, dtype=torch.int32, device=device)
+    query_start_loc[1:] = extend_seq_lens.cumsum(dim=0)
+    if his_mode == HIS_ZEROS:
+        has_initial_state = torch.zeros(B, dtype=torch.bool, device=device)
+    elif his_mode == HIS_PREFIX:
+        has_initial_state = his_src > 0
+    else:  # HIS_SEQ_MINUS_EXT
+        has_initial_state = (his_src[:B] - extend_seq_lens) > 0
+    meta = precompute_helion_extend_metadata(
+        B=B,
+        T=T,
+        W=4,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        query_start_loc=query_start_loc,
+    )
+    return query_start_loc, has_initial_state, meta
+
+
+def _ref_verify(B, draft_token_num, cache_indices):
+    device = cache_indices.device
+    query_start_loc = torch.arange(
+        0, (B + 1) * draft_token_num, draft_token_num, dtype=torch.int32, device=device
+    )
+    has_initial_state = torch.ones(B, dtype=torch.bool, device=device)
+    meta = precompute_helion_extend_metadata(
+        B=B,
+        T=B * draft_token_num,
+        W=4,
+        cache_indices=cache_indices,
+        has_initial_state=has_initial_state,
+        query_start_loc=query_start_loc,
+    )
+    return query_start_loc, has_initial_state, meta
+
+
+def _assert_equal(got, ref):
+    for tag, g, r in (
+        ("query_start_loc", got[0], ref[0]),
+        ("has_initial_state", got[1], ref[1]),
+        ("cache_mask", got[2]["cache_mask"], ref[2]["cache_mask"]),
+        ("safe_idx", got[2]["safe_idx"], ref[2]["safe_idx"]),
+        ("cu", got[2]["cu"], ref[2]["cu"]),
+        ("si", got[2]["si"], ref[2]["si"]),
+    ):
+        assert g.dtype == r.dtype, (tag, g.dtype, r.dtype)
+        assert g.shape == r.shape, (tag, g.shape, r.shape)
+        assert torch.equal(g, r), tag
+
+
+def _cache_indices(b, idx_dtype):
+    ci = torch.randint(0, 4096, (b,), dtype=idx_dtype, device="cuda")
+    pad = torch.rand(b, device="cuda") < 0.25
+    ci[pad] = PAD_SLOT_ID
+    return ci
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("his_mode", [HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT])
+@pytest.mark.parametrize("lens_dtype", [torch.int32, torch.int64])
+def test_extend_matches_unfused(b, his_mode, lens_dtype):
+    torch.manual_seed(b * 10 + his_mode)
+    lens = torch.randint(0, 33, (b,), dtype=lens_dtype, device="cuda")
+    lens[torch.rand(b, device="cuda") < 0.2] = 0  # zero-length sequences
+    T = int(lens.sum().item())
+    cache_indices = _cache_indices(b, torch.int32)
+    if his_mode == HIS_PREFIX:
+        his_src = torch.randint(0, 3, (b,), dtype=lens_dtype, device="cuda")
+    elif his_mode == HIS_SEQ_MINUS_EXT:
+        his_src = lens + torch.randint(0, 2, (b,), dtype=lens_dtype, device="cuda")
+    else:
+        his_src = None
+
+    ref = _ref_extend(b, lens, his_mode, his_src, cache_indices, T)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=T,
+        cache_indices=cache_indices,
+        his_mode=his_mode,
+        extend_seq_lens=lens,
+        his_src=his_src,
+    )
+    assert got is not None
+    _assert_equal(got, ref)
+
+
+@requires_cuda
+@pytest.mark.parametrize("b", BATCH_SIZES)
+@pytest.mark.parametrize("draft_token_num", [1, 9])
+def test_verify_matches_unfused(b, draft_token_num):
+    torch.manual_seed(b)
+    cache_indices = _cache_indices(b, torch.int64)
+    ref = _ref_verify(b, draft_token_num, cache_indices)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=b * draft_token_num,
+        cache_indices=cache_indices,
+        his_mode=HIS_ONES,
+        draft_token_num=draft_token_num,
+    )
+    assert got is not None
+    _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
+    B-1 exactly like the reference's searchsorted + clamp."""
+    b = 5
+    lens = torch.tensor([3, 0, 4, 0, 2], dtype=torch.int64, device="cuda")
+    T = int(lens.sum().item()) + 17
+    cache_indices = _cache_indices(b, torch.int32)
+    seq_lens = lens + 1
+    ref = _ref_extend(b, lens, HIS_SEQ_MINUS_EXT, seq_lens, cache_indices, T)
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=T,
+        cache_indices=cache_indices,
+        his_mode=HIS_SEQ_MINUS_EXT,
+        extend_seq_lens=lens,
+        his_src=seq_lens,
+    )
+    assert got is not None
+    _assert_equal(got, ref)
+
+
+@requires_cuda
+def test_fallback_past_batch_bound():
+    b = 1024  # > _FUSED_EXTEND_MAX_B
+    lens = torch.ones(b, dtype=torch.int64, device="cuda")
+    got = fused_extend_sconv_metadata(
+        B=b,
+        T=b,
+        cache_indices=_cache_indices(b, torch.int32),
+        his_mode=HIS_ZEROS,
+        extend_seq_lens=lens,
+    )
+    assert got is None
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-x"]))