chore: import upstream snapshot with attribution

tests/python/op/test_batch_spec_verify.py

@@ -0,0 +1,163 @@
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+
+from mlc_llm.op.batch_spec_verify import batch_spec_verify
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("nbatch", [32, 64])
+@pytest.mark.parametrize("vocab", [3, 32, 64, 32000, 33, 65, 32001, 128000])
+@pytest.mark.parametrize("plist", [[0.5, 0.5], [1, 0], [0, 1]])
+def test_batch_spec_verify(nbatch, vocab, plist):
+    def numpy_reference(
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    ):
+        nbatch = token_tree_parent_ptr.shape[0]
+        for b in range(nbatch):
+            parent_ptr = token_tree_parent_ptr[b]
+            child_ptr = token_tree_first_child[parent_ptr]
+            while child_ptr != -1:
+                child_token = draft_tokens[child_ptr]
+                p_child = model_probs[parent_ptr, child_token]
+                q_child = draft_probs[child_ptr, child_token]
+                uniform_sample = uniform_samples[child_ptr]
+                if p_child / q_child >= uniform_sample:
+                    parent_ptr = child_ptr
+                    child_ptr = token_tree_first_child[child_ptr]
+                else:
+                    model_probs[parent_ptr, :] = np.maximum(
+                        model_probs[parent_ptr, :] - draft_probs[child_ptr, :], 0.0
+                    )
+                    psum = np.sum(model_probs[parent_ptr, :])
+                    model_probs[parent_ptr, :] /= psum
+                    child_ptr = token_tree_next_sibling[child_ptr]
+            token_tree_parent_ptr[b] = parent_ptr
+
+    np.random.seed(0)
+
+    def gen_chain(num_nodes, base):
+        token_tree_first_child = list()
+        token_tree_next_sibling = list()
+        for i in range(num_nodes):
+            token_tree_first_child.append(base + i + 1 if i + 1 < num_nodes else -1)
+            token_tree_next_sibling.append(-1)
+        return token_tree_first_child, token_tree_next_sibling, base, base + 1
+
+    def gen_full_binary_tree(height, base):
+        token_tree_first_child = list()
+        token_tree_next_sibling = list()
+        num_nodes = 2**height - 1
+        for i in range(num_nodes):
+            token_tree_first_child.append(base + i * 2 + 1 if i * 2 + 1 < num_nodes else -1)
+            token_tree_next_sibling.append(base + i * 2 + 2 if i * 2 + 2 < num_nodes else -1)
+        return token_tree_first_child, token_tree_next_sibling, base, base + 1
+
+    ### Inputs
+    num_nodes = 0
+    token_tree_first_child = list()
+    token_tree_next_sibling = list()
+    token_tree_parent_ptr = list()
+
+    for _ in range(nbatch):
+        choice = np.random.choice(2, 1, p=plist)
+        if choice == 0:
+            nodes_batch = np.random.randint(3, 32)
+            res = gen_chain(nodes_batch, num_nodes)
+            num_nodes += nodes_batch
+        else:
+            height = np.random.randint(3, 5)
+            res = gen_full_binary_tree(height, num_nodes)
+            num_nodes += 2**height - 1
+        token_tree_first_child.extend(res[0])
+        token_tree_next_sibling.extend(res[1])
+        token_tree_parent_ptr.append(res[2])
+
+    token_tree_first_child = np.array(token_tree_first_child).astype("int32")
+    token_tree_next_sibling = np.array(token_tree_next_sibling).astype("int32")
+    token_tree_parent_ptr = np.array(token_tree_parent_ptr).astype("int32")
+
+    draft_probs = np.random.rand(num_nodes, vocab).astype("float32")
+    draft_probs /= np.sum(draft_probs, axis=1, keepdims=True)
+    draft_tokens = np.random.randint(0, vocab, num_nodes).astype("int32")
+    model_probs = np.random.rand(num_nodes, vocab).astype("float32")
+    model_probs /= np.sum(model_probs, axis=1, keepdims=True)
+    uniform_samples = np.random.rand(num_nodes).astype("float32")
+
+    ### TVM Inputs
+    dev = tvm.cuda(0)
+    draft_probs_tvm = tvm.runtime.tensor(draft_probs, dev)
+    draft_tokens_tvm = tvm.runtime.tensor(draft_tokens, dev)
+    model_probs_tvm = tvm.runtime.tensor(model_probs, dev)
+    token_tree_first_child_tvm = tvm.runtime.tensor(token_tree_first_child, dev)
+    token_tree_next_sibling_tvm = tvm.runtime.tensor(token_tree_next_sibling, dev)
+    uniform_samples_tvm = tvm.runtime.tensor(uniform_samples, dev)
+    token_tree_parent_ptr_tvm = tvm.runtime.tensor(token_tree_parent_ptr, dev)
+
+    # print("draft_probs", draft_probs)
+    # print("draft_tokens", draft_tokens)
+    # print("model_probs", model_probs)
+    # print("token_tree_first_child", token_tree_first_child)
+    # print("token_tree_next_sibling", token_tree_next_sibling)
+    # print("uniform_samples", uniform_samples)
+    # print("token_tree_parent_ptr", token_tree_parent_ptr)
+
+    ### Numpy reference
+    numpy_reference(
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    )
+    # print("model_probs", model_probs)
+    # print("token_tree_parent_ptr", token_tree_parent_ptr)
+
+    ### TVM
+    kernel = batch_spec_verify(vocab)
+    mod = tvm.build(kernel, target="cuda")
+    mod(
+        draft_probs_tvm,
+        draft_tokens_tvm,
+        model_probs_tvm,
+        token_tree_first_child_tvm,
+        token_tree_next_sibling_tvm,
+        uniform_samples_tvm,
+        token_tree_parent_ptr_tvm,
+    )
+    # print("model_probs", model_probs_tvm.asnumpy())
+    # print("token_tree_parent_ptr", token_tree_parent_ptr_tvm.asnumpy())
+
+    tvm.testing.assert_allclose(model_probs, model_probs_tvm.asnumpy())
+    tvm.testing.assert_allclose(
+        token_tree_parent_ptr, token_tree_parent_ptr_tvm.asnumpy(), rtol=0, atol=0
+    )
+
+    time_evaluator = mod.time_evaluator(mod.entry_name, dev, number=10, repeat=3)
+    print(f"batch_size: {nbatch}, vocab_size: {vocab}, tree_structure: {plist}")
+    print(
+        time_evaluator(
+            draft_probs_tvm,
+            draft_tokens_tvm,
+            model_probs_tvm,
+            token_tree_first_child_tvm,
+            token_tree_next_sibling_tvm,
+            uniform_samples_tvm,
+            token_tree_parent_ptr_tvm,
+        )
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_mrope.py

@@ -0,0 +1,251 @@
+import numpy as np
+import pytest
+
+tvm = pytest.importorskip("tvm")
+from tvm import relax  # noqa: E402
+from tvm.relax.frontend import nn  # noqa: E402
+from tvm.relax.frontend.nn import spec  # noqa: E402
+from tvm.runtime import tensor as tvm_tensor  # noqa: E402
+
+from mlc_llm.op import (  # noqa: E402
+    MultimodalRotaryEmbedding,
+    VisionPositionMetadata,
+    apply_multimodal_rotary_pos_emb,
+    get_mrope_position_ids,
+)
+
+
+def _numpy_rotate_half(x: np.ndarray) -> np.ndarray:
+    x1, x2 = np.split(x, 2, axis=-1)
+    return np.concatenate([-x2, x1], axis=-1)
+
+
+def _numpy_apply_mrope(
+    q: np.ndarray,
+    k: np.ndarray,
+    position_ids: np.ndarray,
+    theta: float,
+    mrope_section: tuple[int, ...],
+) -> tuple[np.ndarray, np.ndarray]:
+    if position_ids.ndim != 3:
+        raise ValueError(f"position_ids must be rank-3, got shape {position_ids.shape}")
+    if position_ids.shape[0] == 3:
+        position_ids = np.transpose(position_ids, (1, 2, 0))
+    elif position_ids.shape[-1] != 3:
+        raise ValueError(
+            "position_ids must have shape (batch, seq, 3) or (3, batch, seq), "
+            f"got {position_ids.shape}"
+        )
+
+    head_dim = q.shape[-1]
+    inv_freq = 1.0 / (theta ** (np.arange(0, head_dim, 2, dtype=np.float32) / float(head_dim)))
+    pos = np.transpose(position_ids, (2, 0, 1))
+    inv = inv_freq.reshape(1, 1, -1, 1).astype(np.float32)
+    inv = np.broadcast_to(inv, (3, pos.shape[1], inv_freq.size, 1))
+    pos = pos.reshape(3, pos.shape[1], 1, pos.shape[2]).astype(np.float32)
+    freqs = np.matmul(inv, pos)
+    freqs = np.transpose(freqs, (0, 1, 3, 2))
+    emb = np.concatenate([freqs, freqs], axis=-1)
+    cos = np.cos(emb)
+    sin = np.sin(emb)
+    split_sizes = list(mrope_section) * 2
+    split_points = np.cumsum(split_sizes)[:-1]
+    cos_chunks = np.split(cos, split_points, axis=-1)
+    sin_chunks = np.split(sin, split_points, axis=-1)
+    cos = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(cos_chunks)], axis=-1)
+    sin = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(sin_chunks)], axis=-1)
+    cos = np.expand_dims(cos, axis=2)
+    sin = np.expand_dims(sin, axis=2)
+    q_out = q * cos + _numpy_rotate_half(q) * sin
+    k_out = k * cos + _numpy_rotate_half(k) * sin
+    return q_out, k_out
+
+
+def _evaluate_tensor(expr):
+    mod = tvm.IRModule.from_expr(expr)
+    target = tvm.target.Target("llvm")
+    ex = tvm.relax.build(mod, target)
+    vm = tvm.relax.VirtualMachine(ex, tvm.cpu())
+    return vm["main"]().numpy()
+
+
+def _run_mlc_mrope(
+    q_np: np.ndarray,
+    k_np: np.ndarray,
+    position_ids_np: np.ndarray,
+    theta: float,
+    mrope_section: tuple[int, ...],
+) -> tuple[np.ndarray, np.ndarray]:
+    class RopeModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.rotary = MultimodalRotaryEmbedding(q_np.shape[-1], theta, mrope_section)
+
+        def forward(
+            self,
+            q: nn.Tensor,
+            k: nn.Tensor,
+            pos: nn.Tensor,
+        ):
+            """Run MRoPE on test tensors and return rotated query/key outputs."""
+            cos, sin = self.rotary(q, pos)
+            return apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section)
+
+    module = RopeModule()
+    mod, _, _ = module.export_tvm(
+        spec={
+            "forward": {
+                "q": spec.Tensor(q_np.shape, "float32"),
+                "k": spec.Tensor(k_np.shape, "float32"),
+                "pos": spec.Tensor(position_ids_np.shape, "int64"),
+            }
+        },
+        allow_extern=True,
+    )
+    target = tvm.target.Target("llvm")
+    exec_mod = relax.build(mod, target=target)
+    vm = relax.VirtualMachine(exec_mod, tvm.cpu())
+    device = tvm.cpu()
+    q_nd = tvm_tensor(q_np.astype("float32"), device=device)
+    k_nd = tvm_tensor(k_np.astype("float32"), device=device)
+    pos_nd = tvm_tensor(position_ids_np.astype("int64"), device=device)
+    out_q, out_k = vm["forward"](q_nd, k_nd, pos_nd)
+    return out_q.numpy(), out_k.numpy()
+
+
+def test_apply_mrope_matches_numpy_reference():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    batch, seq_len, heads, head_dim = 1, 4, 2, 12
+    rng = np.random.default_rng(0)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    position_ids = np.zeros((batch, seq_len, 3), dtype=np.int64)
+    position_ids[0, :, 0] = np.arange(seq_len)
+    position_ids[0, :, 1] = np.arange(seq_len) * 2
+    position_ids[0, :, 2] = np.arange(seq_len) * 3
+
+    mlc_q, mlc_k = _run_mlc_mrope(q_np, k_np, position_ids, theta, mrope_section)
+    ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k, ref_k, rtol=1e-5, atol=1e-5)
+
+
+def test_get_mrope_position_ids_text_only():
+    input_ids = np.array([[1, 2, 3, 0, 0]], dtype=np.int64)
+    attention_mask = np.array([[1, 1, 1, 0, 0]], dtype=np.int64)
+    meta = VisionPositionMetadata(
+        vision_start_token_id=1000,
+        image_token_id=1001,
+        video_token_id=1002,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    position_ids, deltas = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    expected = attention_mask.cumsum(axis=-1) - 1
+    expected = np.where(attention_mask == 0, 1, expected)
+    expected = np.expand_dims(expected, axis=0).repeat(3, axis=0)
+    np.testing.assert_array_equal(position_ids, expected)
+    np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
+
+
+def test_get_mrope_position_ids_single_image_block():
+    meta = VisionPositionMetadata(
+        vision_start_token_id=5000,
+        image_token_id=5001,
+        video_token_id=6000,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    input_ids = np.array(
+        [[11, 12, 5000, 5001, 21, 22, 23, 24, 31, 32]],
+        dtype=np.int64,
+    )
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
+    position_ids, deltas = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    expected = np.array(
+        [
+            [0, 1, 2, 3, 3, 3, 3, 5, 6, 7],
+            [0, 1, 2, 3, 3, 4, 4, 5, 6, 7],
+            [0, 1, 2, 3, 4, 3, 4, 5, 6, 7],
+        ],
+        dtype=np.int64,
+    ).reshape(3, 1, -1)
+    np.testing.assert_array_equal(position_ids, expected)
+    np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
+
+
+def test_apply_mrope_accepts_3_batch_seq_layout():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    batch, seq_len, heads, head_dim = 1, 4, 2, 12
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+
+    position_ids_bsc = np.zeros((batch, seq_len, 3), dtype=np.int64)
+    position_ids_bsc[0, :, 0] = np.arange(seq_len)
+    position_ids_bsc[0, :, 1] = np.arange(seq_len) * 2
+    position_ids_bsc[0, :, 2] = np.arange(seq_len) * 3
+    position_ids_3bs = np.transpose(position_ids_bsc, (2, 0, 1))
+
+    mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+    mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
+    ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q_bsc, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_bsc, ref_k, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_q_3bs, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_3bs, ref_k, rtol=1e-5, atol=1e-5)
+
+
+def test_get_mrope_position_ids_output_is_directly_usable():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    meta = VisionPositionMetadata(
+        vision_start_token_id=7000,
+        image_token_id=7001,
+        video_token_id=7002,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    input_ids = np.array([[11, 12, 7000, 7001, 21, 22, 23, 24, 31, 32]], dtype=np.int64)
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
+    position_ids_3bs, _ = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    position_ids_bsc = np.transpose(position_ids_3bs, (1, 2, 0))
+
+    batch, seq_len = input_ids.shape
+    heads, head_dim = 2, 12
+    rng = np.random.default_rng(2)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+
+    mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
+    mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q_3bs, mlc_q_bsc, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_3bs, mlc_k_bsc, rtol=1e-5, atol=1e-5)

tests/python/op/test_top_p_pivot.py

@@ -0,0 +1,86 @@
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+
+from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
+
+# mypy: disable-error-code="var-annotated"
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("batch_size", [32, 64])
+@pytest.mark.parametrize("vocab", [3, 32, 64, 128])
+def test_top_p_renorm(batch_size, vocab):
+    top_p = 0.95
+    init_pivots_np = np.array([1 - top_p, 0.02, 0.01]).astype(np.float32)
+    top_p_np = np.array([top_p]).astype(np.float32)
+
+    p_np = np.random.exponential(3, size=(batch_size, vocab)).astype(np.float32)
+    p_np /= np.sum(p_np, axis=-1, keepdims=True)
+    final_pivot_np = np.zeros(batch_size).astype(np.float32)
+    final_lsum_np = np.zeros(batch_size).astype(np.float32)
+
+    dev = tvm.cuda(0)
+    var_prob = tvm.runtime.tensor(p_np, dev)
+    var_init_pivots = tvm.runtime.tensor(init_pivots_np, dev)
+    top_p_global = tvm.runtime.tensor(top_p_np, dev)
+    var_final_pivot = tvm.runtime.tensor(final_pivot_np, dev)
+    var_final_lsum = tvm.runtime.tensor(final_lsum_np, dev)
+
+    kernel = top_p_pivot(init_pivots_np.shape[0])
+    mod = tvm.build(kernel, target="cuda")
+    mod(var_prob, top_p_global, var_init_pivots, var_final_pivot, var_final_lsum)
+
+    final_pivot = var_final_pivot.asnumpy()
+    final_lsum = var_final_lsum.asnumpy()
+
+    renorm_np = p_np.copy()
+    var_renorm = tvm.runtime.tensor(renorm_np, dev)
+
+    kernel_renorm = top_p_renorm()
+    mod_renorm = tvm.build(kernel_renorm, target="cuda")
+    mod_renorm(var_prob, var_final_pivot, var_final_lsum, var_renorm)
+
+    renorm = var_renorm.asnumpy()
+
+    def verify_pivot(probs: np.ndarray, pivot: float, lsum: float, renorm: np.ndarray):
+        sorted_probs = np.sort(probs, axis=-1)[::-1]
+        num_larger_than_pivot = np.sum(sorted_probs >= pivot)
+        filtered_sorted_probs = sorted_probs[:num_larger_than_pivot]
+        min_larger_than_pivot = min(filtered_sorted_probs)
+
+        sum_larger_than_pivot = np.sum(np.where(sorted_probs >= pivot, sorted_probs, 0))
+        sum_larger_than_pivot_exclude_min = np.sum(
+            np.where(filtered_sorted_probs != min_larger_than_pivot, filtered_sorted_probs, 0)
+        )
+
+        probs[probs < pivot] = 0
+        renorm_prob = probs / np.sum(probs, axis=-1, keepdims=True)
+        try:
+            assert sum_larger_than_pivot >= top_p
+            assert sum_larger_than_pivot_exclude_min < top_p
+            assert abs(lsum - sum_larger_than_pivot) < 1e-6
+            assert np.allclose(renorm, renorm_prob, atol=1e-6, rtol=1e-6)
+        except AssertionError:
+            print("Failed")
+            print("probs:", repr(probs))
+            print("pivot:", pivot)
+            print("sorted_probs:", sorted_probs)
+            print("num_larger_than_pivot:", num_larger_than_pivot)
+            print("filtered_sorted_probs:", filtered_sorted_probs)
+            print("min_larger_than_pivot:", min_larger_than_pivot)
+            print("sum_larger_than_pivot:", sum_larger_than_pivot)
+            print("sum_larger_than_pivot_exclude_min:", sum_larger_than_pivot_exclude_min)
+            print("renom_prob:", renorm_prob)
+            print("renorm:", renorm)
+            raise
+
+    for i in range(batch_size):
+        verify_pivot(p_np[i], final_pivot[i], final_lsum[i], renorm[i])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_tree_attn.py

@@ -0,0 +1,243 @@
+import math
+
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+from tvm.relax.frontend.nn.llm import tree_attn
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("nbatch", [1, 4, 32])
+@pytest.mark.parametrize("h_q", [8, 16])
+@pytest.mark.parametrize("h_kv", [4, 8])
+@pytest.mark.parametrize("d", [128])
+@pytest.mark.parametrize("rotary_mode", [0, 1])
+def test_tree_attn(nbatch, h_q, h_kv, d, rotary_mode):
+    np.random.seed(0)
+    np.set_printoptions(linewidth=10000)
+
+    def gen_chain(num_nodes):
+        mask = np.tril(np.ones((num_nodes, num_nodes)))
+        return num_nodes, list(mask.flatten()), np.arange(num_nodes)
+
+    def gen_full_binary_tree(height):
+        mask = list()
+        pos = list()
+        num_nodes = 2**height - 1
+        for i in range(num_nodes):
+            if i == 0:
+                mask_0 = [0] * num_nodes
+                mask_0[0] = 1
+                mask.append(mask_0)
+                pos.append(0)
+            else:
+                mask_i = mask[(i + 1) // 2 - 1].copy()
+                mask_i[i] = 1
+                mask.append(mask_i)
+                pos.append(pos[(i + 1) // 2 - 1] + 1)
+        return num_nodes, list(np.array(mask).flatten()), pos
+
+    ### Inputs
+    num_nodes = 0
+    m_list = list()
+    mn_list = list()
+    mask_list = list()
+    q_pos_list = list()
+
+    mn_list.append(0)
+
+    for _ in range(nbatch):
+        choice = np.random.choice(2, 1, p=[1, 0])
+        if choice == 0:
+            nodes_batch = np.random.randint(3, 32)
+            res = gen_chain(nodes_batch)
+            num_nodes += nodes_batch
+        else:
+            height = np.random.randint(2, 6)
+            res = gen_full_binary_tree(height)
+            num_nodes += 2**height - 1
+        m_list.append(res[0])
+        mn_list.append(res[0] ** 2)
+        mask_list.extend(res[1])
+        q_pos_list.extend(res[2])
+
+    qkv_indptr = np.array(np.cumsum([0, *m_list])).astype(np.int32)
+    m_list = np.array(m_list).astype(np.int32)
+    mn_list = np.array(mn_list).astype(np.int32)
+    mn_list = np.cumsum(mn_list).astype(np.int32)
+    mask_list = np.array(mask_list).astype(np.int32)
+    q_pos_list = np.array(q_pos_list).astype(np.int32)
+
+    # print("qkv_indptr:", qkv_indptr)
+    # print("m_list:", m_list)
+    # print("mn_list:", mn_list)
+    # for num_nodes, base in zip(m_list, mn_list):
+    #     print("num_nodes:", num_nodes)
+    #     print("indptr:", base)
+    #     print(
+    #         "mask:",
+    #         mask_list[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes),
+    #     )
+    #     print("q_pos:", q_pos_list[base : base + num_nodes])
+
+    q = np.random.rand(num_nodes, h_q, d).astype(np.float16)
+    q_indptr = qkv_indptr
+    k = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
+    v = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
+    kv_indptr = qkv_indptr
+    q_rope_position = q_pos_list
+    m_arr = m_list
+    mn_indptr = mn_list
+    mask = mask_list
+    output = np.zeros((num_nodes, h_q, d), dtype=np.float16)
+    lse = np.zeros((num_nodes, h_q), dtype=np.float32)
+    rotary_scale = 1.0
+    rotary_theta = 10000.0
+    attn_score_scaling_factor = 1.0
+
+    ### TVM Inputs
+    dev = tvm.cuda(0)
+    q_tvm = tvm.runtime.tensor(q, dev)
+    q_indptr_tvm = tvm.runtime.tensor(q_indptr, dev)
+    k_tvm = tvm.runtime.tensor(k, dev)
+    v_tvm = tvm.runtime.tensor(v, dev)
+    kv_indptr_tvm = tvm.runtime.tensor(kv_indptr, dev)
+    q_rope_position_tvm = tvm.runtime.tensor(q_rope_position, dev)
+    # m_arr_tvm = tvm.runtime.tensor(m_arr, dev)
+    mn_indptr_tvm = tvm.runtime.tensor(mn_indptr, dev)
+    mask_tvm = tvm.runtime.tensor(mask, dev)
+    output_tvm = tvm.runtime.tensor(output, dev)
+    lse_tvm = tvm.runtime.tensor(lse, dev)
+
+    target = tvm.target.Target("cuda")
+    kernel = tree_attn(h_kv=h_kv, h_q=h_q, d=d, dtype="float16", rope_scaling={}, target=target)
+    mod = tvm.build(kernel, target=target)
+    mod(
+        q_tvm,
+        q_indptr_tvm,
+        k_tvm,
+        v_tvm,
+        kv_indptr_tvm,
+        q_rope_position_tvm,
+        # m_arr_tvm,
+        mn_indptr_tvm,
+        mask_tvm,
+        output_tvm,
+        lse_tvm,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        nbatch,
+    )
+
+    ### Numpy reference
+    def numpy_reference(
+        q,
+        q_indptr,
+        k,
+        v,
+        kv_indptr,
+        q_rope_position,
+        m_arr,
+        mn_indptr,
+        mask,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        output_tvm,
+    ):
+        def rope_freq(s, d, d_range, theta, dtype):
+            freq = s / math.pow(theta, (d * 2 % d_range) / float(d_range))
+            cos_freq = np.cos(freq).astype(dtype)
+            sin_freq = np.sin(freq).astype(dtype)
+            return cos_freq, sin_freq
+
+        def rope(buffer, offset, rotary_dim, theta, scale, dtype):
+            result = buffer.copy()
+            for pos, h, d in np.ndindex(buffer.shape):
+                cos_freq, sin_freq = rope_freq(offset[pos] * scale, d, rotary_dim, theta, dtype)
+                cos = cos_freq * buffer[pos, h, d]
+                sin = sin_freq * (
+                    -buffer[pos, h, d + rotary_dim // 2]
+                    if d < rotary_dim // 2
+                    else buffer[pos, h, d - rotary_dim // 2]
+                )
+                result[pos, h, d] = cos + sin
+            return result
+
+        for i in range(len(m_arr)):
+            num_nodes = m_arr[i]
+            base = mn_indptr[i]
+            q_base = q_indptr[i]
+            kv_base = kv_indptr[i]
+            q_pos = q_rope_position[q_base : q_base + num_nodes]  # (num_nodes,)
+            q_i = q[q_base : q_base + num_nodes]  # (num_nodes, h_q, d)
+            k_i = k[kv_base : kv_base + num_nodes]  # (num_nodes, h_kv, d)
+            v_i = v[kv_base : kv_base + num_nodes]  # (num_nodes, h_kv, d)
+            mask_i = mask[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes)
+
+            if rotary_mode == 1:
+                q_i = rope(q_i, q_pos, d, rotary_theta, rotary_scale, q_i.dtype)
+                k_i = rope(k_i, q_pos, d, rotary_theta, rotary_scale, k_i.dtype)
+
+            # group attention
+            # q: (num_nodes, h_q, d)
+            # k: (num_nodes, h_kv, d)
+            # v: (num_nodes, h_kv, d)
+            group_size = h_q // h_kv
+            q_reshape = q_i.transpose(1, 0, 2)  # (h_q, num_nodes, d)
+            k_reshape = k_i.transpose(1, 2, 0)  # (h_kv, d, num_nodes)
+            v_reshape = v_i.transpose(1, 0, 2)  # (h_kv, num_nodes, d)
+            # expand k_reshape
+            k_reshape = k_reshape.reshape(h_kv, 1, d, num_nodes)
+            k_reshape = np.repeat(k_reshape, group_size, axis=1)
+            k_reshape = k_reshape.reshape(h_q, d, num_nodes)
+            # expand v_reshape
+            v_reshape = v_reshape.reshape(h_kv, 1, num_nodes, d)
+            v_reshape = np.repeat(v_reshape, group_size, axis=1)
+            v_reshape = v_reshape.reshape(h_q, num_nodes, d)
+            # print("q_reshape:", q_reshape.shape)
+            # print("k_reshape:", k_reshape.shape)
+            # print("v_reshape:", v_reshape.shape)
+
+            # qk: (h_q, num_nodes, num_nodes)
+            qk = np.matmul(q_reshape, k_reshape) * attn_score_scaling_factor / math.sqrt(float(d))
+            # softmax(qk, axis=-1), numerical stability
+            qk[:, mask_i == 0] = -np.inf
+            qk_max = np.max(qk, axis=-1, keepdims=True)
+            qk = np.exp(qk - qk_max)
+            qk = qk / np.sum(qk, axis=-1, keepdims=True)
+
+            # attention
+            output_i = np.matmul(qk, v_reshape).transpose(1, 0, 2)  # (num_nodes, h_q, d)
+            # print(output_i)
+
+            tvm.testing.assert_allclose(
+                output_i, output_tvm[q_base : q_base + num_nodes], rtol=1e-3, atol=1e-3
+            )
+
+    numpy_reference(
+        q,
+        q_indptr,
+        k,
+        v,
+        kv_indptr,
+        q_rope_position,
+        m_arr,
+        mn_indptr,
+        mask,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        output_tvm.numpy(),
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_two_stage_softmax.py

@@ -0,0 +1,51 @@
+import numpy as np
+import pytest
+import scipy.special
+import tvm
+from tvm.s_tir import dlight
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+def test_two_stage_softmax():
+    from mlc_llm.compiler_pass.rewrite_softmax import _get_lse_and_softmax_func
+
+    chunk_size = 4096
+    target = tvm.target.Target("cuda")
+    f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(target, chunk_size)
+    mod = tvm.IRModule({"chunk_lse": f_chunk_lse, "softmax_with_chunked_lse": f_softmax_with_lse})
+    with target:
+        mod = dlight.ApplyDefaultSchedule(dlight.gpu.GeneralReduction())(mod)
+
+    runtime_mod = tvm.build(mod, target=target)
+    device = tvm.cuda()
+
+    num_runs = 5
+    vocab_size = 128256
+    for batch_size in [1, 2, 4, 8, 16, 32, 64, 128]:
+        for _ in range(num_runs):
+            x_np = np.random.uniform(low=-10, high=10, size=(batch_size, vocab_size)).astype(
+                "float32"
+            )
+            y_np = scipy.special.softmax(x_np, axis=-1)
+
+            x_nd = tvm.runtime.tensor(x_np, device=device)
+            r_nd = tvm.runtime.empty(
+                (batch_size, (vocab_size + chunk_size - 1) // chunk_size),
+                x_np.dtype,
+                device=device,
+            )
+            y_nd = tvm.runtime.empty(x_np.shape, x_np.dtype, device=device)
+
+            runtime_mod["chunk_lse"](x_nd, r_nd)
+            runtime_mod["softmax_with_chunked_lse"](x_nd, r_nd, y_nd)
+
+            y_nd_arr = y_nd.numpy()
+            np.testing.assert_allclose(y_nd_arr, y_np, atol=1e-6, rtol=1e-6)
+
+        print(f"pass batch size {batch_size}")
+
+
+if __name__ == "__main__":
+    test_two_stage_softmax()