[Spec] Fuse small kenrels under gather_spec_extras (#27233)
This commit is contained in:
Khoa Pham
@@ -0,0 +1,214 @@
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
_OUTPUT_NAMES = ("topk_p", "topk_index", "bonus_tokens", "hidden_states")
|
||||
|
||||
|
||||
def _ref_gather(
|
||||
indices, topk_p_buf, topk_index_buf, output_tokens_buf, hidden_states_buf
|
||||
):
|
||||
"""Reference oracle: the exact torch.compile'd advanced-index gather that the
|
||||
fused Triton kernel replaced (see overlap_utils._gather_spec_extras pre-fusion).
|
||||
A gather is a pure copy, so the kernel must match this bit-for-bit."""
|
||||
topk_p = topk_p_buf[indices]
|
||||
topk_index = topk_index_buf[indices]
|
||||
bonus_tokens = output_tokens_buf[indices]
|
||||
hidden_states = (
|
||||
hidden_states_buf[indices] if hidden_states_buf is not None else None
|
||||
)
|
||||
return topk_p, topk_index, bonus_tokens, hidden_states
|
||||
|
||||
|
||||
def _make_buffers(
|
||||
pool_size,
|
||||
topk,
|
||||
hidden_dim,
|
||||
*,
|
||||
with_hidden,
|
||||
hidden_dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
seed=0,
|
||||
):
|
||||
"""Build FutureMap-shaped relay buffers.
|
||||
|
||||
Mirrors overlap_utils.FutureMap: topk_p / topk_index / hidden_states are
|
||||
2-D (pool_size, width) while output_tokens is 1-D (pool_size,). The width
|
||||
mix (incl. the 1-D buffer -> row width 1) exercises the kernel's per-buffer
|
||||
masking.
|
||||
"""
|
||||
g = torch.Generator(device=device).manual_seed(seed)
|
||||
topk_p_buf = torch.rand(
|
||||
(pool_size, topk), dtype=torch.float32, device=device, generator=g
|
||||
)
|
||||
topk_index_buf = torch.randint(
|
||||
0, 32000, (pool_size, topk), dtype=torch.int64, device=device, generator=g
|
||||
)
|
||||
output_tokens_buf = torch.randint(
|
||||
0, 32000, (pool_size,), dtype=torch.int64, device=device, generator=g
|
||||
)
|
||||
hidden_states_buf = (
|
||||
torch.randn(
|
||||
(pool_size, hidden_dim), dtype=hidden_dtype, device=device, generator=g
|
||||
)
|
||||
if with_hidden
|
||||
else None
|
||||
)
|
||||
return topk_p_buf, topk_index_buf, output_tokens_buf, hidden_states_buf
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
|
||||
class TestGatherSpecExtras(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.device = torch.device("cuda")
|
||||
|
||||
def _assert_matches_reference(self, indices, bufs):
|
||||
"""Run fused kernel + reference on the same inputs and assert every
|
||||
output is identical (dtype, shape, exact values) and that the source
|
||||
buffers are never mutated."""
|
||||
src_snapshots = [None if b is None else b.clone() for b in bufs]
|
||||
|
||||
ref = _ref_gather(indices, *bufs)
|
||||
got = gather_spec_extras(indices, *bufs)
|
||||
|
||||
self.assertEqual(len(got), len(ref))
|
||||
for name, r, o in zip(_OUTPUT_NAMES, ref, got):
|
||||
if r is None:
|
||||
self.assertIsNone(o, f"{name} should be None when buffer is None")
|
||||
continue
|
||||
self.assertIsNotNone(o, f"{name} unexpectedly None")
|
||||
self.assertEqual(o.dtype, r.dtype, f"{name} dtype mismatch")
|
||||
self.assertEqual(tuple(o.shape), tuple(r.shape), f"{name} shape mismatch")
|
||||
self.assertEqual(o.device.type, r.device.type, f"{name} device mismatch")
|
||||
# Pure gather == bit-exact copy, so require zero tolerance.
|
||||
torch.testing.assert_close(
|
||||
o, r, rtol=0, atol=0, msg=f"{name} value mismatch"
|
||||
)
|
||||
|
||||
# The kernel only reads sources; it must not scribble into them.
|
||||
for name, before, buf in zip(_OUTPUT_NAMES, src_snapshots, bufs):
|
||||
if before is None:
|
||||
continue
|
||||
torch.testing.assert_close(
|
||||
buf, before, rtol=0, atol=0, msg=f"source buffer {name} was mutated"
|
||||
)
|
||||
|
||||
def test_matches_reference_across_shapes(self):
|
||||
# (pool_size, m, topk, hidden_dim). Covers: m<pool, m==pool, topk==1,
|
||||
# 1-column blocks, exact-1024-width boundary, wide multi-block widths,
|
||||
# and wide non-power-of-2 widths (partial trailing column block).
|
||||
configs = [
|
||||
(16, 8, 1, 7),
|
||||
(64, 33, 4, 128),
|
||||
(128, 128, 8, 1024),
|
||||
(100, 50, 2, 4096),
|
||||
(257, 200, 16, 4097),
|
||||
(2048, 777, 8, 5120),
|
||||
]
|
||||
for pool_size, m, topk, hidden_dim in configs:
|
||||
for with_hidden in (True, False):
|
||||
with self.subTest(
|
||||
pool_size=pool_size,
|
||||
m=m,
|
||||
topk=topk,
|
||||
hidden_dim=hidden_dim,
|
||||
with_hidden=with_hidden,
|
||||
):
|
||||
bufs = _make_buffers(
|
||||
pool_size,
|
||||
topk,
|
||||
hidden_dim,
|
||||
with_hidden=with_hidden,
|
||||
device=self.device,
|
||||
)
|
||||
indices = torch.randint(
|
||||
0, pool_size, (m,), dtype=torch.int64, device=self.device
|
||||
)
|
||||
self._assert_matches_reference(indices, bufs)
|
||||
|
||||
def test_empty_indices_returns_empty_rows(self):
|
||||
# m == 0 hits the early-return path; outputs must still carry the right
|
||||
# trailing dims / dtypes so downstream concatenation stays valid.
|
||||
for with_hidden in (True, False):
|
||||
with self.subTest(with_hidden=with_hidden):
|
||||
bufs = _make_buffers(
|
||||
32, 4, 256, with_hidden=with_hidden, device=self.device
|
||||
)
|
||||
indices = torch.empty(0, dtype=torch.int64, device=self.device)
|
||||
self._assert_matches_reference(indices, bufs)
|
||||
|
||||
def test_non_contiguous_indices(self):
|
||||
# indices flows from filtered/merged producers and can be strided; the
|
||||
# kernel addresses it linearly and relies on the internal .contiguous().
|
||||
pool_size, m = 256, 64
|
||||
bufs = _make_buffers(pool_size, 8, 512, with_hidden=True, device=self.device)
|
||||
pairs = torch.randint(
|
||||
0, pool_size, (m, 2), dtype=torch.int64, device=self.device
|
||||
)
|
||||
indices = pairs[:, 0]
|
||||
self.assertFalse(indices.is_contiguous(), "test setup: indices must be strided")
|
||||
self._assert_matches_reference(indices, bufs)
|
||||
|
||||
def test_duplicate_indices(self):
|
||||
# Gather (not scatter): repeated source rows are well-defined and must
|
||||
# each produce an identical copy.
|
||||
pool_size, m = 8, 64
|
||||
bufs = _make_buffers(pool_size, 4, 333, with_hidden=True, device=self.device)
|
||||
indices = torch.randint(
|
||||
0, 3, (m,), dtype=torch.int64, device=self.device
|
||||
) # tiny range -> many duplicates
|
||||
self._assert_matches_reference(indices, bufs)
|
||||
|
||||
def test_index_dtype_variants(self):
|
||||
pool_size, m = 128, 50
|
||||
bufs = _make_buffers(pool_size, 8, 1024, with_hidden=True, device=self.device)
|
||||
base = torch.randint(0, pool_size, (m,), device=self.device)
|
||||
for idx_dtype in (torch.int32, torch.int64):
|
||||
with self.subTest(idx_dtype=idx_dtype):
|
||||
self._assert_matches_reference(base.to(idx_dtype), bufs)
|
||||
|
||||
def test_hidden_dtype_variants(self):
|
||||
pool_size, m = 96, 40
|
||||
indices = torch.randint(
|
||||
0, pool_size, (m,), dtype=torch.int64, device=self.device
|
||||
)
|
||||
for hidden_dtype in (torch.bfloat16, torch.float16, torch.float32):
|
||||
with self.subTest(hidden_dtype=hidden_dtype):
|
||||
bufs = _make_buffers(
|
||||
pool_size,
|
||||
8,
|
||||
2048,
|
||||
with_hidden=True,
|
||||
hidden_dtype=hidden_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
self._assert_matches_reference(indices, bufs)
|
||||
|
||||
def test_outputs_do_not_alias_source_buffers(self):
|
||||
pool_size, m = 64, 32
|
||||
bufs = _make_buffers(pool_size, 8, 512, with_hidden=True, device=self.device)
|
||||
indices = torch.randint(
|
||||
0, pool_size, (m,), dtype=torch.int64, device=self.device
|
||||
)
|
||||
outputs = gather_spec_extras(indices, *bufs)
|
||||
for name, out, buf in zip(_OUTPUT_NAMES, outputs, bufs):
|
||||
if out is None or buf is None:
|
||||
continue
|
||||
self.assertNotEqual(
|
||||
out.data_ptr(),
|
||||
buf.data_ptr(),
|
||||
f"{name} output aliases its source buffer",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user