mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
28fb291db1
No behaviour change. `CONTRIBUTING.md` requires `ruff check --no-cache` and `ruff format --check --no-cache` to pass, and the SM90 branch predates that requirement: 72 findings, all in the files added here. Most of it is mechanical (`ruff format`), plus: - annotations for the helper signatures (`torch.Tensor` / `int` / `Generator` / return types); hooks are typed as `Callable[[nn.Module, tuple, object], None]` rather than left bare; - two `if`/`else` assignments become ternaries (SIM108), and the index read in `h20_real_requant_check.py` uses a context manager (SIM115). Also corrects the `sparse_mla_fwd` docstring, which still described the old `0 <= idx <= q_position` masking contract. Passing `is_causal=True` there is what produced the all-NaN layer 0: indices address the assembled `[swa cache | current swa | compressed]` buffer, so the row number is not a valid causal bound and the engine must pass False.
101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
"""Validate LocalEngramGather against a torch reference, single-process and multi-rank.
|
|
|
|
torchrun --nproc_per_node=8 scripts/h20_engram_check.py
|
|
"""
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from deepinfer.h20_ops.engram_local import LocalEngramGather
|
|
|
|
HIDDEN = 256
|
|
N_LOCAL = 4096 # rows per rank
|
|
ENTRIES = 24 # 3 n-gram tables x 8 heads
|
|
TOKENS = 128
|
|
LAYERS = 2
|
|
|
|
|
|
def quantize_table(rows: torch.Tensor, gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""fp8 table + per-32 ue8m0 (stored as fp32) scales, like the checkpoint."""
|
|
groups = rows.float().view(rows.shape[0], HIDDEN // 32, 32)
|
|
amax = groups.abs().amax(-1).clamp_min(1e-30)
|
|
exp = torch.ceil(torch.log2(amax / 448.0)).clamp_min(-126)
|
|
scale = torch.pow(2.0, exp)
|
|
q = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn).view(rows.shape[0], HIDDEN)
|
|
return q, scale
|
|
|
|
|
|
def main() -> None:
|
|
dist.init_process_group("nccl")
|
|
rank, world = dist.get_rank(), dist.get_world_size()
|
|
torch.cuda.set_device(rank)
|
|
dev = torch.device("cuda", rank)
|
|
gen = torch.Generator(device="cpu").manual_seed(0)
|
|
|
|
torch.manual_seed(rank)
|
|
tables, scales = [], []
|
|
for _ in range(LAYERS):
|
|
rows = torch.randn(N_LOCAL, HIDDEN, generator=gen)
|
|
q, s = quantize_table(rows, gen)
|
|
tables.append(q.to(dev))
|
|
scales.append(s.to(dev))
|
|
|
|
gather = LocalEngramGather(group=None, hidden=HIDDEN)
|
|
# Use a real process group so all_to_all runs across ranks.
|
|
gather.group = dist.group.WORLD
|
|
gather.set_tables(tables, scales)
|
|
assert gather.num_local_entries == [N_LOCAL] * LAYERS
|
|
print(f"[rank {rank}] {gather.describe()}")
|
|
|
|
total = world * N_LOCAL
|
|
idx = torch.randint(0, total, (LAYERS, TOKENS, ENTRIES), generator=gen, dtype=torch.int32).to(dev)
|
|
idx[0, 0, 0] = -1 # invalid slot must come back zeroed
|
|
|
|
hooks = gather.fetch(idx)
|
|
got = hooks[0]()
|
|
|
|
# Reference: rebuild the global table from every rank.
|
|
ref_rows = [torch.empty(world * N_LOCAL, HIDDEN, dtype=torch.float8_e4m3fn, device=dev) for _ in range(LAYERS)]
|
|
ref_sf = [torch.empty(world * N_LOCAL, HIDDEN // 32, device=dev) for _ in range(LAYERS)]
|
|
for layer in range(LAYERS):
|
|
dist.all_gather_into_tensor(ref_rows[layer], tables[layer])
|
|
dist.all_gather_into_tensor(ref_sf[layer], scales[layer])
|
|
|
|
def reference(layer: int) -> torch.Tensor:
|
|
flat = idx[layer].reshape(-1).long()
|
|
valid = flat >= 0
|
|
safe = flat.clamp_min(0)
|
|
w = ref_rows[layer].index_select(0, safe).float()
|
|
s = ref_sf[layer].index_select(0, safe)
|
|
vals = w.reshape(-1, HIDDEN // 32, 32) * s.unsqueeze(-1)
|
|
vals = vals.reshape(-1, HIDDEN).to(torch.bfloat16)
|
|
vals[~valid] = 0
|
|
return vals
|
|
|
|
a, sf = got if isinstance(got, tuple) else (got, None)
|
|
ref = reference(0)
|
|
# `a` holds fp8 codes; the per-token 1x128 scales turn them back into values.
|
|
got_vals = (a.float() * sf.float().repeat_interleave(HIDDEN * ENTRIES // sf.shape[1], dim=1)).view(TOKENS, ENTRIES, HIDDEN)
|
|
# Compare in dequantized space: per-128 requantization is the only difference expected.
|
|
ref_deq = ref.float().view(TOKENS, ENTRIES, HIDDEN)
|
|
err = (got_vals - ref_deq).abs().max().item()
|
|
rel = err / max(ref_deq.abs().max().item(), 1e-9)
|
|
print(f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} max_abs={err:.4f} rel={rel:.4f}")
|
|
print(f"[rank {rank}] invalid slot all-zero: {bool((got_vals[0, 0] == 0).all())}")
|
|
|
|
# Non-owner traffic actually exercised?
|
|
owners = torch.div(idx[0].reshape(-1).clamp_min(0), N_LOCAL, rounding_mode="floor")
|
|
remote = int((owners != rank).sum())
|
|
print(f"[rank {rank}] remote entries {remote}/{owners.numel()}")
|
|
|
|
ok = rel < 0.05
|
|
print(f"[rank {rank}] {'PASS' if ok else 'FAIL'}")
|
|
dist.barrier()
|
|
dist.destroy_process_group()
|
|
if not ok:
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|