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

[minimax-m3] Split 1/4: sparse attention ops + JIT kernels + config foundation (#28712)

Browse Source
This commit is contained in:
Xinyuan Tong
2026-06-22 13:10:43 -07:00
committed by GitHub
parent b5e4e289b1
commit 7c23d2255a
51 changed files with 11157 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py
+121
View File
@@ -0,0 +1,121 @@
"""Benchmark: MiniMax-M3 single-stage radix-select decode topk (JIT CUDA) vs the
2-stage split-K Triton baseline (_topk_index_partial_kernel + _topk_index_merge_kernel).
Both consume the decode score tensor [num_heads, batch, max_seqblock] and produce
topk_idx [num_heads, batch, topk]. The JIT kernel is one launch with no
intermediate buffers; the baseline is two launches with split-K partials.
"""
import torch
import triton
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx import (
_topk_index_merge_kernel,
_topk_index_partial_kernel,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=8, suite="base-b-kernel-benchmark-1-gpu-large")
BLOCK_SIZE = 128
TOPK = 16
NUM_HEADS = 1 # per-rank index heads at TP>=4
def _triton_2stage(score, seq_lens):
num_q_heads, batch_size, max_seqblock = score.shape
TOPK_TARGET_GRID = 64
MAX_NUM_TOPK_CHUNKS = 16
t = max(
1,
min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch_size * num_q_heads)),
)
nchunks = 1 << (t.bit_length() - 1)
bt = triton.next_power_of_2(TOPK)
chunk_blocks = (max_seqblock + nchunks - 1) // nchunks
out = torch.empty(
(num_q_heads, batch_size, TOPK), device=score.device, dtype=torch.int32
)
tsp = torch.empty(
nchunks, num_q_heads, batch_size, bt, dtype=torch.float32, device=score.device
)
tip = torch.empty(
nchunks, num_q_heads, batch_size, bt, dtype=torch.int32, device=score.device
)
_topk_index_partial_kernel[(batch_size, num_q_heads, nchunks)](
score,
tsp,
tip,
seq_lens,
BLOCK_SIZE,
TOPK,
chunk_blocks,
score.stride(0),
score.stride(1),
score.stride(2),
tsp.stride(0),
tsp.stride(1),
tsp.stride(2),
tsp.stride(3),
tip.stride(0),
tip.stride(1),
tip.stride(2),
tip.stride(3),
)
_topk_index_merge_kernel[(batch_size, num_q_heads)](
tsp,
tip,
out,
seq_lens,
BLOCK_SIZE,
TOPK,
tsp.stride(0),
tsp.stride(1),
tsp.stride(2),
tsp.stride(3),
tip.stride(0),
tip.stride(1),
tip.stride(2),
tip.stride(3),
out.stride(0),
out.stride(1),
out.stride(2),
NUM_TOPK_CHUNKS=nchunks,
)
return out
def _jit(score, seq_lens):
return minimax_decode_topk(score, seq_lens, BLOCK_SIZE, TOPK)
FN_MAP = {"jit": _jit, "triton_2stage": _triton_2stage}
@marker.parametrize("ctx", [4096, 32768, 131072, 524288], [4096, 524288])
@marker.parametrize("batch", [1, 4, 16, 64, 256], [1, 64])
@marker.benchmark("impl", ["jit", "triton_2stage"])
def benchmark(ctx: int, batch: int, impl: str):
max_seqblock = (524288 + BLOCK_SIZE - 1) // BLOCK_SIZE
nb = min((ctx + BLOCK_SIZE - 1) // BLOCK_SIZE, max_seqblock)
score = torch.full(
(NUM_HEADS, batch, max_seqblock),
float("-inf"),
dtype=torch.float32,
device="cuda",
)
score[:, :, :nb] = torch.randn(NUM_HEADS, batch, nb, device="cuda") * 5.0
score[:, :, nb - 1] = 1e29 # forced local block
seq_lens = torch.full((batch,), ctx, device="cuda", dtype=torch.int32)
return marker.do_bench(
FN_MAP[impl],
input_args=(score, seq_lens),
graph_clone_args=(0, 1), # both read-only inputs
memory_args=(score,),
)
if __name__ == "__main__":
benchmark.run()
test/registered/jit/benchmark/minimax/bench_minimax_qknorm_rope.py
+146
View File
@@ -0,0 +1,146 @@
"""Benchmark: fused MiniMax-M3 Gemma-RMSNorm + partial NeoX RoPE (1 in-place
launch) vs the unfused path (GemmaRMSNorm(q) + GemmaRMSNorm(k) + rotary_emb,
3 launches + intermediates). Main attention branch, per-rank TP8 shape (nq=8, nk=1).
"""
import torch
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.minimax_qknorm_rope import (
minimax_qknorm_rope,
minimax_qknorm_rope_grouped,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
HEAD_DIM, ROTARY_DIM, BASE, EPS, MAXPOS = 128, 64, 5_000_000, 1e-6, 131072
NQ, NK = 64, 4
def _cache():
inv_freq = 1.0 / (
BASE
** (
torch.arange(0, ROTARY_DIM, 2, dtype=torch.float, device="cuda")
/ ROTARY_DIM
)
)
t = torch.arange(MAXPOS, dtype=torch.float, device="cuda")
freqs = torch.outer(t, inv_freq)
return torch.cat([freqs.cos(), freqs.sin()], dim=-1).contiguous()
def _unfused(qkv, wq, wk, cache, positions):
# GemmaRMSNorm (1+w) on q,k head-wise + partial neox rope, in plain torch
# (representative of the separate norm + rope launches).
T = qkv.shape[0]
q, k, v = qkv.split([NQ * HEAD_DIM, NK * HEAD_DIM, NK * HEAD_DIM], dim=-1)
def norm(x, w, nh):
x = x.reshape(T, nh, HEAD_DIM).float()
var = x.pow(2).mean(-1, keepdim=True)
return (x * torch.rsqrt(var + EPS) * (1.0 + w.float())).to(torch.bfloat16)
qn, kn = norm(q, wq, NQ), norm(k, wk, NK)
cs = cache.index_select(0, positions).float()
cos, sin = cs[:, None, :32], cs[:, None, 32:]
def rope(x):
x1, x2 = x[..., :32].float(), x[..., 32:64].float()
o1 = x1 * cos - x2 * sin
o2 = x2 * cos + x1 * sin
return torch.cat([o1, o2, x[..., 64:].float()], dim=-1).to(torch.bfloat16)
return rope(qn), rope(kn)
def _fused(qkv, wq, wk, cache, positions):
minimax_qknorm_rope(qkv, wq, wk, cache, positions, NQ, NK, NK, EPS)
return qkv
FN_MAP = {"fused": _fused, "unfused_torch": _unfused}
@marker.parametrize("T", [1, 16, 64, 256, 1024, 8192], [64, 1024])
@marker.benchmark("impl", ["fused", "unfused_torch"])
def benchmark(T: int, impl: str):
cache = _cache()
wq = (torch.randn(HEAD_DIM, device="cuda") * 0.1).to(torch.bfloat16)
wk = (torch.randn(HEAD_DIM, device="cuda") * 0.1).to(torch.bfloat16)
qkv = torch.randn(T, (NQ + 2 * NK) * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
positions = torch.randint(0, MAXPOS, (T,), device="cuda", dtype=torch.int64)
return marker.do_bench(
FN_MAP[impl],
input_args=(qkv, wq, wk, cache, positions),
graph_clone_args=(0,),
memory_args=None,
)
# --- Combined main + index single launch (the fused qkv+index_qkv GEMM path) ---
# Per-rank TP8 sparse shape: main q=8/k=1/v=1 + index idx_q=1/idx_k=1 (value
# disabled). One grouped launch (q, k, idx_q, idx_k) vs two separate launches.
C_NQ, C_NKV, C_NIQ = 8, 1, 1
C_OFF_Q = 0
C_OFF_K = C_NQ
C_OFF_V = C_NQ + C_NKV
C_OFF_IQ = C_NQ + 2 * C_NKV
C_OFF_IK = C_OFF_IQ + C_NIQ
C_TOTAL_HEADS = C_OFF_IK + 1
def _combined_one(args):
qkv, wq, wk, wiq, wik, cache, positions = args
minimax_qknorm_rope_grouped(
qkv,
[
(wq, C_OFF_Q, C_NQ),
(wk, C_OFF_K, C_NKV),
(wiq, C_OFF_IQ, C_NIQ),
(wik, C_OFF_IK, 1),
],
cache,
positions,
EPS,
)
return qkv
def _combined_two(args):
# Two launches over the same buffer: main (q,k) then index (idx_q, idx_k).
qkv, wq, wk, wiq, wik, cache, positions = args
minimax_qknorm_rope_grouped(
qkv, [(wq, C_OFF_Q, C_NQ), (wk, C_OFF_K, C_NKV)], cache, positions, EPS
)
minimax_qknorm_rope_grouped(
qkv, [(wiq, C_OFF_IQ, C_NIQ), (wik, C_OFF_IK, 1)], cache, positions, EPS
)
return qkv
C_FN_MAP = {"combined_one_launch": _combined_one, "two_launches": _combined_two}
@marker.parametrize("T", [1, 16, 64, 256, 1024, 8192], [64, 1024])
@marker.benchmark("impl", ["combined_one_launch", "two_launches"])
def benchmark_combined(T: int, impl: str):
cache = _cache()
ws = [
(torch.randn(HEAD_DIM, device="cuda") * 0.1).to(torch.bfloat16)
for _ in range(4)
]
qkv = torch.randn(T, C_TOTAL_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
positions = torch.randint(0, MAXPOS, (T,), device="cuda", dtype=torch.int64)
return marker.do_bench(
C_FN_MAP[impl],
input_args=((qkv, *ws, cache, positions),),
graph_clone_args=(0,),
memory_args=None,
)
if __name__ == "__main__":
benchmark.run()
benchmark_combined.run()
test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py
+67
View File
@@ -0,0 +1,67 @@
"""Benchmark: fused MiniMax-M3 KV + index cache store (1 launch) vs the separate
per-buffer index_put_ stores (main K, main V, index K, optional index V)."""
import torch
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
HEAD_DIM = 128
NUM_KV_HEADS = 1
HAS_V = False
N = 1 << 20
DTYPE = torch.bfloat16
def _separate(k, v, kc, vc, ik, ikc, loc, **_):
kc[loc] = k
vc[loc] = v
ikc[loc] = ik
def _fused(k, v, kc, vc, ik, ikc, loc, *, num_kv_heads, head_bytes):
store_kv_index(
k,
v,
kc,
vc,
ik,
ikc,
None,
None,
loc,
num_kv_heads=num_kv_heads,
head_bytes=head_bytes,
)
FN_MAP = {"fused": _fused, "separate": _separate}
@marker.parametrize("T", [16, 64, 256, 1024, 4096, 16384], [256, 4096])
@marker.benchmark("impl", ["fused", "separate"])
def benchmark(T: int, impl: str):
k = torch.randn(T, NUM_KV_HEADS * HEAD_DIM, dtype=DTYPE, device="cuda")
v = torch.randn_like(k)
ik = torch.randn(T, HEAD_DIM, dtype=DTYPE, device="cuda")
kc = torch.zeros(N, NUM_KV_HEADS * HEAD_DIM, dtype=DTYPE, device="cuda")
vc = torch.zeros_like(kc)
ikc = torch.zeros(N, HEAD_DIM, dtype=DTYPE, device="cuda")
loc = torch.randperm(N, device="cuda")[:T]
extra_kwargs = dict(num_kv_heads=NUM_KV_HEADS, head_bytes=HEAD_DIM * DTYPE.itemsize)
return marker.do_bench(
FN_MAP[impl],
input_args=(k, v, kc, vc, ik, ikc, loc),
input_kwargs=extra_kwargs if impl == "fused" else {},
# Read inputs cloned per iter; caches are write targets (kept hot).
graph_clone_args=(0, 1, 4, 6),
memory_args=(k, v, ik, loc),
memory_output=(k, v, ik),
)
if __name__ == "__main__":
benchmark.run()
test/registered/jit/minimax/test_minimax_decode_topk.py
+144
View File
@@ -0,0 +1,144 @@
"""Correctness tests for the MiniMax-M3 single-stage radix-select decode topk.
The kernel selects, per (head, batch) row, the indices of the ``topk`` largest
block scores among the row's first ``num_blocks = ceil(seq_len / block_size)``
entries, front-packing valid block ids and ``-1``-padding the tail. This mirrors
the consumer ``_gqa_share_sparse_decode_kernel`` contract.
"""
import pytest
import torch
from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-b200")
def _ref(score, seq_lens, block_size, topk):
H, B, S = score.shape
out = torch.full((H, B, topk), -1, dtype=torch.int32, device=score.device)
for h in range(H):
for b in range(B):
sl = int(seq_lens[b])
nb = min((sl + block_size - 1) // block_size, S)
if nb <= topk:
for i in range(nb):
out[h, b, i] = i
continue
keff = min(topk, nb)
_, idx = torch.topk(score[h, b, :nb], keff)
out[h, b, :keff] = idx.to(torch.int32)
return out
def _selected_scores_sorted(score, out):
"""Per-row sorted-desc multiset of the scores the kernel selected (tie-robust)."""
H, B, _ = score.shape
rows = []
for h in range(H):
for b in range(B):
sel = out[h, b]
sel = sel[sel >= 0].long()
assert len(sel.unique()) == len(sel), f"duplicate idx h{h} b{b}: {sel}"
rows.append(torch.sort(score[h, b, sel], descending=True).values)
return rows
def _check_contract(out, seq_lens, block_size, topk, S):
H, B, _ = out.shape
for h in range(H):
for b in range(B):
o = out[h, b]
nvalid = int((o >= 0).sum())
# valid entries are front-packed, -1 after
assert torch.all(o[:nvalid] >= 0) and torch.all(o[nvalid:] == -1), o
nb = min((int(seq_lens[b]) + block_size - 1) // block_size, S)
assert nvalid == min(topk, nb)
assert torch.all(o[:nvalid] < nb)
@pytest.mark.parametrize("dtype_sl", [torch.int32, torch.int64])
@pytest.mark.parametrize("H", [1, 2])
@pytest.mark.parametrize("B", [1, 5, 32])
@pytest.mark.parametrize("topk", [16, 32, 64])
@pytest.mark.parametrize("max_ctx", [4096, 131072, 524288])
def test_decode_topk_distinct(dtype_sl, H, B, topk, max_ctx):
torch.manual_seed(1234)
block_size = 128
S = (max_ctx + block_size - 1) // block_size
# distinct scores per row -> exact index-set match against torch.topk
score = torch.empty(H, B, S, dtype=torch.float32, device="cuda")
for h in range(H):
for b in range(B):
score[h, b] = torch.randperm(S, device="cuda").float() + torch.rand(
1, device="cuda"
)
seq_lens = torch.randint(1, max_ctx + 1, (B,), device="cuda", dtype=dtype_sl)
out = minimax_decode_topk(score, seq_lens, block_size, topk)
ref = _ref(score, seq_lens, block_size, topk)
_check_contract(out, seq_lens, block_size, topk, S)
# exact index-set equality (distinct scores)
for h in range(H):
for b in range(B):
assert set(out[h, b][out[h, b] >= 0].tolist()) == set(
ref[h, b][ref[h, b] >= 0].tolist()
)
@pytest.mark.parametrize("kind", ["ties", "negative", "neg_inf_padding", "all_equal"])
def test_decode_topk_adversarial(kind):
torch.manual_seed(7)
block_size = 128
H, B, S, topk = 1, 6, 1024, 16
if kind == "ties":
score = torch.randint(0, 4, (H, B, S), device="cuda").float()
elif kind == "negative":
score = -torch.rand(H, B, S, device="cuda") * 1000 - 1.0
elif kind == "neg_inf_padding":
score = torch.randn(H, B, S, device="cuda")
score[:, :, ::7] = float("-inf") # scattered -inf in valid range
else: # all_equal
score = torch.full((H, B, S), 3.14, dtype=torch.float32, device="cuda")
seq_lens = torch.randint(1, S * block_size, (B,), device="cuda", dtype=torch.int32)
out = minimax_decode_topk(score, seq_lens, block_size, topk)
ref = _ref(score, seq_lens, block_size, topk)
_check_contract(out, seq_lens, block_size, topk, S)
# tie-robust: the multiset of selected scores must match torch.topk's
for a, b in zip(
_selected_scores_sorted(score, out), _selected_scores_sorted(score, ref)
):
torch.testing.assert_close(a, b, rtol=0, atol=0)
@pytest.mark.parametrize("seq_len", [1, 128, 129, 2048, 2049])
def test_decode_topk_small_num_blocks(seq_len):
# num_blocks around / below topk -> naive identity path and boundary.
block_size = 128
H, B, S, topk = 1, 1, 64, 16
score = torch.randn(H, B, S, dtype=torch.float32, device="cuda")
seq_lens = torch.tensor([seq_len], device="cuda", dtype=torch.int32)
out = minimax_decode_topk(score, seq_lens, block_size, topk)
_check_contract(out, seq_lens, block_size, topk, S)
nb = min((seq_len + block_size - 1) // block_size, S)
if nb <= topk:
assert out[0, 0, :nb].tolist() == list(range(nb))
def test_decode_topk_out_param():
block_size = 128
H, B, S, topk = 1, 4, 1024, 16
score = torch.randn(H, B, S, dtype=torch.float32, device="cuda")
seq_lens = torch.full((B,), 100000, device="cuda", dtype=torch.int32)
out = torch.empty((H, B, topk), dtype=torch.int32, device="cuda")
res = minimax_decode_topk(score, seq_lens, block_size, topk, out=out)
assert res is out
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/minimax/test_minimax_decode_topk_page_table.py
+178
View File
@@ -0,0 +1,178 @@
"""Fused decode top-k + page-table transform.
`minimax_decode_topk_page_table` selects the top-k blocks (same as the block-id
`minimax_decode_topk`) and emits the per-query paged page table consumed by the
dense backend (trtllm_mha), instead of block ids. This checks the fused output
end-to-end: trtllm decode over the emitted page table matches the custom
`_gqa_share_sparse_decode_kernel` fed the block-id selection from the same score.
Only the TP>=4 case (num_kv_heads == 1) is covered.
"""
import random
import sys
import pytest
import torch
flashinfer = pytest.importorskip("flashinfer")
from sglang.jit_kernel.minimax_decode_topk import (
minimax_decode_topk,
minimax_decode_topk_page_table,
)
from sglang.srt.layers.attention.minimax_sparse_ops.decode.topk_sparse import (
flash_decode_with_gqa_share_sparse,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=25, suite="base-b-kernel-unit-1-gpu-b200")
dev = "cuda"
def _effective_kv_from_selection(ti, seq_lens, block):
# Ground-truth effective KV length: sum of valid tokens over the selected
# blocks (only the final block can be partial), per query.
bs = seq_lens.shape[0]
out = torch.zeros(bs, dtype=torch.int32, device=ti.device)
for b in range(bs):
sl = int(seq_lens[b])
tot = 0
for c in ti[0, b].tolist():
if c < 0:
continue
tot += min(block, sl - c * block)
out[b] = tot
return out
@pytest.mark.parametrize(
"bs,seq_len",
[
(1, 5000),
(2, 300),
(3, 160),
(4, 2048),
(8, 4096),
(16, 8000),
(2, 40000), # num_blocks=313 -> medium radix path
(1, 90000), # num_blocks=704 -> large compaction path
(1, 480000), # num_blocks=3750 -> large path near kMaxNumBlocks
],
)
@pytest.mark.parametrize("nqh", [8, 16])
def test_fused_page_table_matches_custom(bs, seq_len, nqh):
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
if arch_major < 10:
pytest.skip("trtllm-gen decode is Blackwell (sm100)")
nkv, D, block, topk, ps = 1, 128, 128, 16, 64
torch.manual_seed(bs * 31 + seq_len + nqh)
random.seed(bs * 31 + seq_len)
ppr = (seq_len + ps - 1) // ps
npages = bs * ppr + 4
kf = torch.randn(npages * ps, nkv, D, device=dev, dtype=torch.bfloat16) * 0.5
vf = torch.randn(npages * ps, nkv, D, device=dev, dtype=torch.bfloat16) * 0.5
q = torch.randn(bs, nqh, D, device=dev, dtype=torch.bfloat16) * 0.5
r2t = torch.zeros(bs, seq_len, dtype=torch.int32, device=dev)
sl = torch.full((bs,), seq_len, dtype=torch.int32, device=dev)
sid = torch.arange(bs, dtype=torch.int64, device=dev)
idxs = torch.arange(seq_len, device=dev)
for b in range(bs):
r2t[b] = ((b * ppr + idxs // ps) * ps + idxs % ps).int()
nb = (seq_len + block - 1) // block
score = torch.full((1, bs, nb), -float("inf"), device=dev, dtype=torch.float32)
score[0, :, :nb] = torch.randn(bs, nb, device=dev)
score[0, :, nb - 1] = 1e30 # final (local) block always selected
# block-id selection -> custom kernel (reference)
ti = minimax_decode_topk(score, sl, block, topk)
ref = flash_decode_with_gqa_share_sparse(
q, None, kf, vf, r2t, sl, sid, block, ti, sm_scale=D**-0.5
)
# fused page-table + effective KV length -> trtllm (allocated + returned)
pt, cache = minimax_decode_topk_page_table(score, sl, r2t, sid, block, topk, ps)
# the kernel's effective KV length must match the actual block selection
expect_cache = _effective_kv_from_selection(ti, sl, block)
assert torch.equal(cache, expect_cache), f"{cache} != {expect_cache}"
ws = torch.zeros(128 * 1024 * 1024, dtype=torch.int8, device=dev)
kv = (
kf.view(npages, ps, nkv, D).permute(0, 2, 1, 3),
vf.view(npages, ps, nkv, D).permute(0, 2, 1, 3),
)
o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
query=q,
kv_cache=kv,
workspace_buffer=ws,
block_tables=pt,
seq_lens=cache,
max_seq_len=topk * block,
bmm1_scale=D**-0.5,
bmm2_scale=1.0,
)
cos = torch.nn.functional.cosine_similarity(
ref.float().flatten(), o.float().flatten(), dim=0
).item()
assert cos > 0.999, f"cos={cos}"
@pytest.mark.parametrize("seq_len", [300, 5000, 90000])
@pytest.mark.parametrize("bs", [1, 3])
@pytest.mark.parametrize("nkv", [2, 4])
def test_dp_flattened_page_table(nkv, bs, seq_len):
"""DP attention (num_kv_heads>1): each kv head selects its own blocks, flattened
into bs*nkv pseudo-requests (row = b*nkv + h). Validate the flattened page table
+ effective KV length against the per-head block-id selection, including the
head-minor head-encoded page index (base_page*nkv + h, the index into an HND
cache [num_pages, nkv, ps, D] reshaped to [num_pages*nkv, 1, ps, D])."""
D, block, topk, ps = 128, 128, 16, 64
ppb = block // ps
torch.manual_seed(nkv * 131 + bs * 31 + seq_len)
ppr = (seq_len + ps - 1) // ps
max_kv = seq_len # req_to_token width
r2t = torch.zeros(bs, seq_len, dtype=torch.int32, device=dev)
sl = torch.full((bs,), seq_len, dtype=torch.int32, device=dev)
sid = torch.arange(bs, dtype=torch.int64, device=dev)
idxs = torch.arange(seq_len, device=dev)
for b in range(bs):
r2t[b] = ((b * ppr + idxs // ps) * ps + idxs % ps).int()
nb = (seq_len + block - 1) // block
score = torch.full((nkv, bs, nb), -float("inf"), device=dev, dtype=torch.float32)
score[:, :, :nb] = torch.randn(nkv, bs, nb, device=dev)
score[:, :, nb - 1] = 1e30 # final (local) block always selected
# per-head block-id selection is the reference for the flattened page table
ti = minimax_decode_topk(score, sl, block, topk) # [nkv, bs, topk]
pt, cache = minimax_decode_topk_page_table(score, sl, r2t, sid, block, topk, ps)
msp = topk * ppb
assert pt.shape == (bs * nkv, msp) and cache.shape == (bs * nkv,)
r2t_cpu = r2t.cpu()
for b in range(bs):
for h in range(nkv):
blocks = sorted(c for c in ti[h, b].tolist() if c >= 0)
row = b * nkv + h
# effective KV length = sum of valid tokens over selected blocks
exp_kv = sum(min(block, seq_len - c * block) for c in blocks)
assert (
int(cache[row]) == exp_kv
), f"row {row}: {int(cache[row])} != {exp_kv}"
# page table: each block -> ppb pages via req_to_token, head-minor encoded
for e in range(len(blocks) * ppb):
c = blocks[e // ppb]
tok = c * block + (e % ppb) * ps
if tok >= max_kv:
tok = max_kv - 1
exp = int(r2t_cpu[b, tok]) // ps * nkv + h
assert (
int(pt[row, e]) == exp
), f"row {row} e {e}: {int(pt[row,e])} != {exp}"
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/jit/minimax/test_minimax_qknorm_rope.py
+178
View File
@@ -0,0 +1,178 @@
"""Correctness for the fused MiniMax-M3 Gemma-RMSNorm + partial NeoX RoPE kernel.
Verifies the in-place fused kernel reproduces GemmaRMSNorm((1+w)) + partial NeoX
RoPE to the bf16 round-off floor, leaves V untouched, and matches sglang's RoPE
convention (cos|sin cache, neox pairs (i, i+rotary_dim/2)).
"""
import pytest
import torch
from sglang.jit_kernel.minimax_qknorm_rope import (
minimax_qknorm_rope,
minimax_qknorm_rope_grouped,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
dev = "cuda"
HEAD_DIM, ROTARY_DIM, BASE, EPS = 128, 64, 5_000_000, 1e-6
def _build_cache(max_pos):
inv_freq = 1.0 / (
BASE
** (torch.arange(0, ROTARY_DIM, 2, dtype=torch.float, device=dev) / ROTARY_DIM)
) # [32]
t = torch.arange(max_pos, dtype=torch.float, device=dev)
freqs = torch.outer(t, inv_freq) # [max_pos, 32]
return torch.cat([freqs.cos(), freqs.sin()], dim=-1).contiguous() # [max_pos, 64]
def _ref(q, k, wq, wk, cache, positions, nq, nk):
T = q.shape[0]
def norm(x, w, nh):
x = x.reshape(T, nh, HEAD_DIM).float()
var = x.pow(2).mean(-1, keepdim=True)
return x * torch.rsqrt(var + EPS) * (1.0 + w.float())
cs = cache.index_select(0, positions).float()
cos, sin = cs[:, :32], cs[:, 32:]
def rope(x):
x1, x2 = x[..., :32], x[..., 32:64]
o1 = x1 * cos[:, None, :] - x2 * sin[:, None, :]
o2 = x2 * cos[:, None, :] + x1 * sin[:, None, :]
return torch.cat([o1, o2, x[..., 64:]], dim=-1)
return rope(norm(q, wq, nq)).reshape(T, -1), rope(norm(k, wk, nk)).reshape(T, -1)
@pytest.mark.parametrize("nq,nk", [(8, 1), (64, 8), (16, 2)])
@pytest.mark.parametrize("T", [1, 7, 64, 1024, 4096])
@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
def test_fused_qknorm_rope(nq, nk, T, pos_dtype):
torch.manual_seed(T * 131 + nq)
max_pos = 8192
cache = _build_cache(max_pos)
wq = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
wk = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
q = torch.randn(T, nq * HEAD_DIM, dtype=torch.bfloat16, device=dev)
k = torch.randn(T, nk * HEAD_DIM, dtype=torch.bfloat16, device=dev)
v = torch.randn(T, nk * HEAD_DIM, dtype=torch.bfloat16, device=dev)
positions = torch.randint(0, max_pos, (T,), device=dev, dtype=pos_dtype)
qr, kr = _ref(q, k, wq, wk, cache, positions.long(), nq, nk)
qkv = torch.cat([q, k, v], dim=-1).contiguous()
minimax_qknorm_rope(qkv, wq, wk, cache, positions, nq, nk, nk, EPS)
qf, kf, vf = qkv.split([nq * HEAD_DIM, nk * HEAD_DIM, nk * HEAD_DIM], dim=-1)
floor = (qr.bfloat16().float() - qr.float()).abs().max().item()
assert (qf.float() - qr.float()).abs().max().item() <= 2 * floor + 1e-3
assert (kf.float() - kr.float()).abs().max().item() <= 2 * floor + 1e-3
assert (vf.float() - v.float()).abs().max().item() == 0.0 # V untouched
def test_index_branch_shapes():
# idx_q (nq=num_idx_heads, nk=0) and idx_k (nq=1) in-place calls.
torch.manual_seed(0)
max_pos = 4096
cache = _build_cache(max_pos)
w = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
T, num_idx = 33, 4
positions = torch.randint(0, max_pos, (T,), device=dev, dtype=torch.int64)
for nq in (num_idx, 1):
x = torch.randn(T, nq * HEAD_DIM, dtype=torch.bfloat16, device=dev)
ref, _ = _ref(x, x[:, :HEAD_DIM], w, w, cache, positions, nq, 1)
xc = x.clone()
minimax_qknorm_rope(xc, w, w, cache, positions, nq, 0, 0, EPS)
floor = (ref.bfloat16().float() - ref.float()).abs().max().item()
assert (xc.float() - ref.float()).abs().max().item() <= 2 * floor + 1e-3
def _norm_rope_one(x_heads, w, cache, positions):
# x_heads: [T, nh, HEAD_DIM] fp32; returns same shape, GemmaRMSNorm(1+w)+rope.
var = x_heads.pow(2).mean(-1, keepdim=True)
y = x_heads * torch.rsqrt(var + EPS) * (1.0 + w.float())
cs = cache.index_select(0, positions).float()
cos, sin = cs[:, None, :32], cs[:, None, 32:]
x1, x2 = y[..., :32], y[..., 32:64]
o1 = x1 * cos - x2 * sin
o2 = x2 * cos + x1 * sin
return torch.cat([o1, o2, y[..., 64:]], dim=-1)
@pytest.mark.parametrize("nq,nkv,niq", [(8, 1, 1), (8, 1, 4), (16, 2, 2)])
@pytest.mark.parametrize("idx_v", [0, 1])
@pytest.mark.parametrize("T", [1, 7, 64, 1024])
@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
def test_combined_main_index_grouped(nq, nkv, niq, idx_v, T, pos_dtype):
"""Combined main(q,k,v) + index(idx_q,idx_k,[idx_v]) layout in one launch.
Mirrors the fused qkv+index_qkv GEMM output: a uniform [total_heads, 128]
grid where Q / K main heads and index-Q / index-K heads are normed+roped in
one pass and the V / index-V heads are left untouched.
"""
torch.manual_seed(T * 17 + nq * 3 + niq + idx_v)
max_pos = 8192
cache = _build_cache(max_pos)
positions = torch.randint(0, max_pos, (T,), device=dev, dtype=pos_dtype)
# head layout: q(nq) k(nkv) v(nkv) idx_q(niq) idx_k(1) [idx_v(1)]
off_q = 0
off_k = nq
off_v = nq + nkv
off_iq = nq + 2 * nkv
off_ik = off_iq + niq
total_heads = off_ik + 1 + idx_v
wq = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
wk = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
wiq = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
wik = (torch.randn(HEAD_DIM, device=dev) * 0.1).to(torch.bfloat16)
x = torch.randn(T, total_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
ref = x.float().clone()
ref[:, off_q:off_k] = _norm_rope_one(
ref[:, off_q:off_k], wq, cache, positions.long()
)
ref[:, off_k:off_v] = _norm_rope_one(
ref[:, off_k:off_v], wk, cache, positions.long()
)
ref[:, off_iq:off_ik] = _norm_rope_one(
ref[:, off_iq:off_ik], wiq, cache, positions.long()
)
ref[:, off_ik : off_ik + 1] = _norm_rope_one(
ref[:, off_ik : off_ik + 1], wik, cache, positions.long()
)
qkv = x.reshape(T, total_heads * HEAD_DIM).contiguous()
minimax_qknorm_rope_grouped(
qkv,
[(wq, off_q, nq), (wk, off_k, nkv), (wiq, off_iq, niq), (wik, off_ik, 1)],
cache,
positions,
EPS,
)
out = qkv.reshape(T, total_heads, HEAD_DIM)
floor = (ref.bfloat16().float() - ref).abs().max().item()
assert (out.float() - ref).abs().max().item() <= 2 * floor + 1e-3
# V and index-V heads untouched (bit-exact).
assert (
out[:, off_v:off_iq].float() - x[:, off_v:off_iq].float()
).abs().max() == 0.0
if idx_v:
assert (
out[:, off_ik + 1 :].float() - x[:, off_ik + 1 :].float()
).abs().max() == 0.0
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/minimax/test_minimax_store_kv_index.py
+78
View File
@@ -0,0 +1,78 @@
"""Correctness for the fused MiniMax-M3 KV + index cache store kernel.
Verifies the single fused launch writes the main K/V, the index K, and the
optional index V into their pools at out_cache_loc rows exactly as the separate
index_put_ stores would, for both value modes and int32/int64 indices.
"""
import pytest
import torch
from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
dev = "cuda"
HEAD_DIM = 128
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("num_kv_heads", [1, 4, 8])
@pytest.mark.parametrize("has_v", [False, True])
@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("T", [1, 7, 128, 513])
def test_store_kv_index(dtype, num_kv_heads, has_v, idx_dtype, T):
torch.manual_seed(T * 17 + num_kv_heads)
head_bytes = HEAD_DIM * dtype.itemsize
N = 4096
def rnd(*shape):
return torch.randn(*shape, dtype=dtype, device=dev)
k = rnd(T, num_kv_heads * HEAD_DIM)
v = rnd(T, num_kv_heads * HEAD_DIM)
idx_k = rnd(T, HEAD_DIM)
idx_v = rnd(T, HEAD_DIM) if has_v else None
k_cache = torch.zeros(N, num_kv_heads * HEAD_DIM, dtype=dtype, device=dev)
v_cache = torch.zeros_like(k_cache)
idx_k_cache = torch.zeros(N, HEAD_DIM, dtype=dtype, device=dev)
idx_v_cache = torch.zeros_like(idx_k_cache) if has_v else None
loc = torch.randperm(N, device=dev)[:T].to(idx_dtype)
store_kv_index(
k,
v,
k_cache,
v_cache,
idx_k,
idx_k_cache,
idx_v,
idx_v_cache,
loc,
num_kv_heads=num_kv_heads,
head_bytes=head_bytes,
)
ll = loc.long()
k_ref = torch.zeros_like(k_cache)
v_ref = torch.zeros_like(v_cache)
ik_ref = torch.zeros_like(idx_k_cache)
k_ref[ll], v_ref[ll], ik_ref[ll] = k, v, idx_k
assert torch.equal(k_cache, k_ref)
assert torch.equal(v_cache, v_ref)
assert torch.equal(idx_k_cache, ik_ref)
if has_v:
iv_ref = torch.zeros_like(idx_v_cache)
iv_ref[ll] = idx_v
assert torch.equal(idx_v_cache, iv_ref)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/jit/test_minimax_m3_qk_norm_rope.py
+334
View File
@@ -0,0 +1,334 @@
# SPDX-License-Identifier: Apache-2.0
"""Reference tests for MiniMax-M3 fused Q/K Gemma RMSNorm + RoPE."""
import pytest
import torch
from sglang.srt.utils import is_hip
if not is_hip():
pytest.skip(
"MiniMax-M3 fused Q/K norm + RoPE kernel is ROCm-only.",
allow_module_level=True,
)
if not torch.cuda.is_available():
pytest.skip("Requires a GPU.", allow_module_level=True)
from sglang.jit_kernel.minimax_m3.qk_norm_rope import ( # noqa: E402
qk_gemma_rmsnorm_rope,
sparse_qk_index_gemma_rmsnorm_rope,
sparse_qk_index_gemma_rmsnorm_rope_cache,
)
from sglang.test.ci.ci_register import register_amd_ci # noqa: E402
# ROCm-only fused kernel; runs in the AMD jit-kernel unit suite.
register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
DEVICE = "cuda"
EPS = 1e-6
def _gemma_norm_by_head(x: torch.Tensor, weight: torch.Tensor, head_dim: int):
orig_shape = x.shape
orig_dtype = x.dtype
xh = x.view(x.shape[0], -1, head_dim).float()
var = xh.pow(2).mean(dim=-1, keepdim=True)
out = xh * torch.rsqrt(var + EPS) * (1.0 + weight.float())
return out.to(orig_dtype).reshape(orig_shape)
def _apply_rope_ref(
x: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
head_dim: int,
rotary_dim: int,
is_neox_style: bool,
):
orig_shape = x.shape
xh = x.view(x.shape[0], -1, head_dim)
x_rot = xh[..., :rotary_dim].float()
x_pass = xh[..., rotary_dim:]
cos_sin = cos_sin_cache.index_select(0, positions)
cos, sin = cos_sin.chunk(2, dim=-1)
cos = cos[:, None, :].float()
sin = sin[:, None, :].float()
if is_neox_style:
x1, x2 = x_rot.chunk(2, dim=-1)
y_rot = torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
else:
x1 = x_rot[..., ::2]
x2 = x_rot[..., 1::2]
y_rot = torch.stack((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
y_rot = y_rot.flatten(-2)
return torch.cat((y_rot.to(x.dtype), x_pass), dim=-1).reshape(orig_shape)
def _reference(
q,
k,
q_weight,
k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
):
q_norm = _gemma_norm_by_head(q, q_weight, head_dim)
k_norm = _gemma_norm_by_head(k, k_weight, head_dim)
q_ref = _apply_rope_ref(
q_norm, positions, cos_sin_cache, head_dim, rotary_dim, is_neox_style
)
k_ref = _apply_rope_ref(
k_norm, positions, cos_sin_cache, head_dim, rotary_dim, is_neox_style
)
return q_ref, k_ref
def _sparse_reference(
q,
k,
idx_q,
idx_k,
q_weight,
k_weight,
idx_q_weight,
idx_k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
):
q_ref, k_ref = _reference(
q,
k,
q_weight,
k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
)
idx_q_norm = _gemma_norm_by_head(idx_q, idx_q_weight, head_dim)
idx_k_norm = _gemma_norm_by_head(idx_k, idx_k_weight, head_dim)
idx_q_ref = _apply_rope_ref(
idx_q_norm, positions, cos_sin_cache, head_dim, rotary_dim, is_neox_style
)
idx_k_ref = _apply_rope_ref(
idx_k_norm, positions, cos_sin_cache, head_dim, rotary_dim, is_neox_style
)
return q_ref, k_ref, idx_q_ref, idx_k_ref
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("is_neox_style", [True, False])
@pytest.mark.parametrize(
"num_tokens,q_heads,k_heads,head_dim,rotary_dim",
[(1, 16, 1, 128, 64), (17, 16, 1, 128, 64), (64, 4, 1, 128, 64)],
)
@torch.inference_mode()
def test_qk_gemma_rmsnorm_rope_matches_reference(
dtype, is_neox_style, num_tokens, q_heads, k_heads, head_dim, rotary_dim
):
torch.manual_seed(0)
q_dim = q_heads * head_dim
k_dim = k_heads * head_dim
padding_dim = 37
qkv = torch.randn(
num_tokens, q_dim + k_dim + padding_dim, device=DEVICE, dtype=dtype
)
q, k, _ = qkv.split([q_dim, k_dim, padding_dim], dim=-1)
if num_tokens > 1:
assert not q.is_contiguous()
assert not k.is_contiguous()
q_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
k_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
positions = torch.randint(0, 512, (num_tokens,), device=DEVICE, dtype=torch.long)
cos_sin_cache = torch.randn(512, rotary_dim, device=DEVICE, dtype=dtype)
got_q, got_k = qk_gemma_rmsnorm_rope(
q,
k,
q_weight,
k_weight,
positions,
cos_sin_cache,
EPS,
head_dim,
rotary_dim,
is_neox_style,
)
ref_q, ref_k = _reference(
q,
k,
q_weight,
k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
)
torch.testing.assert_close(got_q, ref_q, atol=3e-2, rtol=3e-2)
torch.testing.assert_close(got_k, ref_k, atol=3e-2, rtol=3e-2)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("is_neox_style", [True, False])
@pytest.mark.parametrize(
"num_tokens,q_heads,k_heads,idx_q_heads,head_dim,rotary_dim",
[(1, 16, 1, 16, 128, 64), (19, 16, 1, 16, 128, 64)],
)
@torch.inference_mode()
def test_sparse_qk_index_gemma_rmsnorm_rope_matches_reference(
dtype,
is_neox_style,
num_tokens,
q_heads,
k_heads,
idx_q_heads,
head_dim,
rotary_dim,
):
torch.manual_seed(1)
q = torch.randn(num_tokens, q_heads * head_dim, device=DEVICE, dtype=dtype)
k = torch.randn(num_tokens, k_heads * head_dim, device=DEVICE, dtype=dtype)
idx_q = torch.randn(num_tokens, idx_q_heads * head_dim, device=DEVICE, dtype=dtype)
idx_k = torch.randn(num_tokens, head_dim, device=DEVICE, dtype=dtype)
q_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
k_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
idx_q_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
idx_k_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
positions = torch.randint(0, 512, (num_tokens,), device=DEVICE, dtype=torch.long)
cos_sin_cache = torch.randn(512, rotary_dim, device=DEVICE, dtype=dtype)
got = sparse_qk_index_gemma_rmsnorm_rope(
q,
k,
idx_q,
idx_k,
q_weight,
k_weight,
idx_q_weight,
idx_k_weight,
positions,
cos_sin_cache,
EPS,
head_dim,
rotary_dim,
is_neox_style,
)
ref = _sparse_reference(
q,
k,
idx_q,
idx_k,
q_weight,
k_weight,
idx_q_weight,
idx_k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
)
for got_tensor, ref_tensor in zip(got, ref):
torch.testing.assert_close(got_tensor, ref_tensor, atol=3e-2, rtol=3e-2)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("is_neox_style", [True, False])
@torch.inference_mode()
def test_sparse_qk_index_gemma_rmsnorm_rope_cache_matches_reference(
dtype, is_neox_style
):
torch.manual_seed(2)
num_tokens, q_heads, k_heads, idx_q_heads = 11, 16, 1, 16
head_dim, rotary_dim = 128, 64
q = torch.randn(num_tokens, q_heads * head_dim, device=DEVICE, dtype=dtype)
k = torch.randn(num_tokens, k_heads * head_dim, device=DEVICE, dtype=dtype)
v = torch.randn(num_tokens, k_heads * head_dim, device=DEVICE, dtype=dtype)
idx_q = torch.randn(num_tokens, idx_q_heads * head_dim, device=DEVICE, dtype=dtype)
idx_k = torch.randn(num_tokens, head_dim, device=DEVICE, dtype=dtype)
q_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
k_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
idx_q_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
idx_k_weight = torch.randn(head_dim, device=DEVICE, dtype=torch.float32)
positions = torch.randint(0, 512, (num_tokens,), device=DEVICE, dtype=torch.long)
cos_sin_cache = torch.randn(512, rotary_dim, device=DEVICE, dtype=dtype)
out_cache_loc = torch.randperm(64, device=DEVICE, dtype=torch.int64)[:num_tokens]
k_cache = torch.empty(64, k_heads, head_dim, device=DEVICE, dtype=dtype)
v_cache = torch.empty(64, k_heads, head_dim, device=DEVICE, dtype=dtype)
idx_k_cache = torch.empty(64, 1, head_dim, device=DEVICE, dtype=dtype)
got = sparse_qk_index_gemma_rmsnorm_rope_cache(
q,
k,
v,
idx_q,
idx_k,
k_cache,
v_cache,
idx_k_cache,
out_cache_loc,
q_weight,
k_weight,
idx_q_weight,
idx_k_weight,
positions,
cos_sin_cache,
EPS,
head_dim,
rotary_dim,
is_neox_style,
)
ref = _sparse_reference(
q,
k,
idx_q,
idx_k,
q_weight,
k_weight,
idx_q_weight,
idx_k_weight,
positions,
cos_sin_cache,
head_dim,
rotary_dim,
is_neox_style,
)
for got_tensor, ref_tensor in zip(got, ref):
torch.testing.assert_close(got_tensor, ref_tensor, atol=3e-2, rtol=3e-2)
torch.testing.assert_close(
k_cache.index_select(0, out_cache_loc),
ref[1].view(num_tokens, k_heads, head_dim),
atol=3e-2,
rtol=3e-2,
)
torch.testing.assert_close(
v_cache.index_select(0, out_cache_loc),
v.view(num_tokens, k_heads, head_dim),
atol=0,
rtol=0,
)
torch.testing.assert_close(
idx_k_cache.index_select(0, out_cache_loc),
ref[3].view(num_tokens, 1, head_dim),
atol=3e-2,
rtol=3e-2,
)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v", "-s"]))