mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
558 lines
26 KiB
Python
558 lines
26 KiB
Python
"""CPU correctness tests: Neuron pure-torch tile_kernels vs the upstream tile_kernels.torch references."""
|
|
|
|
import os
|
|
import sys
|
|
|
|
VENDOR = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "neuron_ops", "vendor")
|
|
sys.path.insert(0, VENDOR)
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tile_kernels
|
|
from tile_kernels import quant
|
|
from tile_kernels import torch as ref
|
|
from tile_kernels.quant.common import get_packed_ue8m0_torch_dtype
|
|
|
|
torch.set_num_threads(4)
|
|
|
|
|
|
def bits(t):
|
|
if t.dtype in (torch.float8_e4m3fn,):
|
|
return t.view(torch.uint8)
|
|
return t
|
|
|
|
|
|
def assert_same(a, b, check_stride=True):
|
|
assert a.dtype == b.dtype, (a.dtype, b.dtype)
|
|
assert a.shape == b.shape, (a.shape, b.shape)
|
|
if check_stride and a.numel() > 1:
|
|
assert a.stride() == b.stride(), (a.stride(), b.stride())
|
|
if a.dtype.is_floating_point and a.dtype != torch.float8_e4m3fn:
|
|
assert torch.equal(torch.nan_to_num(a), torch.nan_to_num(b)), (a - b).abs().max()
|
|
else:
|
|
assert torch.equal(bits(a), bits(b))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# config / surface
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_surface():
|
|
assert tile_kernels.config.is_ascend() is False
|
|
assert tile_kernels.get_num_vec_cores() == 2 * tile_kernels.get_num_sms()
|
|
assert tile_kernels.get_max_ub_per_vector_core() > 0 and tile_kernels.config.get_max_ub_per_vector_core() > 0
|
|
prev = tile_kernels.get_num_sms()
|
|
tile_kernels.set_num_sms(3)
|
|
assert tile_kernels.get_num_sms() == 3 and tile_kernels.get_num_vec_cores() == 6
|
|
tile_kernels.set_num_sms(prev)
|
|
tile_kernels.set_pdl(True)
|
|
tile_kernels.set_pdl(False)
|
|
tile_kernels.use_deterministic_algorithms(True)
|
|
tile_kernels.set_token_alignment(64)
|
|
assert tile_kernels.get_token_alignment() == 64
|
|
tile_kernels.set_token_alignment(128)
|
|
assert tile_kernels.utils.align(5, 4) == 8
|
|
assert tile_kernels.utils.dtype_to_str(torch.bfloat16) == 'bf16'
|
|
assert get_packed_ue8m0_torch_dtype() == torch.int16
|
|
from tile_kernels.config import is_ascend # noqa
|
|
from tile_kernels.mhc.expand_kernel import expand_to_mhc_fwd # noqa
|
|
from tile_kernels.mhc.head_compute_mix_kernel import mhc_head_compute_mix_fwd # noqa
|
|
from tile_kernels.mhc.norm_fn_kernel import mhc_fn_normw_merge_fwd, mhc_reduce_partials_and_rmsnorm_fwd # noqa
|
|
from tile_kernels.mhc.post_kernel import mhc_post_fwd # noqa
|
|
from tile_kernels.mhc.pre_apply_mix_kernel import mhc_pre_apply_mix_fwd # noqa
|
|
from tile_kernels.mhc.pre_split_mixes_kernel import mhc_pre_split_mixes_fwd # noqa
|
|
from tile_kernels.mhc.sinkhorn_kernel import mhc_sinkhorn_fwd # noqa
|
|
from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd # noqa
|
|
from tile_kernels.moe import moe_topk_gate # noqa
|
|
from tile_kernels.quant import swiglu_forward_and_per_token_cast, swiglu_forward, per_token_cast, batched_transpose_weight_sf # noqa
|
|
from tile_kernels.transform import apply_rotary # noqa
|
|
from tile_kernels.torch import cast, cast_back # noqa
|
|
from tile_kernels import get_max_ub_per_vector_core, get_num_vec_cores # noqa
|
|
for name in ('engram_hash', 'engram_gate_fwd', 'fused_weight'):
|
|
assert callable(getattr(tile_kernels.engram, name))
|
|
assert callable(tile_kernels.torch.engram.make_offsets)
|
|
for name in ('per_token_cast', 'per_token_cast_back', 'per_token_cast_and_cast_back'):
|
|
assert callable(getattr(quant, name))
|
|
import importlib
|
|
|
|
for mod in ('benchmark', 'gpu_mem', 'random', 'tilelang_warning', 'xdist_failfast'):
|
|
importlib.import_module(f'tile_kernels.testing.pytest.{mod}')
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# per_token_cast
|
|
# ---------------------------------------------------------------------------
|
|
|
|
LAYOUTS = [
|
|
# round_sf, packed, col_major, e4m3_sf
|
|
(False, False, False, False),
|
|
(True, False, False, False),
|
|
(True, True, False, False),
|
|
(True, True, True, False),
|
|
(True, False, True, False),
|
|
(False, False, True, False),
|
|
(False, False, False, True),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize('layout', LAYOUTS)
|
|
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
|
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
|
|
@pytest.mark.parametrize('m,hidden', [(1, 128), (7, 256), (33, 512), (64, 192)])
|
|
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
|
def test_per_token_cast(layout, fmt, npc, m, hidden, dtype):
|
|
round_sf, packed, col, e4m3_sf = layout
|
|
if hidden % npc:
|
|
pytest.skip('hidden % npc')
|
|
torch.manual_seed(m * 1000 + hidden + npc)
|
|
x = (torch.randn(m, hidden) * torch.logspace(-3, 3, hidden)).to(dtype)
|
|
x[0, :npc] = 0 # all-zero group exercises the amax clamp
|
|
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
|
|
q, sf = quant.per_token_cast(x, fmt, npc, **kw)
|
|
rq, rsf = ref.cast(x, fmt, (1, npc), **kw)
|
|
assert_same(q, rq)
|
|
assert_same(sf, rsf)
|
|
# sf_only / precomputed sf variants
|
|
sf2 = quant.per_token_cast_with_sf_only(x, fmt, npc, use_tma_aligned_col_major_sf=col, round_sf=round_sf, use_packed_ue8m0=packed, use_e4m3_sf=e4m3_sf)
|
|
assert_same(sf2, rsf)
|
|
if not packed and not col:
|
|
q3 = quant.per_token_cast_with_precomputed_sf(x, fmt, npc, rsf)
|
|
rq3 = ref.cast(x, fmt, (1, npc), sf=rsf.float())
|
|
assert_same(q3, rq3)
|
|
# out= in place
|
|
if not col:
|
|
out_x = torch.empty_like(q)
|
|
out_sf = torch.empty_like(sf)
|
|
rx, rs = quant.per_token_cast(x, fmt, npc, out=(out_x, out_sf), **kw)
|
|
assert rx.data_ptr() == out_x.data_ptr() and rs.data_ptr() == out_sf.data_ptr()
|
|
assert_same(out_x, rq)
|
|
assert_same(out_sf, rsf)
|
|
|
|
|
|
def test_per_token_cast_ascend_layout_bytes():
|
|
"""Packed col-major sf: byte (k % 2) of storage element [k // 2, m] is exponent of group k (Ascend get_sf_index)."""
|
|
torch.manual_seed(0)
|
|
m, hidden, npc = 5, 256, 32
|
|
x = torch.randn(m, hidden, dtype=torch.bfloat16) * 10
|
|
_, sf = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
|
assert sf.dtype == torch.int16 and sf.shape == (m, hidden // npc // 2) and sf.stride() == (1, m)
|
|
storage = sf.T.contiguous().view(torch.uint8) # (P, 2m) bytes
|
|
_, sf_rm = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=False)
|
|
exp = (sf_rm.view(torch.int32) >> 23)
|
|
for mi in range(m):
|
|
for k in range(hidden // npc):
|
|
assert storage[k // 2, 2 * mi + k % 2].item() == exp[mi, k].item()
|
|
# row-major packed: byte 2p + j of row m = group 2p + j
|
|
_, sfp = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True)
|
|
assert sfp.shape == (m, hidden // npc // 2) and sfp.is_contiguous()
|
|
assert torch.equal(sfp.view(torch.uint8).to(torch.int32), exp)
|
|
|
|
|
|
@pytest.mark.parametrize('in_layout', ['fp32_128', 'packed_col_32', 'packed_row_32', 'fp32_1x32'])
|
|
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
|
def test_per_token_cast_with_input_sf(in_layout, fmt):
|
|
torch.manual_seed(1)
|
|
m, hidden = 256, 256
|
|
w = torch.randn(m, hidden) * 3
|
|
if in_layout == 'fp32_128':
|
|
xq, xsf = ref.cast(w, 'e4m3', (128, 128), round_sf=True)
|
|
blk = (128, 128)
|
|
elif in_layout == 'fp32_1x32':
|
|
xq, xsf = ref.cast(w, 'e4m3', (1, 32))
|
|
blk = (1, 32)
|
|
elif in_layout == 'packed_col_32':
|
|
xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
|
blk = (1, 32)
|
|
else:
|
|
xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True)
|
|
blk = (1, 32)
|
|
kw = dict(round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True)
|
|
q, sf = quant.per_token_cast((xq, xsf), fmt, 32, x_block_size=blk, **kw)
|
|
rq, rsf = ref.cast((xq, xsf), fmt, (1, 32), x_block_size=blk, **kw)
|
|
assert_same(q, rq)
|
|
assert_same(sf, rsf)
|
|
|
|
|
|
def test_per_token_cast_empty():
|
|
x = torch.empty(0, 128, dtype=torch.bfloat16)
|
|
q, sf = quant.per_token_cast(x, 'e4m3', 32, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
|
assert q.shape == (0, 128) and sf.shape == (0, 2) and sf.dtype == torch.int16
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# cast_back / cast_and_cast_back / per_block_cast
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize('layout', LAYOUTS)
|
|
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
|
@pytest.mark.parametrize('npc', [16, 32, 128])
|
|
@pytest.mark.parametrize('out_fmt', ['bf16', 'fp32'])
|
|
def test_per_token_cast_back(layout, fmt, npc, out_fmt):
|
|
round_sf, packed, col, e4m3_sf = layout
|
|
torch.manual_seed(npc)
|
|
x = torch.randn(19, 256, dtype=torch.bfloat16) * 5
|
|
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
|
|
q, sf = ref.cast(x, fmt, (1, npc), **kw)
|
|
out = quant.per_token_cast_back((q, sf), out_fmt, npc)
|
|
rout = ref.cast_back((q, sf), out_fmt, (1, npc))
|
|
assert_same(out, rout)
|
|
|
|
|
|
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
|
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
|
|
@pytest.mark.parametrize('sf_mode', ['plain', 'round', 'e4m3'])
|
|
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
|
def test_per_token_cast_and_cast_back(fmt, npc, sf_mode, dtype):
|
|
torch.manual_seed(npc)
|
|
base = torch.randn(23, 384, dtype=dtype) * 7
|
|
big = torch.zeros(23, 512, dtype=dtype)
|
|
big[:, :384] = base
|
|
x = big[:, :384] # strided rows
|
|
round_sf, e4m3_sf = sf_mode == 'round', sf_mode == 'e4m3'
|
|
quant.per_token_cast_and_cast_back(x, fmt, npc, round_sf, e4m3_sf)
|
|
q, sf = ref.cast(base, fmt, (1, npc), round_sf=round_sf, use_e4m3_sf=e4m3_sf)
|
|
expected = ref.cast_back((q, sf), 'fp32', (1, npc)).to(dtype)
|
|
assert torch.equal(x, expected)
|
|
assert torch.equal(big[:, 384:], torch.zeros(23, 128, dtype=dtype))
|
|
|
|
|
|
@pytest.mark.parametrize('block', [(128, 128), (32, 32), (1, 128)])
|
|
@pytest.mark.parametrize('layout', LAYOUTS[:6])
|
|
def test_per_block_cast(block, layout):
|
|
round_sf, packed, col, _ = layout
|
|
torch.manual_seed(2)
|
|
x = torch.randn(256, 384, dtype=torch.bfloat16)
|
|
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col)
|
|
q, sf = quant.per_block_cast(x, 'e4m3', block, **kw)
|
|
rq, rsf = ref.cast(x, 'e4m3', block, **kw)
|
|
assert_same(q, rq)
|
|
assert_same(sf, rsf)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# swiglu
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
|
@pytest.mark.parametrize('with_weights', [False, True])
|
|
@pytest.mark.parametrize('clamp', [None, 3.0])
|
|
@pytest.mark.parametrize('with_psum', [False, True])
|
|
def test_swiglu_forward(dtype, with_weights, clamp, with_psum):
|
|
torch.manual_seed(3)
|
|
tile_kernels.set_token_alignment(8)
|
|
n, hidden = 64, 256
|
|
x = (torch.randn(n, hidden * 2) * 3).to(dtype)
|
|
w = torch.rand(n) if with_weights else None
|
|
rsf = 2.5 if with_weights else None
|
|
psum = torch.tensor([5, 8, 8, 20, 37, 44], dtype=torch.int32) if with_psum else None # aligned-prefix layout
|
|
for fmt in ('fp32', 'bf16'):
|
|
cnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
|
|
rcnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
|
|
out = quant.swiglu_forward(x, fmt, psum, w, rsf, clamp, cnt)
|
|
rout = ref.swiglu_forward(x, psum, w, rsf, clamp, rcnt)
|
|
assert torch.equal(out, rout.to(out.dtype))
|
|
if cnt is not None:
|
|
assert torch.equal(cnt, rcnt), (cnt, rcnt)
|
|
# fused fp8 cast (deepinfer get_fp8_cast_params)
|
|
kw = dict(round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
|
for npc in (32, 128):
|
|
q, sf = quant.swiglu_forward_and_per_token_cast(x, 'e4m3', npc, psum, w, rsf, clamp_value=clamp, **kw)
|
|
rq, rsf_ = ref.cast(ref.swiglu_forward(x, psum, w, rsf, clamp), 'e4m3', (1, npc), **kw)
|
|
assert_same(q, rq)
|
|
assert_same(sf, rsf_)
|
|
tile_kernels.set_token_alignment(128)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# batched_transpose_weight_sf
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize('block', [32, 128])
|
|
def test_batched_transpose_weight_sf(block):
|
|
torch.manual_seed(4)
|
|
|
|
def make(*shape):
|
|
e = torch.randint(110, 140, shape, dtype=torch.int32)
|
|
return (e << 23).view(torch.float32)
|
|
|
|
sfs = [make(384 // block, 640 // block), make(3, 256 // block, 896 // block), make(128 // block, 128 // block)]
|
|
flags = [False, True, True]
|
|
outs = quant.batched_transpose_weight_sf(sfs, block, flags, outs=[None] * 3)
|
|
refs = ref.batched_transpose_weight_sf(sfs, block, flags)
|
|
for (a0, a1), (r0, r1) in zip(outs, refs):
|
|
assert_same(a0, r0)
|
|
assert_same(a1, r1)
|
|
# in-place rewrite with new values
|
|
sfs2 = [make(*s.shape) for s in sfs]
|
|
outs2 = quant.batched_transpose_weight_sf(sfs2, block, flags, outs=outs)
|
|
refs2 = ref.batched_transpose_weight_sf(sfs2, block, flags)
|
|
for (a0, a1), (o0, o1), (r0, r1) in zip(outs2, outs, refs2):
|
|
assert a0.data_ptr() == o0.data_ptr() and a1.data_ptr() == o1.data_ptr()
|
|
assert_same(a0, r0)
|
|
assert_same(a1, r1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# rope
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize('rot_dim', [64, 128])
|
|
@pytest.mark.parametrize('ndim', [3, 4])
|
|
@pytest.mark.parametrize('interleaved', [False, True])
|
|
@pytest.mark.parametrize('conjugate', [False, True])
|
|
@pytest.mark.parametrize('use_positions', [False, True])
|
|
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
|
def test_apply_rotary(rot_dim, ndim, interleaved, conjugate, use_positions, dtype):
|
|
torch.manual_seed(5)
|
|
seqlen_ro, s, h, hk, b = 64, 9, 4, 2, 2
|
|
cache = torch.randn(seqlen_ro, rot_dim)
|
|
cos, sin = cache[:, : rot_dim // 2], cache[:, rot_dim // 2 :]
|
|
lead = (s,) if ndim == 3 else (b, s)
|
|
full_q = torch.randn(*lead, h, rot_dim + 32).to(dtype)
|
|
q = full_q[..., 32:] # tail partial rope, strided
|
|
k = torch.randn(*lead, hk, rot_dim).to(dtype)
|
|
pos = torch.randint(0, seqlen_ro - 4, lead, dtype=torch.int32) if use_positions else None
|
|
off = 3
|
|
rq = ref.rotary_embedding_ref(q.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
|
|
rk = ref.rotary_embedding_ref(k.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
|
|
head = full_q[..., :32].clone()
|
|
tile_kernels.transform.apply_rotary(q, cache, k, pos, interleaved, conjugate, off)
|
|
tol = dict(atol=0, rtol=0) if dtype == torch.float32 else dict(atol=1e-2, rtol=1e-2)
|
|
torch.testing.assert_close(q.float(), rq.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
|
|
torch.testing.assert_close(k.float(), rk.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
|
|
assert torch.equal(full_q[..., :32], head)
|
|
# bitwise vs the Ascend op order: out1 = fma(first, cos, second * -sin), out2 = fma(first, sin, second * cos)
|
|
def fma(a, b, c):
|
|
return (a.double() * b.double() + c.double()).float()
|
|
|
|
q2 = full_q[..., 32:].clone()
|
|
q0 = torch.randn_like(q2.float()).to(dtype)
|
|
q2.copy_(q0)
|
|
tile_kernels.transform.apply_rotary(q2, cache, None, pos, interleaved, conjugate, off)
|
|
idx = (pos.long() + off) if pos is not None else (torch.arange(s) + off).expand(*lead[:-1], s)
|
|
c_, s_ = cos[idx].unsqueeze(-2), sin[idx].unsqueeze(-2)
|
|
s_ = -s_ if conjugate else s_
|
|
xf = q0.float()
|
|
f, sc = (xf[..., 0::2], xf[..., 1::2]) if interleaved else (xf[..., : rot_dim // 2], xf[..., rot_dim // 2 :])
|
|
c_, s_ = c_.expand_as(f), s_.expand_as(f)
|
|
o1, o2 = fma(f, c_, sc * -s_), fma(f, s_, sc * c_)
|
|
exp = torch.stack([o1, o2], -1).flatten(-2) if interleaved else torch.cat([o1, o2], -1)
|
|
assert torch.equal(q2, exp.to(dtype))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# moe
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize('num_shared', [0, 1, 2])
|
|
@pytest.mark.parametrize('with_map', [False, True])
|
|
@pytest.mark.parametrize('with_mask', [False, True])
|
|
@pytest.mark.parametrize('with_image', [False, True])
|
|
@pytest.mark.parametrize('with_fix', [False, True])
|
|
def test_moe_topk_gate(num_shared, with_map, with_mask, with_image, with_fix):
|
|
torch.manual_seed(6)
|
|
t, e, k = 37, 64, 6
|
|
logits = torch.randn(t, e) * 2
|
|
logits[3, 10] = logits[3, 11] # tie
|
|
bias = torch.randn(e) * 0.1
|
|
kwargs = dict(num_topk=k, use_shared_as_routed=num_shared > 0, num_shared_experts=num_shared, routed_scaling_factor=2.5, ep_rank=3, scoring_func='sqrtsoftplus', bias=bias)
|
|
if with_mask:
|
|
kwargs['mask'] = torch.rand(t) > 0.3
|
|
if with_image:
|
|
kwargs['image_bias'] = torch.randn(e) * 0.1
|
|
kwargs['image_token_mask'] = torch.rand(t) > 0.5
|
|
if with_map:
|
|
nl = e + num_shared
|
|
dup = 3
|
|
count = torch.randint(1, dup + 1, (nl,), dtype=torch.int32)
|
|
kwargs['to_physical_map'] = torch.randint(0, 500, (nl, dup), dtype=torch.int32)
|
|
kwargs['logical_count'] = count
|
|
unm = unm_ref = None
|
|
if with_fix:
|
|
unm = torch.stack([torch.randperm(e)[:k] for _ in range(t)]).to(torch.int64)
|
|
unm_ref = unm.clone()
|
|
kwargs['fix_routing_mask'] = torch.rand(t) > 0.5
|
|
idx, w = tile_kernels.moe.moe_topk_gate(logits, unmapped_topk_idx=unm, **kwargs)
|
|
ridx, rw = ref.moe_topk_gate_forward(logits, unmapped_topk_idx=unm_ref, **kwargs)
|
|
assert torch.equal(idx, ridx)
|
|
torch.testing.assert_close(w, rw, atol=1e-6, rtol=1e-6)
|
|
if with_fix:
|
|
assert torch.equal(unm, unm_ref)
|
|
|
|
|
|
def test_moe_topk_gate_out_scores_force_random():
|
|
torch.manual_seed(7)
|
|
t, e, k = 16, 32, 4
|
|
logits = torch.randn(t, e)
|
|
mask = torch.ones(t, dtype=torch.bool)
|
|
mask[0] = False
|
|
fr = torch.zeros(t, dtype=torch.bool)
|
|
fr[1:4] = True
|
|
out = (torch.empty(t, k, dtype=torch.int64), torch.empty(t, k))
|
|
scores = torch.empty(t, e)
|
|
unm = torch.empty(t, k, dtype=torch.int64)
|
|
idx, w = tile_kernels.moe.moe_topk_gate(logits, k, False, 0, 1.0, 0, mask=mask, force_random=fr, out=out, scores=scores, unmapped_topk_idx=unm)
|
|
assert idx.data_ptr() == out[0].data_ptr()
|
|
assert (idx[0] == -1).all() and (w[0] == 0).all() and (scores[0] == 0).all() and (unm[0] == -1).all()
|
|
assert (unm[1:4] == -1).all() and (scores[1:4] == 0).all()
|
|
for r in range(1, 4):
|
|
assert len(set(idx[r].tolist())) == k and idx[r].min() >= 0 and idx[r].max() < e
|
|
ridx, rw = ref.moe_topk_gate_forward(logits[4:], k, False, 0, 1.0, 0, 'sqrtsoftplus')
|
|
assert torch.equal(idx[4:], ridx)
|
|
torch.testing.assert_close(scores[4:], torch.nn.functional.softplus(logits[4:]).sqrt())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# engram
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_engram_hash():
|
|
torch.manual_seed(8)
|
|
t, n, layers, tables = 29, 4, 2, 3
|
|
ids = torch.randint(0, 100000, (t, n), dtype=torch.int32)
|
|
mult = torch.randint(1, 2**40, (layers, n), dtype=torch.int64)
|
|
vs = torch.randint(1000, 50000, (layers, n - 1, tables), dtype=torch.int32)
|
|
off = tile_kernels.torch.engram.make_offsets(vs)
|
|
mask = torch.rand(t) > 0.7
|
|
for m in (None, mask):
|
|
out = tile_kernels.engram.engram_hash(ids, mult, vs, off, m)
|
|
rout = ref.engram_hash_ref(ids, mult, vs, off, m)
|
|
assert out.dtype == torch.int32 and torch.equal(out, rout)
|
|
|
|
|
|
@pytest.mark.parametrize('inplace', [False, True])
|
|
@pytest.mark.parametrize('with_mask', [False, True])
|
|
def test_engram_gate(inplace, with_mask):
|
|
torch.manual_seed(9)
|
|
t, hc, d = 13, 4, 256
|
|
hs = torch.randn(t, hc, d, dtype=torch.bfloat16)
|
|
kv = torch.randn(t, hc + 1, d, dtype=torch.bfloat16)
|
|
wh = torch.randn(hc, d, dtype=torch.bfloat16)
|
|
we = torch.randn(hc, d, dtype=torch.bfloat16)
|
|
mask = (torch.rand(t) > 0.5) if with_mask else None
|
|
wf = tile_kernels.engram.fused_weight(wh, we)
|
|
assert wf.dtype == torch.float32 and torch.equal(wf, wh.float() * we.float())
|
|
r_out, r_dot, r_g, r_rx, r_rk = ref.engram_gate_ref(hs, kv, wh, we, 1e-4, 1e-6, save_for_backward=True, image_token_mask=mask)
|
|
hs_in = hs.clone()
|
|
out, dot, g, rx, rk = tile_kernels.engram.engram_gate_fwd(hs_in, kv, wf, 1e-6, 1e-4, save_for_backward=True, image_token_mask=mask, inplace=inplace)
|
|
if inplace:
|
|
assert out.data_ptr() == hs_in.data_ptr()
|
|
sel = ~mask if mask is not None else slice(None)
|
|
assert (out.float() - r_out.float()).abs().max() <= 0.0625
|
|
assert torch.equal(out[mask], hs[mask]) if mask is not None else True
|
|
torch.testing.assert_close(dot[sel], r_dot[sel], atol=1e-3, rtol=1e-4)
|
|
torch.testing.assert_close(g[sel], r_g[sel], atol=1e-5, rtol=1e-5)
|
|
torch.testing.assert_close(rx[sel], r_rx[sel], atol=1e-6, rtol=1e-5)
|
|
torch.testing.assert_close(rk[sel], r_rk[sel], atol=1e-6, rtol=1e-5)
|
|
out2, *rest = tile_kernels.engram.engram_gate_fwd(hs.clone(), kv, wf, 1e-6, 1e-4, save_for_backward=False, image_token_mask=mask)
|
|
assert all(r is None for r in rest)
|
|
assert torch.equal(out2, out)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# mhc
|
|
# ---------------------------------------------------------------------------
|
|
|
|
MHC = 4
|
|
|
|
|
|
def test_mhc_expand_and_mix():
|
|
torch.manual_seed(10)
|
|
n, h = 11, 256
|
|
x = torch.randn(n, h, dtype=torch.bfloat16)
|
|
o = torch.empty(n, MHC, h, dtype=torch.bfloat16)
|
|
tile_kernels.mhc.expand_to_mhc_fwd(x, o)
|
|
assert torch.equal(o, ref.expand_to_mhc_ref(x, MHC))
|
|
|
|
inp = torch.randn(n, MHC)
|
|
scale, base = torch.randn(1), torch.randn(MHC)
|
|
out = torch.empty_like(inp)
|
|
tile_kernels.mhc.mhc_head_compute_mix_fwd(inp, scale, base, out, 1e-6)
|
|
torch.testing.assert_close(out, ref.mhc_head_compute_mix_ref(inp, scale, base, 1e-6), atol=0, rtol=0)
|
|
|
|
m3 = MHC * (MHC + 2)
|
|
mixes = torch.randn(1, n, m3)
|
|
scale3, base3 = torch.randn(3), torch.randn(m3)
|
|
pre, post, comb = torch.empty(n, MHC), torch.empty(n, MHC), torch.empty(n, MHC * MHC)
|
|
tile_kernels.mhc.mhc_pre_split_mixes_fwd(mixes.view(n, m3), scale3, base3, pre, post, comb, 2.0, 1e-6)
|
|
rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale3, base3, MHC, 2.0, 1e-6)
|
|
torch.testing.assert_close(pre, rpre.view(n, MHC), atol=1e-7, rtol=1e-6)
|
|
torch.testing.assert_close(post, rpost.view(n, MHC), atol=1e-7, rtol=1e-6)
|
|
torch.testing.assert_close(comb, rcomb.reshape(n, MHC * MHC), atol=1e-7, rtol=1e-6)
|
|
|
|
xr = torch.randn(n, MHC, h, dtype=torch.bfloat16)
|
|
mix = torch.rand(n, MHC)
|
|
o2 = torch.empty(n, h, dtype=torch.bfloat16)
|
|
tile_kernels.mhc.mhc_pre_apply_mix_fwd(xr, mix, o2)
|
|
torch.testing.assert_close(o2.float(), ref.mhc_pre_apply_mix_ref(xr.float(), mix.unsqueeze(-1)).float(), atol=1e-2, rtol=1e-2)
|
|
|
|
cm = torch.randn(n, MHC, MHC)
|
|
cmo = torch.empty_like(cm)
|
|
tile_kernels.mhc.mhc_sinkhorn_fwd(cm, cmo, 10, 1e-6)
|
|
torch.testing.assert_close(cmo, ref.sinkhorn_normalize_ref(cm, 10, 1e-6), atol=1e-6, rtol=1e-5)
|
|
|
|
|
|
def test_mhc_post():
|
|
torch.manual_seed(11)
|
|
a, b, h = 2, 5, 256
|
|
x = torch.randn(a, b, h, dtype=torch.bfloat16)
|
|
res = torch.randn(a, b, MHC, h, dtype=torch.bfloat16)
|
|
post = torch.rand(a, b, MHC, 1)
|
|
comb = torch.rand(a, b, MHC, MHC)
|
|
out = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb)
|
|
rout = ref.mhc_post_ref(x, res, post, comb)
|
|
assert out.dtype == torch.bfloat16 and out.shape == res.shape
|
|
assert (out.float() - rout.float()).abs().max() <= 0.0625
|
|
out2 = torch.empty_like(res)
|
|
r = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb, out=out2)
|
|
assert r.data_ptr() == out2.data_ptr() and torch.equal(out2, out)
|
|
|
|
|
|
@pytest.mark.parametrize('n_splits', [1, 4, 16])
|
|
def test_mhc_norm_fn(n_splits):
|
|
torch.manual_seed(12)
|
|
t, h = 9, 256
|
|
m3 = MHC * (MHC + 2)
|
|
residual = torch.randn(1, t, MHC, h, dtype=torch.bfloat16)
|
|
fn = torch.randn(m3, MHC * h) * 0.02
|
|
normw = torch.rand(MHC * h) + 0.5
|
|
merged = torch.empty_like(fn)
|
|
tile_kernels.mhc.mhc_fn_normw_merge_fwd(fn, normw, merged)
|
|
assert torch.equal(merged, fn * normw)
|
|
om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual, merged, n_splits) # (S, T, 1, M3), (S, T, 1)
|
|
om, sq, out = torch.empty(t, 1, m3), torch.empty(t, 1), torch.empty(t, m3)
|
|
tile_kernels.mhc.mhc_reduce_partials_and_rmsnorm_fwd(om_s, sq_s, om, sq, out, MHC * h, 1e-6, n_splits)
|
|
rout = ref.mhc_pre_norm_fn_ref(residual, fn, normw, 1e-6).view(t, m3)
|
|
torch.testing.assert_close(out, rout, atol=1e-4, rtol=1e-4)
|
|
torch.testing.assert_close(om, om_s.sum(0))
|
|
torch.testing.assert_close(sq, sq_s.sum(0))
|
|
|
|
|
|
@pytest.mark.parametrize('n_splits', [1, 16])
|
|
def test_mhc_pre_big_fuse(n_splits):
|
|
torch.manual_seed(13)
|
|
t, h = 7, 256
|
|
m3 = MHC * (MHC + 2)
|
|
residual = torch.randn(t, MHC, h, dtype=torch.bfloat16)
|
|
fn = torch.randn(m3, MHC * h) * 0.02
|
|
om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual.unsqueeze(0), fn, n_splits)
|
|
gemm_mul, gemm_sq = om_s.view(n_splits, t, m3), sq_s.view(n_splits, t)
|
|
scale, base = torch.rand(3), torch.randn(m3) * 0.1
|
|
post, comb, layer_in = torch.empty(t, MHC), torch.empty(t, MHC * MHC), torch.empty(t, h, dtype=torch.bfloat16)
|
|
tile_kernels.mhc.mhc_pre_big_fuse_fwd(gemm_mul, gemm_sq, scale, base, residual, post, comb, layer_in, 1e-6, 1e-6, 1e-6, 2.0, 20, n_splits=n_splits)
|
|
mixes = ref.mhc_pre_norm_fn_ref(residual.unsqueeze(0), fn, None, 1e-6) # (1, T, M3)
|
|
rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale, base, MHC, 2.0, 1e-6)
|
|
rcomb = ref.sinkhorn_normalize_ref(rcomb, 20, 1e-6)
|
|
rli = ref.mhc_pre_apply_mix_ref(residual.unsqueeze(0).float(), rpre)
|
|
torch.testing.assert_close(post, rpost.view(t, MHC), atol=1e-5, rtol=1e-5)
|
|
torch.testing.assert_close(comb, rcomb.reshape(t, MHC * MHC), atol=1e-5, rtol=1e-4)
|
|
torch.testing.assert_close(layer_in.float(), rli.view(t, h).float(), atol=2e-2, rtol=1e-2)
|