mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +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>
90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
"""Reference quant/dequant ported from vendor/FlashMLA/tests/quant.py (CPU, no kernelkit)."""
|
|
|
|
import torch
|
|
|
|
META = {
|
|
"v41_fp8": (512, 448, 64, 32, 16),
|
|
"v41_fp4": (512, 448, 64, 16, 32),
|
|
}
|
|
BPT = {"v41_fp8": 528, "v41_fp4": 288}
|
|
|
|
_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32)
|
|
|
|
|
|
def _cast_scale_inv_to_ue8m0(scales_inv, out_dtype=torch.float32):
|
|
return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype)
|
|
|
|
|
|
def _quantize_to_e2m1(x):
|
|
x = x.float()
|
|
mags = _E2M1
|
|
sign = (torch.signbit(x)).to(torch.uint8) << 3
|
|
a = torch.nan_to_num(x.abs(), nan=0.0, posinf=6.0).clamp_max(6.0)
|
|
mids = (mags[:-1] + mags[1:]) / 2
|
|
code = torch.bucketize(a, mids, right=True)
|
|
on_tie = (a.unsqueeze(-1) == mids).any(dim=-1)
|
|
tie_code = torch.bucketize(a, mids, right=False)
|
|
code = torch.where(on_tie, tie_code + (tie_code & 1), code)
|
|
return (sign | code.to(torch.uint8)).to(torch.uint8)
|
|
|
|
|
|
def quantize_k_cache(input_k_cache, fmt):
|
|
"""input_k_cache: (num_blocks, block_size, 1, 512) bf16 -> uint8 (nb, bs, 1, bpt)."""
|
|
d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
|
|
num_blocks, block_size, h_k, _ = input_k_cache.shape
|
|
assert h_k == 1
|
|
x = input_k_cache.squeeze(2)
|
|
bpt = BPT[fmt]
|
|
result = torch.empty((num_blocks, block_size, bpt), dtype=torch.uint8)
|
|
|
|
if fmt == "v41_fp8":
|
|
data = result[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
|
|
sf = result[..., d_nope + d_rope :]
|
|
for t in range(16):
|
|
sl = x[..., t * tile_size : (t + 1) * tile_size]
|
|
inv = torch.abs(sl).max(dim=-1).values.float() / 448.0
|
|
inv = _cast_scale_inv_to_ue8m0(inv)
|
|
# ue8m0 byte = bf16 bits >> 7 (Ascend path of tests/quant.py)
|
|
sf[:, :, t] = (inv.view(torch.int32) // 2**23).to(torch.uint8)
|
|
data[:, :, t * tile_size : (t + 1) * tile_size] = (sl.float() / inv.unsqueeze(-1)).to(
|
|
torch.float8_e4m3fn
|
|
)
|
|
else:
|
|
data = result[..., : d // 2]
|
|
sf = result[..., d // 2 :].view(torch.float8_e4m3fn)
|
|
xf = x.float()
|
|
amax = torch.nan_to_num(xf.abs(), nan=float("inf")).view(num_blocks, block_size, num_tiles, tile_size).amax(-1)
|
|
scale = torch.clamp(amax / 6.0, 2.0**-9, 448.0).to(torch.float8_e4m3fn)
|
|
normalized = xf.view(num_blocks, block_size, num_tiles, tile_size) / scale.float().unsqueeze(-1)
|
|
codes = _quantize_to_e2m1(normalized).view(num_blocks, block_size, d)
|
|
data[...] = codes[..., 0::2] | (codes[..., 1::2] << 4)
|
|
sf[...] = scale
|
|
return result.view(num_blocks, block_size, 1, bpt)
|
|
|
|
|
|
def dequantize_k_cache(quant_k_cache, fmt):
|
|
d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
|
|
num_blocks, block_size, h_k, _ = quant_k_cache.shape
|
|
assert h_k == 1
|
|
out = torch.empty((num_blocks, block_size, d), dtype=torch.bfloat16)
|
|
q = quant_k_cache.view(num_blocks, block_size, -1)
|
|
if fmt == "v41_fp8":
|
|
raw = q[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
|
|
sf = q[..., d_nope + d_rope :]
|
|
for t in range(16):
|
|
cur = raw[..., t * tile_size : (t + 1) * tile_size].to(torch.bfloat16)
|
|
scales = (sf[:, :, t].to(torch.int32) * 2**7).to(torch.int16).view(torch.bfloat16)
|
|
out[..., t * tile_size : (t + 1) * tile_size] = cur * scales.unsqueeze(-1)
|
|
else:
|
|
data = q[..., : d // 2]
|
|
sf = q[..., d // 2 :].view(torch.float8_e4m3fn)
|
|
codes = torch.empty((num_blocks, block_size, d), dtype=torch.uint8)
|
|
codes[..., 0::2] = data & 0xF
|
|
codes[..., 1::2] = data >> 4
|
|
val = _E2M1[(codes & 7).long()]
|
|
val = torch.where((codes & 8) != 0, -val, val).view(num_blocks, block_size, num_tiles, tile_size)
|
|
out[...] = (val.to(torch.bfloat16) * sf.to(torch.bfloat16).unsqueeze(-1)).view(
|
|
num_blocks, block_size, d
|
|
)
|
|
return out.view(num_blocks, block_size, 1, d)
|