mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
02e099fff8
- NeuronGraphRunner: the CudaGraphRunner contract over torch.compile(dynamic=False), one graph per batch size (DEEPINFER_ENFORCE_EAGER=0, backend DEEPINFER_NEURON_COMPILE_BACKEND). Layer cache views are re-sliced from their root allocations inside the graph, so AOT never regenerates aliased mutated inputs with as_strided. - Decode MoE with fixed shapes: packed all-gather, dense local experts, fp32 reduce-scatter (neuron_ops/static_moe.py, deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a). - Graph-break and compile fixes: hide torch.cuda during compile, torch.device equality for FX codegen, hash-based force_random routing, Engram hash/fetch on the host (exact int64), env vars resolved before tracing, Movable annotation, no .item()/nonzero/as_strided writes on the decode path, KV stores route padding slots to a scratch block kept past the pool, native-dtype gather under compile. - fp8: trn2 reads float8_e4m3fn as IEEE e4m3 (max 240). Activations/KV quantize to 224 (DEEPINFER_NEURON_FP8_MAX), checkpoint fp8 weights are rescaled at load (bytes / 2, ue8m0 + 1), and neuronx-cc gets --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 via --internal-hlo2tensorizer-options. The uint8-carrier backend (neuron_ops/neuron_compile.py, "deepinfer_neuron") stays as a fallback. - Tests: tests/neuron/compile; run_all.sh runs kernel unit tests at the e4m3fn range (448) and the end-to-end stages at 224. Docs: vllm-neuron PR #71 analysis and compile notes. Host: tiny model 9/9 vs golden (all prompts) eager, compiled path (aot_eager) matches eager; CPU suite green. Device: the whole decode graph still crashes neuronx-cc (NCC_IINL902, being bisected). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
"""trn2 e4m3 range handling (deepinfer/neuron_ops/fp8.py, issue #24), on the host."""
|
|
|
|
import torch
|
|
|
|
from deepinfer.neuron_ops import fp8
|
|
|
|
FP8 = torch.float8_e4m3fn
|
|
|
|
|
|
def test_halve_matches_torch_cast():
|
|
b = torch.arange(256, dtype=torch.int32).to(torch.uint8)
|
|
ref = (b.view(FP8).float() / 2).to(FP8).view(torch.uint8)
|
|
got = fp8.halve_e4m3fn_bytes(b)
|
|
nan = (b & 0x7F) == 0x7F
|
|
assert torch.equal(got[~nan], ref[~nan])
|
|
assert bool(((got[nan] & 0x7F) == 0x7F).all())
|
|
|
|
|
|
def test_halved_values_fit_hardware_e4m3():
|
|
b = torch.arange(256, dtype=torch.int32).to(torch.uint8)
|
|
got = fp8.halve_e4m3fn_bytes(b)
|
|
finite = (b & 0x7F) != 0x7F
|
|
assert bool((((got[finite] >> 3) & 15) < 15).all()) # exponent field 15 is inf/NaN on trn2
|
|
|
|
|
|
def test_rescale_keeps_dequantized_values():
|
|
torch.manual_seed(0)
|
|
w = torch.nn.Parameter((torch.randn(64, 128) * 100).clamp(-448, 448).to(FP8), requires_grad=False)
|
|
w.sf = torch.randint(110, 130, (64, 4), dtype=torch.uint8) # ue8m0 per 32 columns
|
|
before = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1)
|
|
lossy = ((w.data.view(torch.uint8) >> 3) & 15) <= 1 # halving rounds only below the smallest normal
|
|
model = torch.nn.Module()
|
|
model.w = w
|
|
if fp8.FP8_MAX > 240:
|
|
return
|
|
assert fp8.rescale_fp8_weights(model) == 1
|
|
after = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1)
|
|
assert torch.equal(after[~lossy], before[~lossy])
|
|
ulp = 2.0**-9 * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1) # one smallest-subnormal step, after rescale
|
|
assert bool(((after - before).abs()[lossy] <= ulp[lossy]).all())
|
|
assert float(w.float().abs().max()) <= 224.0
|