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>
694 lines
31 KiB
Python
694 lines
31 KiB
Python
"""CPU correctness tests for the pure-torch deep_gemm port.
|
|
|
|
Reference implementations and data generators are ported from
|
|
DeepGEMM-Ascend/tests (generators.py, test_layout.py, test_attention.py,
|
|
test_hyperconnection.py), adapted to CPU and small shapes.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
import sys
|
|
import random
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
|
|
|
import torch
|
|
|
|
import deep_gemm
|
|
|
|
torch.manual_seed(0)
|
|
random.seed(0)
|
|
|
|
FAILURES = []
|
|
|
|
# bf16 outputs: the dominant error is rounding D to bf16 (~1e-6 in calc_diff terms).
|
|
BF16_TOL = 1e-5
|
|
|
|
|
|
def check(name, cond, extra=""):
|
|
if cond:
|
|
print(f" [ok] {name}")
|
|
else:
|
|
print(f" [FAIL] {name} {extra}")
|
|
FAILURES.append(name)
|
|
|
|
|
|
def calc_diff(x, y):
|
|
x, y = x.double(), y.double()
|
|
denom = (x * x + y * y).sum().item()
|
|
if denom == 0:
|
|
return 0.0
|
|
return 1 - 2 * (x * y).sum().item() / denom
|
|
|
|
|
|
def ceil_div(x, y):
|
|
return (x + y - 1) // y
|
|
|
|
|
|
def aligned(x, a):
|
|
return ceil_div(x, a) * a
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# generators (ported from DeepGEMM-Ascend/tests/generators.py)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def ceil_to_ue8m0(x):
|
|
bits = x.abs().float().view(torch.int32)
|
|
exp = ((bits // 2 ** 23) % 256) + (bits % 2 ** 23).bool().int()
|
|
return (exp.clamp(1, 254) << 23).view(torch.float32)
|
|
|
|
|
|
def quantize_to_fp4_raw(x):
|
|
ax = x.abs()
|
|
code = torch.zeros_like(x, dtype=torch.uint8)
|
|
for boundary in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0):
|
|
code += (ax > boundary).to(torch.uint8)
|
|
sign = (x < 0) & (code != 0)
|
|
values = torch.tensor((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0), dtype=torch.float32)
|
|
repr_f32 = values[code.long()]
|
|
repr_f32 = torch.where(sign, -repr_f32, repr_f32)
|
|
code = code | (sign.to(torch.uint8) * 8)
|
|
packed = (code[..., ::2] & 0x0F) | ((code[..., 1::2] & 0x0F) << 4)
|
|
return packed.contiguous().view(torch.int8), repr_f32
|
|
|
|
|
|
def quantize_last_dim(x, quant_dtype, gran_k=32):
|
|
"""Quantize along the last dim in gran_k blocks; returns (data, sf_fp32, ref_f32)."""
|
|
k = x.shape[-1]
|
|
assert k % gran_k == 0
|
|
blocks = x.float().reshape(*x.shape[:-1], k // gran_k, gran_k)
|
|
sf_div = 6.0 if quant_dtype == "e2m1" else 448.0
|
|
sf = ceil_to_ue8m0(blocks.abs().amax(dim=-1).clamp(1e-4) / sf_div)
|
|
scaled = (blocks / sf[..., None]).reshape_as(x)
|
|
if quant_dtype == "e2m1":
|
|
data, repr_f32 = quantize_to_fp4_raw(scaled)
|
|
else:
|
|
data = scaled.to(torch.float8_e4m3fn)
|
|
repr_f32 = data.float()
|
|
ref = (repr_f32.reshape(*x.shape[:-1], k // gran_k, gran_k) * sf[..., None]).reshape_as(x)
|
|
return data, sf.contiguous(), ref
|
|
|
|
|
|
def pack_sf_row_major(sf_fp32):
|
|
"""Packed UE8M0 int16 in K-major order (what per_token_cast emits without the col-major view)."""
|
|
exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
|
|
if exponent.shape[-1] % 2:
|
|
exponent = torch.nn.functional.pad(exponent, (0, 1))
|
|
return exponent.contiguous().view(torch.int16)
|
|
|
|
|
|
def pack_sf_reference(sf_fp32, mn, gran_mn):
|
|
"""test_layout.py::reference - fp32 sf -> canonical int16 MN-major."""
|
|
exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
|
|
if exponent.shape[-1] % 2:
|
|
exponent = torch.nn.functional.pad(exponent, (0, 1))
|
|
sf = exponent.contiguous().view(torch.int16)
|
|
return (
|
|
sf.transpose(-2, -1)
|
|
.repeat_interleave(gran_mn, -1)[..., :mn]
|
|
.contiguous()
|
|
.transpose(-2, -1)
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 1. SF layout
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def make_sf(shape, dtype):
|
|
if dtype == torch.float32:
|
|
exponent = torch.randint(1, 255, shape, dtype=torch.int32)
|
|
return (exponent * (1 << 23)).view(torch.float32)
|
|
return torch.randint(-32768, 32768, shape, dtype=torch.int16)
|
|
|
|
|
|
def with_major(sf, mn_major):
|
|
if mn_major and sf.size(-2) > 1 and sf.size(-1) > 1:
|
|
sf = sf.transpose(-2, -1).contiguous().transpose(-2, -1)
|
|
assert sf.stride(-2) == 1
|
|
else:
|
|
assert sf.stride(-1) == 1
|
|
return sf
|
|
|
|
|
|
def test_sf_layout():
|
|
print("test_sf_layout")
|
|
for dtype in (torch.float32, torch.int16):
|
|
for mn_major in (False, True):
|
|
for gran_mn in (1, 32):
|
|
for num_groups in (None, 3):
|
|
mn, k = 128, 256
|
|
sf_k_div = 32 if dtype == torch.float32 else 64
|
|
shape = (ceil_div(mn, gran_mn), ceil_div(k, sf_k_div))
|
|
if num_groups is not None:
|
|
shape = (num_groups, *shape)
|
|
sf = with_major(make_sf(shape, dtype), mn_major)
|
|
out = deep_gemm.transform_sf_into_required_layout(
|
|
sf, mn, k, (gran_mn, gran_mn, 32), num_groups=num_groups, is_sfa=True
|
|
)
|
|
expected = pack_sf_reference(sf, mn, gran_mn) if dtype == torch.float32 else (
|
|
sf.transpose(-2, -1).repeat_interleave(gran_mn, -1)[..., :mn].contiguous().transpose(-2, -1)
|
|
)
|
|
tag = f"dtype={dtype}, mn_major={mn_major}, gran_mn={gran_mn}, groups={num_groups}"
|
|
ok = (
|
|
out.dtype == torch.int16
|
|
and out.shape == expected.shape
|
|
and out.stride(-2) == 1
|
|
and out.stride(-1) == out.size(-2)
|
|
and torch.equal(out, expected)
|
|
)
|
|
check(f"sf layout ({tag})", ok, f"{out.shape=} {expected.shape=} {out.stride()=}")
|
|
# Canonical input is returned unchanged (no copy).
|
|
mn, k = 64, 128
|
|
canonical = torch.zeros((ceil_div(k, 64), mn), dtype=torch.int16).transpose(0, 1)
|
|
out = deep_gemm.transform_sf_into_required_layout(canonical, mn, k, (1, 1, 32), is_sfa=True)
|
|
check("canonical sf passthrough", out.data_ptr() == canonical.data_ptr())
|
|
# Ascend/Neuron col-major sf is unpadded, so get_tma_aligned_size is the identity.
|
|
check("get_tma_aligned_size (unpadded)", deep_gemm.get_tma_aligned_size(5, 2) == 5 and deep_gemm.get_tma_aligned_size(9, 4) == 9)
|
|
# Unpadded col-major int16 sf (strides (1, MN)) is the canonical layout.
|
|
mn2, k2 = 37, 256
|
|
col_major = torch.zeros((ceil_div(k2, 64), mn2), dtype=torch.int16).transpose(0, 1)
|
|
out2 = deep_gemm.transform_sf_into_required_layout(col_major, mn2, k2, (1, 1, 32), is_sfa=False)
|
|
check("unpadded col-major sf passthrough", out2.data_ptr() == col_major.data_ptr() and out2.stride() == (1, mn2))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 2. dense GEMMs
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_bf16_gemm():
|
|
print("test_bf16_gemm")
|
|
m, n, k = 65, 128, 256
|
|
a = torch.randn(m, k, dtype=torch.bfloat16)
|
|
b = torch.randn(n, k, dtype=torch.bfloat16)
|
|
ref = a.float() @ b.float().T
|
|
d = torch.empty(m, n, dtype=torch.bfloat16)
|
|
deep_gemm.bf16_gemm_nt(a, b, d, None)
|
|
check("bf16_gemm_nt", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))
|
|
|
|
d32 = torch.empty(m, n, dtype=torch.float32)
|
|
deep_gemm.cublaslt_gemm_nt(a, b, d32, None)
|
|
check("cublaslt_gemm_nt == bf16 ref", calc_diff(d32, ref) < 1e-10)
|
|
|
|
# accumulate: c aliases d
|
|
acc = torch.randn(m, n, dtype=torch.float32)
|
|
d_acc = acc.clone()
|
|
deep_gemm.bf16_gemm_nt(a, b, d_acc, d_acc)
|
|
check("bf16_gemm_nt accumulate", calc_diff(d_acc, ref + acc) < 1e-10)
|
|
|
|
# separate [1, N] bias
|
|
bias = torch.randn(1, n, dtype=torch.float32)
|
|
d_bias = torch.empty(m, n, dtype=torch.float32)
|
|
deep_gemm.bf16_gemm_nt(a, b, d_bias, bias)
|
|
check("bf16_gemm_nt bias", calc_diff(d_bias, ref + bias) < 1e-10)
|
|
|
|
# tn / nn variants
|
|
at = a.T.contiguous()
|
|
d_tn = torch.empty(m, n, dtype=torch.float32)
|
|
deep_gemm.bf16_gemm_tn(at, b.T.contiguous(), d_tn, None)
|
|
check("bf16_gemm_tn", calc_diff(d_tn, ref) < 1e-10)
|
|
d_nn = torch.empty(m, n, dtype=torch.float32)
|
|
deep_gemm.bf16_gemm_nn(a, b.T.contiguous(), d_nn, None)
|
|
check("bf16_gemm_nn", calc_diff(d_nn, ref) < 1e-10)
|
|
# strided output view
|
|
storage = torch.empty(m, n + 7, dtype=torch.bfloat16)
|
|
view = storage[:, :n]
|
|
deep_gemm.bf16_gemm_nt(a, b, view, None)
|
|
check("bf16_gemm_nt strided out", calc_diff(view.float(), ref) < BF16_TOL, calc_diff(view.float(), ref))
|
|
|
|
|
|
def run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode, m=64, n=128, k=256, acc=False, gran_b=1):
|
|
a_bf16 = torch.randn(m, k, dtype=torch.bfloat16)
|
|
b_bf16 = torch.randn(n, k, dtype=torch.bfloat16)
|
|
a_data, a_sf, a_ref = quantize_last_dim(a_bf16, a_fmt)
|
|
b_data, b_sf, b_ref = quantize_last_dim(b_bf16, b_fmt)
|
|
if sf_mode == "packed":
|
|
a_sf_in = pack_sf_reference(a_sf, m, 1)
|
|
b_sf_in = pack_sf_reference(b_sf, n, 1)
|
|
elif sf_mode == "packed_row_major":
|
|
a_sf_in = pack_sf_row_major(a_sf)
|
|
b_sf_in = pack_sf_row_major(b_sf)
|
|
else:
|
|
a_sf_in, b_sf_in = a_sf, b_sf
|
|
ref = a_ref @ b_ref.T
|
|
d = torch.empty(m, n, dtype=torch.bfloat16)
|
|
c = None
|
|
if acc:
|
|
base = torch.randn(m, n, dtype=torch.bfloat16)
|
|
d.copy_(base)
|
|
c = d
|
|
ref = ref + base.float()
|
|
deep_gemm.fp8_fp4_gemm_nt((a_data, a_sf_in), (b_data, b_sf_in), d, c, recipe=(1, 1, 32))
|
|
return calc_diff(d.float(), ref)
|
|
|
|
|
|
def test_fp8_fp4_gemm():
|
|
print("test_fp8_fp4_gemm")
|
|
for a_fmt, b_fmt in (("e4m3", "e4m3"), ("e4m3", "e2m1"), ("e2m1", "e2m1")):
|
|
for sf_mode in ("fp32", "packed", "packed_row_major"):
|
|
diff = run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode)
|
|
check(f"fp8_fp4_gemm_nt ({a_fmt}x{b_fmt}, sf={sf_mode})", diff < BF16_TOL, f"{diff=}")
|
|
diff = run_fp8_fp4_gemm_case("e4m3", "e2m1", "packed", acc=True)
|
|
check("fp8_fp4_gemm_nt accumulate", diff < BF16_TOL, f"{diff=}")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 3. m-grouped contiguous GEMMs
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def make_psum_layout(expected_m, num_groups, alignment):
|
|
layout, psum = [], 0
|
|
for _ in range(num_groups):
|
|
size = max(int(random.uniform(0.7, 1.3) * expected_m), 0)
|
|
layout.append(psum + size)
|
|
psum = aligned(psum + size, alignment)
|
|
return psum, layout
|
|
|
|
|
|
def iter_psum_slice(layout, alignment):
|
|
last = 0
|
|
for ptr in layout:
|
|
yield slice(last, ptr)
|
|
last = aligned(ptr, alignment)
|
|
|
|
|
|
def test_m_grouped():
|
|
print("test_m_grouped")
|
|
alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()
|
|
check("mk alignment == 256", alignment == 256)
|
|
deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
|
|
num_groups, expected_m, n, k = 3, 200, 64, 128
|
|
total_m, layout = make_psum_layout(expected_m, num_groups, alignment)
|
|
grouped = torch.tensor(layout, dtype=torch.int32)
|
|
spans = list(iter_psum_slice(layout, alignment))
|
|
|
|
# bf16
|
|
a = torch.randn(total_m, k, dtype=torch.bfloat16)
|
|
for span in spans: # zero the inter-group padding
|
|
pass
|
|
pad_mask = torch.ones(total_m, dtype=torch.bool)
|
|
for span in spans:
|
|
pad_mask[span] = False
|
|
a[pad_mask] = 0
|
|
b = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
|
|
ref = torch.zeros(total_m, n, dtype=torch.float32)
|
|
for i, span in enumerate(spans):
|
|
ref[span] = a[span].float() @ b[i].float().T
|
|
d = torch.empty(total_m, n, dtype=torch.bfloat16)
|
|
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped, use_psum_layout=True, expected_m_for_psum_layout=expected_m)
|
|
check("m_grouped_bf16_gemm_nt_contiguous", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))
|
|
check("m_grouped bf16 zero padding", bool((d.float()[pad_mask] == 0).all()))
|
|
|
|
# fp8 x fp4 (the DeepInfer mxfp4 MoE path)
|
|
a_bf16 = torch.randn(total_m, k, dtype=torch.bfloat16)
|
|
a_bf16[pad_mask] = 0
|
|
a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
|
|
b_bf16 = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
|
|
b_data, b_sf, b_ref = quantize_last_dim(b_bf16, "e2m1")
|
|
a_sf_packed = pack_sf_reference(a_sf, total_m, 1)
|
|
b_sf_packed = pack_sf_reference(b_sf, n, 1)
|
|
ref = torch.zeros(total_m, n, dtype=torch.float32)
|
|
for i, span in enumerate(spans):
|
|
ref[span] = a_ref[span] @ b_ref[i].T
|
|
d = torch.empty(total_m, n, dtype=torch.bfloat16)
|
|
deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
|
|
(a_data, a_sf_packed), (b_data, b_sf_packed), d, grouped,
|
|
recipe_a=(1, 32), recipe_b=(1, 32), use_psum_layout=True, expected_m_for_psum_layout=expected_m,
|
|
)
|
|
diff = calc_diff(d.float(), ref)
|
|
check("m_grouped_fp8_fp4_gemm_nt_contiguous", diff < BF16_TOL, f"{diff=}")
|
|
check("m_grouped fp8xfp4 zero padding", bool((d.float()[pad_mask] == 0).all()))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 4. einsum
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_einsum():
|
|
print("test_einsum")
|
|
b_dim, h, r, dd = 7, 4, 128, 64
|
|
a_bf16 = torch.randn(b_dim, h, r, dtype=torch.bfloat16)
|
|
w_bf16 = torch.randn(h, dd, r, dtype=torch.bfloat16)
|
|
ref = torch.einsum("bhr,hdr->bhd", a_bf16.float(), w_bf16.float())
|
|
out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
|
|
deep_gemm.einsum("bhr,hdr->bhd", a_bf16, w_bf16, out, None)
|
|
check("einsum bf16 bhr,hdr->bhd", calc_diff(out.float(), ref) < BF16_TOL, calc_diff(out.float(), ref))
|
|
|
|
a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
|
|
w_data, w_sf, w_ref = quantize_last_dim(w_bf16, "e4m3")
|
|
a_sf_packed = pack_sf_reference(a_sf, a_sf.shape[-2], 1) # [b, h, r/64] canonical per b
|
|
w_sf_packed = pack_sf_reference(w_sf, w_sf.shape[-2], 1)
|
|
ref = torch.einsum("bhr,hdr->bhd", a_ref, w_ref)
|
|
out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
|
|
deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf_packed), (w_data, w_sf_packed), out, recipe=(1, 1, 32))
|
|
diff = calc_diff(out.float(), ref)
|
|
check("fp8_einsum bhr,hdr->bhd", diff < BF16_TOL, f"{diff=}")
|
|
# the fp32-sf input form must agree too
|
|
out2 = torch.empty_like(out)
|
|
deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf), (w_data, w_sf), out2, recipe=(1, 1, 32))
|
|
check("fp8_einsum fp32 sf == packed sf", torch.equal(out, out2))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 5. MQA logits (ported from tests/test_attention.py)
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def ceil_to_ue8m0_with_byte(x):
|
|
bits = x.abs().float().view(torch.int32)
|
|
exp = (((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int()).clamp(1, 254)
|
|
return (exp << 23).view(torch.float32), exp.to(torch.uint8)
|
|
|
|
|
|
def quantize_mqa_tensor(x, use_fp4):
|
|
sf_div = 6.0 if use_fp4 else 448.0
|
|
sf_shape = (*x.shape[:-1], x.shape[-1] // 32, 32)
|
|
sf, sf_byte = ceil_to_ue8m0_with_byte(x.float().reshape(sf_shape).abs().amax(dim=-1).clamp(1e-4) / sf_div)
|
|
x_scaled = (x.float().reshape(sf_shape) / sf[..., None]).reshape_as(x)
|
|
if use_fp4:
|
|
x_quant, x_repr = quantize_to_fp4_raw(x_scaled)
|
|
else:
|
|
x_quant = x_scaled.to(torch.float8_e4m3fn)
|
|
x_repr = x_quant.float()
|
|
x_ref = (x_repr.reshape(sf_shape) * sf[..., None]).reshape_as(x)
|
|
sf_packed = sf_byte[..., ::2].to(torch.int16) | (sf_byte[..., 1::2].to(torch.int16) << 8)
|
|
return x_quant, sf_packed.contiguous(), x_ref
|
|
|
|
|
|
def ref_mqa_logits(q, kv, weights, ks, ke):
|
|
seq_len, num_heads, _ = q.shape
|
|
seq_len_kv = kv.shape[0]
|
|
score = torch.einsum("qhd,kd->qhk", q.float(), kv.float())
|
|
out = (score.relu() * weights.float()[:, :, None]).sum(dim=1)
|
|
pos = torch.arange(seq_len_kv)
|
|
mask = (pos[None, :] >= ks[:, None]) & (pos[None, :] < ke[:, None])
|
|
return out.masked_fill(~mask, float("-inf"))
|
|
|
|
|
|
def test_mqa_logits(use_fp4):
|
|
name = "fp4" if use_fp4 else "fp8"
|
|
seq_len, seq_len_kv, num_heads, head_dim = 9, 512, 8, 128
|
|
q_bf16 = torch.randn(seq_len, num_heads, head_dim, dtype=torch.bfloat16)
|
|
kv_bf16 = torch.randn(seq_len_kv, head_dim, dtype=torch.bfloat16)
|
|
weights = torch.randn(seq_len, num_heads, dtype=torch.bfloat16)
|
|
a = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
|
|
b = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
|
|
ks, ke = torch.minimum(a, b), torch.maximum(a, b)
|
|
max_seqlen_k = max(1, int((ke - ks).max()))
|
|
|
|
q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
|
|
kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
|
|
ref = ref_mqa_logits(q_ref, kv_ref, weights, ks, ke).to(torch.bfloat16)
|
|
|
|
out = deep_gemm.fp8_fp4_mqa_logits((q, q_sf), (kv, kv_sf), weights, ks, ke, max_seqlen_k)
|
|
check(f"mqa logits {name} dtype/stride", out.dtype == torch.bfloat16
|
|
and out.shape == (seq_len, max_seqlen_k)
|
|
and out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
|
|
dense = torch.full((seq_len, seq_len_kv), float("-inf"), dtype=torch.bfloat16)
|
|
for i, (s, e) in enumerate(zip(ks.tolist(), ke.tolist())):
|
|
dense[i, s:e] = out[i, : e - s]
|
|
neg = torch.isinf(ref.float()) & (ref.float() < 0)
|
|
check(f"mqa logits {name} -inf mask", torch.equal(torch.isinf(dense.float()) & (dense.float() < 0), neg))
|
|
diff = calc_diff(dense.float().masked_fill(neg, 0), ref.float().masked_fill(neg, 0))
|
|
check(f"mqa logits {name} values", diff < 5e-4, f"{diff=}")
|
|
return q, q_sf, q_ref, kv, kv_sf, kv_ref, weights, ks, ke
|
|
|
|
|
|
def make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4):
|
|
data_bytes = 64 if use_fp4 else 128
|
|
sf_bytes = 4
|
|
token_bytes = data_bytes + sf_bytes
|
|
data = kv.contiguous().view(torch.uint8).reshape(num_pages, page_kv * data_bytes)
|
|
sf = kv_sf.contiguous().view(torch.uint8).reshape(num_pages, page_kv * sf_bytes)
|
|
fused = torch.cat([data, sf], dim=1)
|
|
return fused.reshape(num_pages, page_kv, 1, token_bytes)
|
|
|
|
|
|
def ref_paged_mqa_logits(q, kv, weights, context_lens, block_table, max_context_len):
|
|
num_q = q.shape[0]
|
|
page_size = kv.shape[1]
|
|
out = torch.zeros(num_q, max_context_len)
|
|
for q_idx, (context_len,) in enumerate(context_lens.tolist()):
|
|
pages = ceil_div(context_len, page_size)
|
|
keys = kv[block_table[q_idx, :pages].long()].flatten(0, 1).float()
|
|
scores = torch.einsum("hd,kd->hk", q[q_idx, 0].float(), keys)
|
|
logits = (scores.relu() * weights[q_idx, :, None].float()).sum(0)
|
|
out[q_idx, :context_len] = logits[:context_len]
|
|
return out
|
|
|
|
|
|
def test_paged_mqa_logits(use_fp4):
|
|
name = "fp4" if use_fp4 else "fp8"
|
|
batch, num_heads, head_dim, page_kv = 3, 8, 128, 64
|
|
context_lens_host = [70, 128, 33]
|
|
pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
|
|
num_pages = sum(pages_per_request)
|
|
q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
|
|
kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
|
|
weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
|
|
q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
|
|
kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
|
|
fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)
|
|
|
|
perm = torch.randperm(num_pages, dtype=torch.int32)
|
|
max_pages = max(pages_per_request)
|
|
block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
|
|
offset = 0
|
|
for i, count in enumerate(pages_per_request):
|
|
block_table[i, :count] = perm[offset:offset + count]
|
|
offset += count
|
|
context_lens = torch.tensor(context_lens_host, dtype=torch.int32).view(-1, 1)
|
|
indices = torch.arange(batch, dtype=torch.int32)
|
|
max_context_len = max(context_lens_host)
|
|
|
|
# Reference uses the (permuted) page order through the block table.
|
|
ref = ref_paged_mqa_logits(q_ref, kv_ref, weights, context_lens, block_table, max_context_len)
|
|
meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens, num_heads, indices)
|
|
out = deep_gemm.fp8_fp4_paged_mqa_logits(
|
|
(q, q_sf), fused, weights, context_lens, block_table, meta, max_context_len, indices, False
|
|
)
|
|
valid = torch.arange(max_context_len)[None] < context_lens.view(-1, 1)
|
|
diff = calc_diff(out.float().masked_fill(~valid, 0), ref.masked_fill(~valid, 0))
|
|
check(f"paged mqa logits {name}", diff < 5e-4, f"{diff=}")
|
|
check(f"paged mqa logits {name} stride", out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
|
|
return (q, q_sf), fused, weights, context_lens, block_table, indices, max_context_len, out
|
|
|
|
|
|
def test_sparse_mqa_logits(use_fp4):
|
|
"""Sparse logits must match the dense logits at the selected tokens."""
|
|
name = "fp4" if use_fp4 else "fp8"
|
|
num_q, num_kv, num_heads, head_dim = 6, 512, 8, 128
|
|
sparse_block_kv, capacity = 8, 16
|
|
q_bf16 = torch.randn(num_q, num_heads, head_dim, dtype=torch.bfloat16)
|
|
kv_bf16 = torch.randn(num_kv, head_dim, dtype=torch.bfloat16)
|
|
weights = torch.randn(num_q, num_heads, dtype=torch.bfloat16)
|
|
q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
|
|
kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
|
|
|
|
rng = random.Random(1)
|
|
starts, ends, blocks = [], [], []
|
|
for _ in range(num_q):
|
|
start = rng.randrange(0, 4) * sparse_block_kv + rng.randrange(0, sparse_block_kv)
|
|
end = rng.randrange(start + 1, num_kv)
|
|
first_block = start // sparse_block_kv
|
|
available = ceil_div(max(0, end - start), sparse_block_kv)
|
|
count = min(available, capacity)
|
|
chosen = sorted(rng.sample(range(first_block, first_block + available), count))
|
|
chosen += [chosen[-1] if chosen else 0] * (capacity - count)
|
|
starts.append(start)
|
|
ends.append(end)
|
|
blocks.append(chosen)
|
|
ks = torch.tensor(starts, dtype=torch.int32)
|
|
ke = torch.tensor(ends, dtype=torch.int32)
|
|
sparse_blocks = torch.tensor(blocks, dtype=torch.int32)
|
|
|
|
meta = deep_gemm.get_sparse_mqa_logits_metadata(ks, ke, num_kv, sparse_blocks, q.dtype, sparse_block_kv)
|
|
sparse = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
|
|
check(f"sparse mqa {name} shape/stride",
|
|
sparse.shape == (num_q, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
|
|
and sparse.stride(0) * sparse.element_size() % 512 == 0)
|
|
|
|
full = deep_gemm.fp8_fp4_mqa_logits(
|
|
(q, q_sf), (kv, kv_sf), weights, ks, ke, max(1, int((ke - ks).max()))
|
|
)
|
|
offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
|
|
token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + (ks % sparse_block_kv)[:, None, None] + offsets).flatten(1).long()
|
|
num_blocks = torch.tensor([min(capacity, ceil_div(max(0, e - s), sparse_block_kv)) for s, e in zip(starts, ends)])
|
|
valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
|
|
valid &= token_indices >= ks[:, None]
|
|
valid &= token_indices < ke[:, None]
|
|
full_cols = (token_indices - ks[:, None]).clamp(0, full.shape[1] - 1)
|
|
check(f"sparse mqa {name} == dense logits",
|
|
torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))
|
|
|
|
# Determinism / reuse
|
|
sparse2 = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
|
|
check(f"sparse mqa {name} repeatable", torch.equal(sparse[valid], sparse2[valid]))
|
|
|
|
|
|
def test_paged_sparse_mqa_logits(use_fp4):
|
|
name = "fp4" if use_fp4 else "fp8"
|
|
num_heads, head_dim, page_kv = 8, 128, 64
|
|
sparse_block_kv, capacity = 8, 16
|
|
context_lens_host = [70, 128, 33]
|
|
batch = len(context_lens_host)
|
|
pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
|
|
num_pages = sum(pages_per_request)
|
|
q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
|
|
kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
|
|
weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
|
|
q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
|
|
kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
|
|
fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)
|
|
|
|
perm = torch.randperm(num_pages, dtype=torch.int32)
|
|
max_pages = max(pages_per_request)
|
|
block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
|
|
offset = 0
|
|
for i, count in enumerate(pages_per_request):
|
|
block_table[i, :count] = perm[offset:offset + count]
|
|
offset += count
|
|
context_lens = torch.tensor(context_lens_host, dtype=torch.int32)
|
|
indices = torch.arange(batch, dtype=torch.int32)
|
|
|
|
rng = random.Random(2)
|
|
blocks = []
|
|
for context_len in context_lens_host:
|
|
available = ceil_div(context_len, sparse_block_kv)
|
|
count = min(available, capacity)
|
|
chosen = sorted(rng.sample(range(available), count))
|
|
blocks.append(chosen + [chosen[-1]] * (capacity - count))
|
|
sparse_blocks = torch.tensor(blocks, dtype=torch.int32)
|
|
|
|
meta = deep_gemm.get_paged_sparse_mqa_logits_metadata(
|
|
context_lens, block_table, indices, page_kv, sparse_blocks, q.dtype, sparse_block_kv
|
|
)
|
|
sparse = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
|
|
check(f"paged sparse {name} shape/stride",
|
|
sparse.shape == (batch, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
|
|
and sparse.stride(0) * sparse.element_size() % 512 == 0)
|
|
|
|
dense_meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens.view(-1, 1), num_heads, indices)
|
|
full = deep_gemm.fp8_fp4_paged_mqa_logits(
|
|
(q, q_sf), fused, weights, context_lens.view(-1, 1), block_table, dense_meta,
|
|
max(context_lens_host), indices, False
|
|
)
|
|
offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
|
|
token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + offsets).flatten(1).long()
|
|
num_blocks = torch.tensor([min(capacity, ceil_div(c, sparse_block_kv)) for c in context_lens_host])
|
|
valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
|
|
valid &= token_indices < context_lens[:, None]
|
|
full_cols = token_indices.clamp(0, full.shape[1] - 1)
|
|
check(f"paged sparse {name} == dense logits",
|
|
torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))
|
|
|
|
# REUSE: the KV snapshot taken on the first call survives a KV overwrite.
|
|
first = sparse.clone()
|
|
fused.zero_()
|
|
again = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
|
|
check(f"paged sparse {name} reuses the KV snapshot", torch.equal(first[valid], again[valid]))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 6. hyperconnection prenorm
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_hc_prenorm():
|
|
print("test_hc_prenorm")
|
|
m, n, k = 65, 24, 256
|
|
a = torch.randn(m, k, dtype=torch.bfloat16)
|
|
b = torch.randn(n, k, dtype=torch.float32)
|
|
d = torch.empty(m, n, dtype=torch.float32)
|
|
sqr_sum = torch.empty(m, dtype=torch.float32)
|
|
deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum)
|
|
ref_d = a.float() @ b.T
|
|
ref_sqr = a.float().square().sum(-1)
|
|
check("tf32_hc_prenorm_gemm d", calc_diff(d, ref_d) < 1e-6, calc_diff(d, ref_d))
|
|
check("tf32_hc_prenorm_gemm sqr_sum", calc_diff(sqr_sum, ref_sqr) < 1e-6)
|
|
d_split = torch.empty(1, m, n, dtype=torch.float32)
|
|
sqr_split = torch.empty(1, m, dtype=torch.float32)
|
|
deep_gemm.tf32_hc_prenorm_gemm(a, b, d_split, sqr_split, num_splits=1)
|
|
check("tf32_hc_prenorm_gemm num_splits=1", torch.equal(d_split[0], d) and torch.equal(sqr_split[0], sqr_sum))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 7. API surface
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def test_api_surface():
|
|
print("test_api_surface")
|
|
required = [
|
|
"SymmBuffer", "bf16_gemm_nt", "bf16_mega_gate", "bf16_mega_moe", "cublaslt_gemm_nt",
|
|
"fp8_einsum", "fp8_fp4_gemm_nt", "fp8_fp4_mega_moe", "fp8_fp4_mqa_logits",
|
|
"fp8_fp4_paged_mqa_logits", "fp8_fp4_paged_sparse_mqa_logits", "fp8_fp4_sparse_mqa_logits",
|
|
"get_block_m_for_mega_moe", "get_mk_alignment_for_contiguous_layout", "get_num_sms",
|
|
"get_paged_mqa_logits_metadata", "get_paged_sparse_mqa_logits_metadata",
|
|
"get_sparse_mqa_logits_metadata", "get_theoretical_mk_alignment_for_contiguous_layout",
|
|
"get_tma_aligned_size", "m_grouped_bf16_gemm_nt_contiguous",
|
|
"m_grouped_fp8_fp4_gemm_nt_contiguous", "mega_mhc", "set_mk_alignment_for_contiguous_layout",
|
|
"set_num_sms", "set_pdl", "transform_sf_into_required_layout",
|
|
"transform_weights_for_mega_moe", "use_deterministic_algorithms", "tf32_hc_prenorm_gemm",
|
|
]
|
|
missing = [name for name in required if not hasattr(deep_gemm, name)]
|
|
check("all required symbols exported", not missing, f"missing={missing}")
|
|
for name in ("bf16_mega_gate", "bf16_mega_moe", "fp8_fp4_mega_moe", "mega_mhc",
|
|
"transform_weights_for_mega_moe", "get_block_m_for_mega_moe"):
|
|
try:
|
|
getattr(deep_gemm, name)()
|
|
check(f"{name} raises NotImplementedError", False)
|
|
except NotImplementedError:
|
|
check(f"{name} raises NotImplementedError", True)
|
|
except Exception as exc: # noqa: BLE001
|
|
check(f"{name} raises NotImplementedError", False, repr(exc))
|
|
prev = deep_gemm.get_num_sms()
|
|
check("default num_sms == 8 (tile_kernels.get_device_num_sms)", prev == 8, f"{prev=}")
|
|
deep_gemm.set_num_sms(prev + 4)
|
|
check("set/get_num_sms", deep_gemm.get_num_sms() == prev + 4)
|
|
deep_gemm.set_num_sms(prev)
|
|
deep_gemm.use_deterministic_algorithms(True)
|
|
check("use_deterministic_algorithms", deep_gemm.get_deterministic_algorithms() is True)
|
|
deep_gemm.use_deterministic_algorithms(False)
|
|
deep_gemm.set_pdl(True)
|
|
|
|
class _Group:
|
|
def size(self):
|
|
return 2
|
|
|
|
def rank(self):
|
|
return 0
|
|
|
|
buf = deep_gemm.SymmBuffer(
|
|
_Group(), num_experts=8, num_max_tokens_per_rank=16, num_topk=4, hidden=256,
|
|
intermediate_hidden=128, num_shared_experts=1, mma_type="fp8xfp4", activation="swiglu",
|
|
device="cpu",
|
|
)
|
|
check("SymmBuffer views", buf.x.shape == (16, 256) and buf.x_sf.shape == (16, 4)
|
|
and buf.topk_idx.shape == (16, 4) and buf.num_bytes > 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_sf_layout()
|
|
test_bf16_gemm()
|
|
test_fp8_fp4_gemm()
|
|
test_m_grouped()
|
|
test_einsum()
|
|
print("test_mqa_logits")
|
|
test_mqa_logits(False)
|
|
test_mqa_logits(True)
|
|
print("test_paged_mqa_logits")
|
|
test_paged_mqa_logits(False)
|
|
test_paged_mqa_logits(True)
|
|
print("test_sparse_mqa_logits")
|
|
test_sparse_mqa_logits(False)
|
|
test_sparse_mqa_logits(True)
|
|
print("test_paged_sparse_mqa_logits")
|
|
test_paged_sparse_mqa_logits(False)
|
|
test_paged_sparse_mqa_logits(True)
|
|
test_hc_prenorm()
|
|
test_api_surface()
|
|
print()
|
|
if FAILURES:
|
|
print(f"FAILED ({len(FAILURES)}): {FAILURES}")
|
|
sys.exit(1)
|
|
print("ALL TESTS PASSED")
|