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>
317 lines
15 KiB
Python
317 lines
15 KiB
Python
"""CPU correctness tests for the pure-torch flash_mla vendor package.
|
|
|
|
Run: python -I test_flash_mla.py (PYTHONPATH must contain the vendor dir)
|
|
"""
|
|
|
|
import sys
|
|
|
|
import torch
|
|
|
|
import flash_mla
|
|
from flash_mla import flash_mla_sparse_fwd, flash_mla_with_kvcache, get_mla_metadata
|
|
from flash_mla.flash_mla_interface import FlashMLASchedMeta
|
|
|
|
import ref_attn
|
|
import ref_quant
|
|
|
|
torch.manual_seed(0)
|
|
FAILS = []
|
|
|
|
|
|
def check(name, got, exp, atol=2e-2, rtol=2e-2, exact=False):
|
|
assert got.shape == exp.shape, f"{name}: shape {tuple(got.shape)} != {tuple(exp.shape)}"
|
|
g, e = got.float(), exp.float()
|
|
if exact:
|
|
ok = torch.equal(g, e)
|
|
msg = "bitwise"
|
|
else:
|
|
same_inf = ((g == e) | ((g != g) & (e != e))).float()
|
|
finite = torch.isfinite(g) & torch.isfinite(e)
|
|
ok = bool(torch.all(torch.where(finite, (g - e).abs() <= atol + rtol * e.abs(), same_inf.bool())))
|
|
msg = f"max abs diff {(g - e).abs()[finite].max().item() if finite.any() else 0.0:.3e}"
|
|
print(f"[{'ok ' if ok else 'FAIL'}] {name} ({msg})")
|
|
if not ok:
|
|
FAILS.append(name)
|
|
|
|
|
|
def make_indices(rows, topk, s_kv, invalid_frac=0.2, invalid_value=-1):
|
|
idx = torch.randint(0, s_kv, (rows, topk), dtype=torch.int32)
|
|
mask = torch.rand(rows, topk) < invalid_frac
|
|
idx = torch.where(mask, torch.full_like(idx, invalid_value), idx)
|
|
return idx
|
|
|
|
|
|
# ---------------------------------------------------------------- prefill
|
|
def test_prefill(h_q=64, s_q=7, s_kv=41, topk=24, with_sink=True, with_len=True, d=512):
|
|
q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
|
|
# also exercise the ">= s_kv" invalid convention of the prefill kernel
|
|
idx[0, 0, 0] = s_kv + 3
|
|
sink = torch.randn(h_q) if with_sink else None
|
|
tl = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32) if with_len else None
|
|
sm = d**-0.5
|
|
|
|
out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, sm, 512, sink, tl)
|
|
r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q, kv, idx, sm, tl, sink)
|
|
tag = f"prefill h{h_q} sink={with_sink} len={with_len}"
|
|
check(f"{tag} out", out, r_out)
|
|
check(f"{tag} max_logits", ml, r_ml)
|
|
check(f"{tag} lse", lse, r_lse)
|
|
assert out.dtype == torch.bfloat16 and lse.dtype == torch.float32 and ml.dtype == torch.float32
|
|
|
|
|
|
def test_prefill_lonely():
|
|
h_q, s_q, s_kv, topk, d = 64, 3, 16, 8, 512
|
|
q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
|
|
idx = torch.full((s_q, 1, topk), -1, dtype=torch.int32)
|
|
out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5)
|
|
check("prefill lonely out==0", out, torch.zeros_like(out), exact=True)
|
|
check("prefill lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)
|
|
check("prefill lonely max_logits==-inf", ml, torch.full_like(ml, float("-inf")), exact=True)
|
|
|
|
|
|
# ---------------------------------------------------------------- kv cache
|
|
def test_kv_dequant():
|
|
for fmt in ("v41_fp8", "v41_fp4"):
|
|
raw = (torch.randn(5, 8, 1, 512) / 4).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, fmt)
|
|
exp = ref_quant.dequantize_k_cache(cache, fmt).view(5, 8, 512)
|
|
got = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1), fmt)
|
|
check(f"dequant {fmt}", got, exp, exact=True)
|
|
# relative error against the bf16 source
|
|
err = ((got.float() - raw.view(5, 8, 512).float()).abs() / raw.view(5, 8, 512).float().abs().clamp_min(1e-3)).median()
|
|
print(f" {fmt} median rel err vs bf16 source: {err.item():.4f}")
|
|
# byte-view equivalence: int8 / float8 views of the same bytes must dequantize identically
|
|
got_i8 = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1).view(torch.int8), fmt)
|
|
check(f"dequant {fmt} int8 view", got_i8, exp, exact=True)
|
|
|
|
|
|
# ---------------------------------------------------------------- decode
|
|
def test_decode(h_q=64, b=5, topk=20, block=4, nblocks=9, with_sink=True, with_len=True, extra=False, strided=False):
|
|
d, s_q = 512, 1
|
|
raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
|
|
if strided:
|
|
# non-contiguous block stride, like deepinfer/worker/cache_engine's as_strided views
|
|
pad = torch.zeros(nblocks, block * 528 + 64, dtype=torch.uint8)
|
|
pad[:, : block * 528] = cache.reshape(nblocks, -1)
|
|
cache = pad.as_strided((nblocks, block, 1, 528), (pad.stride(0), 528, 528, 1))
|
|
kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
|
|
|
|
q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
|
|
sink = torch.randn(h_q) if with_sink else None
|
|
tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32) if with_len else None
|
|
sm = d**-0.5
|
|
|
|
scopes = [(kv_bf16, idx, tl)]
|
|
e_cache = e_idx = e_tl = None
|
|
if extra:
|
|
e_raw = (torch.randn(6, block, 1, d) / 8).to(torch.bfloat16)
|
|
e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp4")
|
|
e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp4").view(-1, d)
|
|
e_topk = 12
|
|
e_idx = make_indices(b * s_q, e_topk, 6 * block).view(b, s_q, e_topk)
|
|
e_tl = torch.randint(1, e_topk + 1, (b,), dtype=torch.int32)
|
|
scopes.append((e_kv, e_idx, e_tl))
|
|
|
|
meta, none = get_mla_metadata()
|
|
assert none is None and isinstance(meta, FlashMLASchedMeta)
|
|
out, lse = flash_mla_with_kvcache(
|
|
q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
|
|
tile_scheduler_metadata=meta, softmax_scale=sm, causal=False, is_fp8_kvcache=True,
|
|
indices=idx, attn_sink=sink, extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
|
|
topk_length=tl, extra_topk_length=e_tl,
|
|
)
|
|
r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, scopes, sm, sink)
|
|
tag = f"decode h{h_q} sink={with_sink} len={with_len} extra={extra} strided={strided}"
|
|
check(f"{tag} out", out, r_out)
|
|
check(f"{tag} lse", lse, r_lse)
|
|
assert meta.have_initialized and meta.config.topk == topk and meta.config.page_block_size == block
|
|
# reuse with a different topk must be rejected
|
|
try:
|
|
flash_mla_with_kvcache(
|
|
q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
|
|
tile_scheduler_metadata=meta, softmax_scale=sm, indices=idx[..., :-1],
|
|
)
|
|
print("[FAIL] sched_meta reuse check")
|
|
FAILS.append("sched_meta reuse")
|
|
except AssertionError:
|
|
print("[ok ] sched_meta reuse rejected on topk change")
|
|
|
|
|
|
def test_decode_multi_token():
|
|
"""s_q > 1 (MTP/DFlash style): topk_length is per request, broadcast over s_q."""
|
|
d, h_q, b, s_q, topk, block, nblocks = 512, 64, 3, 4, 16, 8, 7
|
|
raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
|
|
kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
|
|
q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
|
|
tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
|
|
sm = d**-0.5
|
|
out, lse = flash_mla_with_kvcache(
|
|
q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
|
|
tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=sm, indices=idx, topk_length=tl,
|
|
)
|
|
r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv_bf16, idx, tl)], sm)
|
|
check("decode s_q=4 out", out, r_out)
|
|
check("decode s_q=4 lse", lse, r_lse)
|
|
|
|
|
|
def test_decode_lonely():
|
|
d, h_q, b, topk, block, nblocks = 512, 64, 2, 8, 4, 4
|
|
raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
|
|
q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
|
|
idx = torch.full((b, 1, topk), -1, dtype=torch.int32)
|
|
out, lse = flash_mla_with_kvcache(
|
|
q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
|
|
tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
|
|
attn_sink=torch.randn(h_q),
|
|
)
|
|
check("decode lonely out==0", out, torch.zeros_like(out), exact=True)
|
|
check("decode lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)
|
|
|
|
|
|
def test_sink_infinities():
|
|
d, h_q, s_q, s_kv, topk = 512, 64, 4, 20, 8
|
|
q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
|
|
idx = torch.randint(0, s_kv, (s_q, 1, topk), dtype=torch.int32)
|
|
sink = torch.full((h_q,), float("-inf"))
|
|
sink[: h_q // 2] = float("inf")
|
|
out, _, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, sink)
|
|
base, _, base_lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, None)
|
|
check("sink=+inf -> out 0", out[:, : h_q // 2], torch.zeros_like(out[:, : h_q // 2]), exact=True)
|
|
check("sink=-inf -> unchanged", out[:, h_q // 2 :], base[:, h_q // 2 :], atol=1e-2, rtol=1e-2)
|
|
check("sink does not change lse", lse, base_lse, exact=True)
|
|
|
|
|
|
def test_dense_unsupported():
|
|
try:
|
|
flash_mla.flash_attn_varlen_func(None, None, None, None, None, 0, 0)
|
|
FAILS.append("dense should raise")
|
|
print("[FAIL] dense entry point should raise")
|
|
except RuntimeError as e:
|
|
print(f"[ok ] dense entry point raises: {str(e)[:60]}...")
|
|
|
|
|
|
def test_fused_gated():
|
|
import flash_mla.fused_norm_rope_attn_rope_cast as fused
|
|
|
|
for fn, args in ((fused.permute_q_b_proj, ((None, None), 64, 512)), (fused.permute_wv_proj, ((None, None), 1, 512))):
|
|
try:
|
|
fn(*args)
|
|
FAILS.append("permute should raise")
|
|
except RuntimeError as e:
|
|
print(f"[ok ] {fn.__name__} raises: {e}")
|
|
|
|
# the fused attention core itself (norm + rope + attn + rope), checked against the test ref
|
|
d, h_q, s_q, s_kv, topk = 512, 64, 5, 32, 16
|
|
q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
|
|
positions = torch.randint(0, 1024, (s_q,), dtype=torch.int32)
|
|
cache = torch.randn(1024, 64)
|
|
rms = fused._rms_norm_scale_factor(q, 1e-4)
|
|
q_roped = fused._rope(q, False, cache, positions, 64).to(torch.bfloat16)
|
|
out, ml, lse = flash_mla_sparse_fwd(q_roped, kv, idx, d**-0.5, 512, None, None, _rms_norm_scale_factor=rms)
|
|
out = fused._rope(out, True, cache, positions, 64)
|
|
r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q_roped, kv, idx, d**-0.5, None, None, rms=rms)
|
|
# reference O-RoPE, computed independently
|
|
exp = r_out.float().clone()
|
|
cos = cache[positions.long(), :32].unsqueeze(1)
|
|
sin = -cache[positions.long(), 32:].unsqueeze(1)
|
|
rp = exp[..., -64:]
|
|
x0, x1 = rp[..., 0::2].clone(), rp[..., 1::2].clone()
|
|
rp[..., 0::2] = x0 * cos - x1 * sin
|
|
rp[..., 1::2] = x0 * sin + x1 * cos
|
|
check("fused norm+rope+attn+rope out", out, exp, atol=3e-2, rtol=3e-2)
|
|
check("fused lse", lse, r_lse)
|
|
|
|
|
|
def test_extra_fp8_cache():
|
|
"""SWA main cache + a second fp8 (528 B) cache as `extra`, the non-fp4 DeepInfer variant."""
|
|
d, h_q, b, topk, block = 512, 64, 3, 12, 4
|
|
raw = (torch.randn(5, block, 1, d) / 8).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
|
|
kv = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
|
|
e_raw = (torch.randn(7, block, 1, d) / 8).to(torch.bfloat16)
|
|
e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp8")
|
|
e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp8").view(-1, d)
|
|
q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(b, topk, 5 * block).view(b, 1, topk)
|
|
e_idx = make_indices(b, 9, 7 * block).view(b, 1, 9)
|
|
out, lse = flash_mla_with_kvcache(
|
|
q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
|
|
tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
|
|
extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
|
|
)
|
|
r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv, idx, None), (e_kv, e_idx, None)], d**-0.5)
|
|
check("decode extra fp8 cache out", out, r_out)
|
|
check("decode extra fp8 cache lse", lse, r_lse)
|
|
|
|
|
|
def test_attention_backend_integration():
|
|
"""Drive deepinfer/model_executor/layers/attention_backend.py (head padding path included)."""
|
|
import importlib.util
|
|
import os
|
|
|
|
path = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "model_executor", "layers", "attention_backend.py")
|
|
spec = importlib.util.spec_from_file_location("ab", path)
|
|
ab = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(ab)
|
|
|
|
d, h_q, s_q, s_kv, topk = 512, 40, 6, 30, 16 # 40 heads -> padded to 64
|
|
q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
|
|
kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
|
|
idx = make_indices(s_q, topk, s_kv, invalid_value=-1)
|
|
sink = torch.randn(h_q)
|
|
ilen = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32)
|
|
out, lse = ab.prefill_attn_varlen_sparse(q, kv, idx, d**-0.5, attn_sink=sink, index_length=ilen)
|
|
pq = torch.cat([q, q.new_zeros(s_q, 24, d)], dim=1)
|
|
ps = torch.cat([sink, sink.new_zeros(24)])
|
|
r_out, _, r_lse = ref_attn.ref_sparse_attn_fwd(pq, kv, idx.unsqueeze(1), d**-0.5, ilen, ps)
|
|
check("attention_backend prefill out (head pad 40->64)", out, r_out[:, :h_q])
|
|
check("attention_backend prefill lse", lse, r_lse[:, :h_q])
|
|
|
|
b, block, nblocks = 4, 8, 6
|
|
raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
|
|
cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
|
|
kvb = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
|
|
qd = (torch.randn(b, h_q, d) / 8).to(torch.bfloat16)
|
|
didx = make_indices(b, topk, nblocks * block).view(b, topk)
|
|
dlen = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
|
|
out, lse = ab.decode_mla_mqa_sparse(
|
|
qd, cache, 512, didx, FlashMLASchedMeta(), d**-0.5, attn_sink=sink, index_length=dlen
|
|
)
|
|
pqd = torch.cat([qd, qd.new_zeros(b, 24, d)], dim=1).view(b, 1, 64, d)
|
|
r_out, r_lse = ref_attn.ref_sparse_attn_decode(pqd, [(kvb, didx.view(b, 1, topk), dlen)], d**-0.5, ps)
|
|
check("attention_backend decode out (head pad)", out, r_out.view(b, 64, d)[:, :h_q])
|
|
check("attention_backend decode lse", lse, r_lse.view(b, 64)[:, :h_q])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_kv_dequant()
|
|
for h in (64, 128):
|
|
test_prefill(h_q=h)
|
|
test_prefill(with_sink=False, with_len=False)
|
|
test_prefill_lonely()
|
|
test_decode()
|
|
test_decode(with_sink=False, with_len=False)
|
|
test_decode(extra=True)
|
|
test_decode(strided=True)
|
|
test_decode(h_q=128)
|
|
test_decode_multi_token()
|
|
test_decode_lonely()
|
|
test_sink_infinities()
|
|
test_extra_fp8_cache()
|
|
test_attention_backend_integration()
|
|
test_dense_unsupported()
|
|
test_fused_gated()
|
|
print("\nFAILURES:", FAILS if FAILS else "none")
|
|
sys.exit(1 if FAILS else 0)
|