mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
Neuron: test suite for new boxes, device-vs-host differential tool
- 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>
This commit is contained in:
whn09
parent
d2389f8a9d
commit
760e3cadc0
216
examples/neuron/device_diff.py
Normal file
216
examples/neuron/device_diff.py
Normal file
@ -0,0 +1,216 @@
|
||||
"""Device-vs-host differential check for the Neuron kernels.
|
||||
|
||||
Runs the tiny checkpoint on the NeuronCore through run_deepinfer_tiny.py and, for every call of a vendor
|
||||
kernel (deep_gemm, tile_kernels, flash_mla, deep_select) or a tilelang_ops/neuron kernel, re-runs the same
|
||||
pure-torch implementation on host copies of the arguments and compares the results, including arguments
|
||||
mutated in place. Prints one line per kernel with the worst relative error and the argument layouts of
|
||||
the worst call, so a kernel that misbehaves on the device (and only there) stands out.
|
||||
|
||||
NEURON_RT_VISIBLE_CORES=63 CUDA_VISIBLE_DEVICES=63 python examples/neuron/device_diff.py \\
|
||||
--prompts short --out ~/di/neuron_diff [run_deepinfer_tiny.py args...]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import functools
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
_STATS: dict[str, dict[str, Any]] = collections.OrderedDict()
|
||||
_ACTIVE = [True]
|
||||
|
||||
|
||||
def _to_host(x: Any) -> Any:
|
||||
if isinstance(x, torch.Tensor):
|
||||
# Keep strides: a few kernels depend on the caller's (e.g. column-major) layout.
|
||||
return torch.empty_strided(x.shape, x.stride(), dtype=x.dtype, device="cpu").copy_(x.cpu()) if x.device.type != "cpu" else x.clone()
|
||||
if isinstance(x, (list, tuple)):
|
||||
return type(x)(_to_host(v) for v in x)
|
||||
if isinstance(x, dict):
|
||||
return {k: _to_host(v) for k, v in x.items()}
|
||||
return x
|
||||
|
||||
|
||||
def _tensors(x: Any, prefix: str = "") -> list[tuple[str, torch.Tensor]]:
|
||||
if isinstance(x, torch.Tensor):
|
||||
return [(prefix, x)]
|
||||
if isinstance(x, (list, tuple)):
|
||||
return [t for i, v in enumerate(x) for t in _tensors(v, f"{prefix}[{i}]")]
|
||||
if isinstance(x, dict):
|
||||
return [t for k, v in x.items() for t in _tensors(v, f"{prefix}.{k}")]
|
||||
return []
|
||||
|
||||
|
||||
def _rel(dev: torch.Tensor, host: torch.Tensor) -> float:
|
||||
a = dev.detach().cpu()
|
||||
b = host.detach()
|
||||
if a.shape != b.shape:
|
||||
return float("nan")
|
||||
if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu) or not a.is_floating_point():
|
||||
if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
|
||||
a, b = a.view(torch.uint8), b.view(torch.uint8)
|
||||
neq = (a != b).float().mean().item() if a.numel() else 0.0
|
||||
return neq # fraction of differing elements / bytes
|
||||
a, b = a.float(), b.float()
|
||||
finite = torch.isfinite(b) & torch.isfinite(a)
|
||||
if not finite.all():
|
||||
if not torch.equal(torch.isfinite(a), torch.isfinite(b)):
|
||||
return float("inf")
|
||||
a, b = a[finite], b[finite]
|
||||
den = b.norm().item()
|
||||
return (a - b).norm().item() / den if den > 0 else (a - b).norm().item()
|
||||
|
||||
|
||||
def _layout(args: Any, kwargs: Any) -> str:
|
||||
parts = []
|
||||
for name, t in _tensors(args, "a") + _tensors(kwargs, "k"):
|
||||
parts.append(f"{name}:{str(t.dtype).replace('torch.', '')}{list(t.shape)}s{list(t.stride())}")
|
||||
return " ".join(parts)[:400]
|
||||
|
||||
|
||||
def _record(name: str, rel: float, where: str, layout: str) -> None:
|
||||
st = _STATS.setdefault(name, {"calls": 0, "worst": 0.0, "where": "", "layout": ""})
|
||||
st["calls"] += 1
|
||||
if not (rel <= st["worst"]): # also catches nan/inf
|
||||
st["worst"], st["where"], st["layout"] = rel, where, layout
|
||||
|
||||
|
||||
def wrap(name: str, fn: Any) -> Any:
|
||||
if getattr(fn, "_device_diff", False):
|
||||
return fn
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
if not _ACTIVE[0]:
|
||||
return fn(*args, **kwargs)
|
||||
h_args, h_kwargs = _to_host(args), _to_host(kwargs)
|
||||
layout = _layout(args, kwargs)
|
||||
out = fn(*args, **kwargs)
|
||||
_ACTIVE[0] = False # nested kernels inside the host replay are not checked again
|
||||
try:
|
||||
h_out = fn(*h_args, **h_kwargs)
|
||||
except Exception as e: # noqa: BLE001
|
||||
_record(name, float("nan"), f"host replay failed: {type(e).__name__}: {e}"[:200], layout)
|
||||
return out
|
||||
finally:
|
||||
_ACTIVE[0] = True
|
||||
worst, where = 0.0, ""
|
||||
for (n, d), (_, h) in zip(_tensors(out, "out"), _tensors(h_out, "out")):
|
||||
r = _rel(d, h)
|
||||
if not (r <= worst):
|
||||
worst, where = r, n
|
||||
for (n, d), (_, h) in zip(_tensors(args, "a") + _tensors(kwargs, "k"), _tensors(h_args, "a") + _tensors(h_kwargs, "k")):
|
||||
if d.device.type == "cpu":
|
||||
continue
|
||||
r = _rel(d, h) # in-place written arguments
|
||||
if not (r <= worst):
|
||||
worst, where = r, f"inplace {n}"
|
||||
_record(name, worst, where, layout)
|
||||
return out
|
||||
|
||||
wrapper._device_diff = True
|
||||
return wrapper
|
||||
|
||||
|
||||
def wrap_factory(name: str, factory: Any) -> Any:
|
||||
"""tilelang_ops: ``xxx_tl(static...)`` returns the kernel; wrap the kernel it returns."""
|
||||
|
||||
@functools.wraps(factory)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
return wrap(name, factory(*args, **kwargs))
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def install() -> None:
|
||||
import deep_gemm
|
||||
import deep_select.interface
|
||||
import flash_mla
|
||||
import flash_mla.flash_mla_interface
|
||||
import tile_kernels
|
||||
import tile_kernels.engram
|
||||
import tile_kernels.moe
|
||||
import tile_kernels.quant
|
||||
import tile_kernels.transform
|
||||
|
||||
replaced: dict[int, Any] = {}
|
||||
|
||||
def patch_module(mod: types.ModuleType, prefix: str) -> None:
|
||||
for attr, val in list(vars(mod).items()):
|
||||
if attr.startswith("_") or not callable(val) or isinstance(val, type):
|
||||
continue
|
||||
if getattr(val, "__module__", "") is None or not str(getattr(val, "__module__", "")).startswith(prefix.split(".")[0]):
|
||||
continue
|
||||
w = replaced.get(id(val)) or wrap(f"{prefix}.{attr}", val)
|
||||
replaced[id(val)] = w
|
||||
setattr(mod, attr, w)
|
||||
|
||||
for mod, prefix in [
|
||||
(deep_gemm, "deep_gemm"),
|
||||
(flash_mla, "flash_mla"),
|
||||
(flash_mla.flash_mla_interface, "flash_mla"),
|
||||
(deep_select.interface, "deep_select"),
|
||||
(tile_kernels.quant, "tile_kernels.quant"),
|
||||
(tile_kernels.transform, "tile_kernels.transform"),
|
||||
(tile_kernels.moe, "tile_kernels.moe"),
|
||||
(tile_kernels.engram, "tile_kernels.engram"),
|
||||
]:
|
||||
patch_module(mod, prefix)
|
||||
for sub in ("expand_kernel", "head_compute_mix_kernel", "norm_fn_kernel", "post_kernel", "pre_apply_mix_kernel",
|
||||
"pre_split_mixes_kernel", "sinkhorn_kernel", "pre_big_fuse_kernel"):
|
||||
m = __import__(f"tile_kernels.mhc.{sub}", fromlist=["x"])
|
||||
patch_module(m, "tile_kernels.mhc")
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops import ops
|
||||
|
||||
for attr in [a for a in dir(ops) if a.endswith("_tl")]:
|
||||
setattr(ops, attr, wrap_factory(f"tl.{attr}", getattr(ops, attr)))
|
||||
|
||||
# Model modules bound some kernels with `from X import f`: rebind those names too.
|
||||
import importlib
|
||||
|
||||
for modname in ["deepinfer.model_executor.models.v4_modules.mhc_ops", "deepinfer.model_executor.models.v4_modules.mhc_modeling",
|
||||
"deepinfer.model_executor.models.v4_modules.mlp", "deepinfer.model_executor.models.v4_modules.hc",
|
||||
"deepinfer.model_executor.models.v4_modules.attn", "deepinfer.model_executor.models.v4_modules.index_attn",
|
||||
"deepinfer.model_executor.layers.attention_backend", "deepinfer.model_executor.layers.rotary_embedding",
|
||||
"deepinfer.model_executor.layers.sampler", "deepinfer.model_executor.layers.sampler_utils"]:
|
||||
mod = importlib.import_module(modname)
|
||||
for attr, val in list(vars(mod).items()):
|
||||
if id(val) in replaced:
|
||||
setattr(mod, attr, replaced[id(val)])
|
||||
|
||||
|
||||
def report() -> None:
|
||||
print("\n[device_diff] kernel calls worst-rel where / layout of worst call")
|
||||
for name, st in sorted(_STATS.items(), key=lambda kv: -(kv[1]["worst"] if kv[1]["worst"] == kv[1]["worst"] else 1e9)):
|
||||
print(f"[device_diff] {name:48s} {st['calls']:6d} {st['worst']:9.2e} {st['where']} {st['layout']}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
import run_deepinfer_tiny as R
|
||||
|
||||
argv = sys.argv[1:]
|
||||
if "--device" not in argv:
|
||||
argv = ["--device", "neuron", *argv]
|
||||
R.setup_env("neuron")
|
||||
import deepinfer
|
||||
|
||||
deepinfer.PLATFORM # noqa: B018 - initialize the platform (device mapping) before patching
|
||||
install()
|
||||
sys.argv = [sys.argv[0], *argv]
|
||||
try:
|
||||
R.main()
|
||||
finally:
|
||||
report()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -291,6 +291,10 @@ def run_deepinfer(args):
|
||||
print("[probe] using P-bf16 sparse attention (proposed fix)")
|
||||
wrap(mlp_mod, "moe_topk_gate", lambda o, a, k: (cap.put("topk_idx", o[0].float()), cap.put("topk_weights", o[1])))
|
||||
|
||||
# Fixed KV pool: the engine's profile run (a max_num_batched_tokens dummy prefill) is slow on the host and
|
||||
# exhausts HBM on the NeuronCore with the pure-torch kernels.
|
||||
R.install_workarounds(32)
|
||||
|
||||
from deepinfer.entrypoints.llm import LLM
|
||||
from deepinfer.sampling_params import SamplingParams
|
||||
|
||||
|
||||
250
tests/neuron/deep_ep/refs.py
Normal file
250
tests/neuron/deep_ep/refs.py
Normal file
@ -0,0 +1,250 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
def get_global_seed() -> int:
|
||||
return 0
|
||||
|
||||
|
||||
def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
|
||||
topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor],
|
||||
num_max_tokens_per_rank: int, num_experts: int):
|
||||
"""
|
||||
The reference implementation of dispatching tokens to experts across multiple ranks.
|
||||
|
||||
Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`).
|
||||
|
||||
Arguments:
|
||||
x: input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden]`, `[num_tokens, num_sf_packs]`).
|
||||
topk_idx: top-k expert indices for each token, `[num_tokens, num_topk]`.
|
||||
topk_weights: top-k weights for each token, can be None, `[num_tokens, num_topk]`.
|
||||
num_max_tokens_per_rank: the maximum number of tokens per rank, must be `>=` the actual number of
|
||||
tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank`, since we calculate
|
||||
`src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
|
||||
num_experts: the total number of experts across all ranks.
|
||||
|
||||
Returns:
|
||||
recv_x, recv_topk_idx, recv_topk_weights: received tokens, top-k indices and top-k weights after
|
||||
dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are
|
||||
set to -1.
|
||||
recv_src_token_idx: received `src_token_global_idx` for each received token.
|
||||
num_recv_tokens_per_rank: number of received tokens from each rank, `[num_ranks]`.
|
||||
"""
|
||||
rank_idx = dist.get_rank()
|
||||
num_ranks = dist.get_world_size()
|
||||
|
||||
assert num_experts % num_ranks == 0
|
||||
num_experts_per_rank = num_experts // num_ranks
|
||||
|
||||
# Unpack SF
|
||||
use_fp8 = isinstance(x, tuple)
|
||||
x, sf = x if use_fp8 else (x, None)
|
||||
|
||||
num_tokens, hidden = x.size()
|
||||
num_tokens_, num_topk = topk_idx.size()
|
||||
assert num_tokens == num_tokens_
|
||||
num_sf_packs = 0
|
||||
if sf is not None:
|
||||
num_tokens__, num_sf_packs = sf.size()
|
||||
assert num_tokens == num_tokens__
|
||||
if topk_weights is not None:
|
||||
num_tokens__, num_topk_ = topk_weights.size()
|
||||
assert num_tokens == num_tokens__
|
||||
assert num_topk == num_topk_
|
||||
|
||||
# Prepare per-peer send buffers
|
||||
send_x_list = []
|
||||
send_sf_list = []
|
||||
send_topk_idx_list = []
|
||||
send_topk_weights_list = []
|
||||
send_src_token_idx_list = []
|
||||
num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device)
|
||||
for dst_rank_idx in range(num_ranks):
|
||||
expert_start_idx = dst_rank_idx * num_experts_per_rank
|
||||
expert_end_idx = expert_start_idx + num_experts_per_rank
|
||||
|
||||
# Get the indices of tokens
|
||||
mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
|
||||
indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
|
||||
num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel()
|
||||
|
||||
# Select the data for tokens
|
||||
x_to_send = x[indices_to_send]
|
||||
sf_to_send = sf[indices_to_send] if use_fp8 else None
|
||||
topk_idx_to_send = topk_idx[indices_to_send]
|
||||
topk_weights_to_send = topk_weights[indices_to_send]
|
||||
masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx),
|
||||
topk_idx_to_send, torch.full_like(topk_idx_to_send, -1))
|
||||
|
||||
send_x_list.append(x_to_send)
|
||||
send_sf_list.append(sf_to_send)
|
||||
send_topk_idx_list.append(masked_topk_idx)
|
||||
send_topk_weights_list.append(topk_weights_to_send)
|
||||
send_src_token_idx_list.append(indices_to_send)
|
||||
|
||||
send_x = torch.cat(send_x_list, dim=0)
|
||||
send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None
|
||||
send_topk_idx = torch.cat(send_topk_idx_list, dim=0)
|
||||
send_topk_weights = torch.cat(send_topk_weights_list, dim=0)
|
||||
send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int)
|
||||
send_src_token_idx += rank_idx * num_max_tokens_per_rank
|
||||
|
||||
# Exchange size
|
||||
num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device)
|
||||
dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
num_recv_tokens = int(num_recv_tokens_per_rank.sum().item())
|
||||
|
||||
# Exchange main data
|
||||
num_send_tokens_per_rank = num_send_tokens_per_rank.tolist()
|
||||
num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist()
|
||||
recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device)
|
||||
recv_sf = torch.empty((num_recv_tokens, num_sf_packs), dtype=sf.dtype, device=x.device) if use_fp8 else None
|
||||
recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device)
|
||||
recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device)
|
||||
recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device)
|
||||
dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
if use_fp8:
|
||||
# gloo has no int16 collective: move the SF bytes (bit-exact, same semantics).
|
||||
dist.all_to_all_single(recv_sf.view(torch.uint8), send_sf.contiguous().view(torch.uint8),
|
||||
num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
||||
|
||||
# Mask top-k indices
|
||||
expert_start_idx = rank_idx * num_experts_per_rank
|
||||
expert_end_idx = expert_start_idx + num_experts_per_rank
|
||||
mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx)
|
||||
recv_topk_idx = recv_topk_idx - expert_start_idx
|
||||
recv_topk_idx.masked_fill_(~mask, -1)
|
||||
|
||||
# Pack SF
|
||||
recv_x = (recv_x, recv_sf) if use_fp8 else recv_x
|
||||
|
||||
return (recv_x, recv_topk_idx, recv_topk_weights,
|
||||
recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int))
|
||||
|
||||
|
||||
def generate_pre_combine_data(src_token_global_idx: torch.Tensor,
|
||||
num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor:
|
||||
"""
|
||||
Generate data needed for combine from `src_token_global_idx`.
|
||||
Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
|
||||
The generated data (denoted as `y`) of the i-th token has a shape of `[num_topk, hidden]`, with
|
||||
|
||||
`y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))`
|
||||
|
||||
where `P=131071` is a large prime, `token_seeds` is calculated via
|
||||
`token_seeds = src_token_global_idx[i] * num_topk + j`, and
|
||||
`max_seed = num_ranks * num_max_tokens_per_rank * num_topk`.
|
||||
|
||||
Arguments:
|
||||
src_token_global_idx: source token global indices, `[num_tokens]`.
|
||||
|
||||
Returns:
|
||||
the generated data, `[num_tokens, num_topk, hidden]`.
|
||||
"""
|
||||
num_ranks = dist.get_world_size()
|
||||
token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
|
||||
torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0)) # [num_tokens, num_topk]
|
||||
max_seed = num_ranks * num_max_tokens_per_rank * num_topk
|
||||
result = torch.sin(
|
||||
(((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
|
||||
torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) +
|
||||
math.sin(float(get_global_seed()))
|
||||
)
|
||||
return result.to(torch.bfloat16)
|
||||
|
||||
|
||||
def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
"""
|
||||
Accumulate `data` in order along the num_topk dimension.
|
||||
|
||||
Arguments:
|
||||
data: data to be accumulated, `[num_tokens, num_topk, hidden]`.
|
||||
initial_value: initial value for accumulation, `[num_tokens, hidden]`.
|
||||
|
||||
Returns:
|
||||
the result, `[num_tokens, hidden]`.
|
||||
"""
|
||||
num_topk = data.shape[1]
|
||||
if initial_value is None:
|
||||
result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device)
|
||||
else:
|
||||
result = initial_value.clone()
|
||||
for i in range(num_topk):
|
||||
result += data[:, i, :].float()
|
||||
return result.to(data.dtype)
|
||||
|
||||
|
||||
def combine(y: torch.Tensor, topk_idx: torch.Tensor, num_experts: int,
|
||||
bias: Optional[Union[Tuple[torch.Tensor, ...], torch.Tensor]],
|
||||
reduce_in_local: bool) -> torch.Tensor:
|
||||
"""
|
||||
The reference implementation of (possibly intra-rank reduced) combining tokens.
|
||||
|
||||
Arguments:
|
||||
y: input tokens to be combined, `[num_tokens, num_topk, hidden]`.
|
||||
topk_idx: `[num_tokens, num_topk]`.
|
||||
num_experts: the total number of experts across all ranks.
|
||||
bias: 0, 1 or 2 `[num_tokens, hidden]` final bias added to the output.
|
||||
reduce_in_local: whether to do reduction within each rank first.
|
||||
- `True` -> reduce selections landing on the same rank before the global accumulation.
|
||||
- `False` -> equivalent to `allow_multiple_reduction` being False.
|
||||
|
||||
Returns:
|
||||
the combined result, `[num_tokens, hidden]`.
|
||||
"""
|
||||
num_ranks = dist.get_world_size()
|
||||
num_tokens, hidden = y.shape[0], y.shape[2]
|
||||
num_topk = y.shape[1]
|
||||
|
||||
def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Perform in-place grouped reduction on `data_to_reduce` according to `group_id`.
|
||||
The summation within each group is performed in strict order along the `num_topk` dimension.
|
||||
The result for each group is stored at the rightmost token of that group, and other tokens are
|
||||
set to zero.
|
||||
|
||||
Arguments:
|
||||
data_to_reduce: data to be reduced, `[num_tokens, num_topk, hidden]`.
|
||||
group_id: group IDs for each token, `[num_tokens, num_topk]`.
|
||||
"""
|
||||
# Shuffle to make tokens with the same group_id contiguous
|
||||
group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
|
||||
# transformed_src_indices[i, j] = i * num_topk + src_indices[i, j]
|
||||
transformed_src_indices = (
|
||||
(src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten())
|
||||
data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden)
|
||||
# Perform segmented reduce within each group
|
||||
cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device)
|
||||
for i in range(num_topk):
|
||||
is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \
|
||||
if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1]
|
||||
cur_accum_buf += data_to_reduce[:, i, :].float()
|
||||
# For one token, if `is_segment_break` is True,
|
||||
# save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]`
|
||||
segment_break_token_indices = torch.where(is_segment_break)[0]
|
||||
data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype)
|
||||
cur_accum_buf[segment_break_token_indices] = 0.0
|
||||
non_segment_break_token_indices = torch.where(~is_segment_break)[0]
|
||||
data_to_reduce[non_segment_break_token_indices, i] = 0.0
|
||||
# Unshuffle
|
||||
# noinspection PyShadowingNames
|
||||
result = torch.empty_like(data_to_reduce)
|
||||
result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden)
|
||||
return result.view(num_tokens, num_topk, hidden)
|
||||
|
||||
num_experts_per_rank = num_experts // num_ranks
|
||||
if reduce_in_local:
|
||||
y = grouped_reduce(y, topk_idx // num_experts_per_rank)
|
||||
bias_sum = None
|
||||
if isinstance(bias, tuple):
|
||||
assert len(bias) <= 2
|
||||
for bias_tensor in bias:
|
||||
bias_sum = bias_tensor.float() if bias_sum is None else bias_sum + bias_tensor.float()
|
||||
elif bias is not None:
|
||||
bias_sum = bias.float()
|
||||
result = ordered_accumulate(y, bias_sum)
|
||||
return result
|
||||
103
tests/neuron/deep_ep/test_bucket.py
Normal file
103
tests/neuron/deep_ep/test_bucket.py
Normal file
@ -0,0 +1,103 @@
|
||||
"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*, CPU + gloo).
|
||||
|
||||
Covers what `model_executor/parallel_utils/cp.py` uses: a session, `session.allocate`,
|
||||
in-place and out-of-place `all_gather`, and the NotImplementedError stubs.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
||||
import deep_ep # noqa: E402
|
||||
|
||||
# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
|
||||
# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
|
||||
DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
|
||||
|
||||
|
||||
def worker(local_rank: int, num_ranks: int, args):
|
||||
os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
|
||||
os.environ.setdefault('MASTER_PORT', str(args.port))
|
||||
if DEVICE == 'neuron':
|
||||
# Must be set before the Neuron runtime initializes (first device op) in this process.
|
||||
os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
|
||||
dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
|
||||
group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
|
||||
rank = group.rank()
|
||||
|
||||
alignment = deep_ep.get_num_allocation_alignment()
|
||||
assert alignment == 2 << 20 and deep_ep.get_num_rdma_alignment() == 512
|
||||
|
||||
# Plan-based construction (upstream BufferAllocator path).
|
||||
plan = deep_ep.BufferAllocator()
|
||||
gathered = plan.allocate((num_ranks, 2, 8), torch.float32)
|
||||
buffer = deep_ep.BucketBuffer(group, plan, explicitly_destroy=True, device=torch.device(DEVICE))
|
||||
assert buffer.num_bytes == plan.num_bytes
|
||||
|
||||
src = torch.arange(16, dtype=torch.float32, device=DEVICE).view(2, 8) + rank * 100
|
||||
ref = torch.stack([src - rank * 100 + r * 100 for r in range(num_ranks)])
|
||||
|
||||
# Out-of-place: registered destination.
|
||||
gathered.fill_(float('nan'))
|
||||
result = buffer.all_gather(src, dsts=gathered, group=group).wait()
|
||||
assert result.ndim == 1 and torch.equal(result, ref.reshape(-1))
|
||||
|
||||
# In-place: source is this rank's shard of a registered region.
|
||||
gathered.fill_(float('nan'))
|
||||
gathered[rank].copy_(src)
|
||||
result = buffer.all_gather(gathered[rank], group=group).wait()
|
||||
assert torch.equal(result, ref.reshape(-1))
|
||||
buffer.barrier(group=group)
|
||||
buffer.destroy()
|
||||
|
||||
# Session path (what CPManager uses).
|
||||
buffer = deep_ep.BucketBuffer(group, 4 * alignment, explicitly_destroy=True,
|
||||
device=torch.device(DEVICE))
|
||||
with buffer.session() as session:
|
||||
region = session.allocate((num_ranks, 2, 8), torch.float32, group=group)
|
||||
shard = region[rank]
|
||||
shard.copy_(src)
|
||||
handle = buffer.all_gather(shard)
|
||||
out = handle.wait()
|
||||
assert torch.equal(out.view(num_ranks, 2, 8).flatten(0, 1), ref.reshape(-1, 8))
|
||||
|
||||
# Unregistered source is staged into the session storage.
|
||||
unregistered = (src + 1000).contiguous()
|
||||
out = buffer.all_gather(unregistered).wait()
|
||||
ref2 = torch.stack([src - rank * 100 + r * 100 + 1000 for r in range(num_ranks)])
|
||||
assert torch.equal(out, ref2.reshape(-1))
|
||||
|
||||
# Batched
|
||||
outs = buffer.all_gather([shard, shard]).wait()
|
||||
assert len(outs) == 2 and all(torch.equal(o, ref.reshape(-1)) for o in outs)
|
||||
|
||||
for name in ('reduce_scatter', 'all_reduce'):
|
||||
try:
|
||||
getattr(buffer, name)(shard)
|
||||
except NotImplementedError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError(f'{name} should raise NotImplementedError')
|
||||
assert 'all_gather' not in buffer.__dict__
|
||||
buffer.destroy()
|
||||
|
||||
if rank == 0:
|
||||
print(f'Bucket tests passed (world_size={num_ranks})', flush=True)
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--num-processes', type=int, default=2)
|
||||
parser.add_argument('--port', type=int, default=29551)
|
||||
args = parser.parse_args()
|
||||
if args.num_processes == 1:
|
||||
worker(0, 1, args)
|
||||
else:
|
||||
mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
|
||||
131
tests/neuron/deep_ep/test_engram.py
Normal file
131
tests/neuron/deep_ep/test_engram.py
Normal file
@ -0,0 +1,131 @@
|
||||
"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py (CPU + gloo)."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
||||
import deep_ep # noqa: E402
|
||||
from deep_ep.utils.math import ceil_div, per_token_cast_to_fp8 # noqa: E402
|
||||
|
||||
# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
|
||||
# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
|
||||
DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
|
||||
|
||||
|
||||
def run(rank: int, num_ranks: int, group, args, use_fp8: bool):
|
||||
dtype = torch.float8_e4m3fn if use_fp8 else torch.bfloat16
|
||||
num_entries_per_layer = [int(v) for v in args.num_entries_per_layer.split(',')]
|
||||
num_layers = len(num_entries_per_layer)
|
||||
num_sf_packs = 2 * ceil_div(args.hidden, 128) if use_fp8 else 0
|
||||
|
||||
num_bytes, num_rdma_bytes = deep_ep.EngramBuffer.get_storage_size_hint(
|
||||
group, num_entries_per_layer, args.hidden, args.num_tokens,
|
||||
args.num_entries_per_token, dtype, num_sf_packs)
|
||||
num_qps, qp_depth, restrict = deep_ep.EngramBuffer.get_theoretical_config(
|
||||
group, num_layers, args.num_tokens, args.num_entries_per_token)
|
||||
assert (num_qps, qp_depth, restrict) == (0, 0, False)
|
||||
|
||||
buffer = deep_ep.EngramBuffer(
|
||||
group, num_gpu_bytes=num_bytes, num_rdma_storage_size=num_rdma_bytes,
|
||||
use_cpu_rdma_storage=True, num_allocated_qps=num_qps, qp_depth=qp_depth,
|
||||
restrict_rd_atomic=restrict, allow_hybrid_mode=True, explicitly_destroy=True)
|
||||
buffer.set_config(num_entries_per_layer, args.hidden, args.num_tokens,
|
||||
args.num_entries_per_token, dtype, num_sf_packs)
|
||||
|
||||
torch.manual_seed(7 + rank)
|
||||
local_storages, sfs = [], ([] if use_fp8 else None)
|
||||
for num_entries in num_entries_per_layer:
|
||||
local_bf16 = torch.randn((num_entries, args.hidden), dtype=torch.bfloat16, device=DEVICE)
|
||||
if use_fp8:
|
||||
local_storage, local_sf = per_token_cast_to_fp8(local_bf16)
|
||||
local_sf = local_sf.view(torch.int16).contiguous()
|
||||
sf = torch.empty((num_ranks * num_entries, num_sf_packs), dtype=local_sf.dtype, device=DEVICE)
|
||||
if num_ranks == 1:
|
||||
sf.copy_(local_sf)
|
||||
else:
|
||||
# gloo has no int16 collective; gather the bytes instead (DeepInfer's
|
||||
# NgramEmbedding.gather_sf gathers int16 directly on the accelerator).
|
||||
dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8), group=group)
|
||||
sfs.append(sf)
|
||||
else:
|
||||
local_storage = local_bf16
|
||||
local_storages.append(local_storage.contiguous())
|
||||
buffer.write(local_storages, sfs=sfs)
|
||||
|
||||
# Reference global tables
|
||||
global_bits = []
|
||||
for local_storage in local_storages:
|
||||
bits = torch.empty((num_ranks * local_storage.shape[0], args.hidden * dtype.itemsize),
|
||||
dtype=torch.uint8, device=DEVICE)
|
||||
local_bits = local_storage.view(torch.uint8).view(local_storage.shape[0], -1).contiguous()
|
||||
if num_ranks == 1:
|
||||
bits.copy_(local_bits)
|
||||
else:
|
||||
dist.all_gather_into_tensor(bits, local_bits, group=group)
|
||||
global_bits.append(bits)
|
||||
|
||||
def generate_indices():
|
||||
return torch.stack([
|
||||
torch.randint(0, num_ranks * num_entries, (args.num_tokens, args.num_entries_per_token),
|
||||
dtype=torch.int, device=DEVICE)
|
||||
for num_entries in num_entries_per_layer])
|
||||
|
||||
for col_major in ((False, True) if use_fp8 else (False,)):
|
||||
for layer_order in (range(num_layers), reversed(range(num_layers))):
|
||||
indices = generate_indices()
|
||||
ref_data = [global_bits[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1)
|
||||
for l in range(num_layers)]
|
||||
ref_sfs = [sfs[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1)
|
||||
for l in range(num_layers)] if use_fp8 else None
|
||||
|
||||
hooks = buffer.fetch(indices, use_tma_aligned_col_major_sf=col_major)
|
||||
assert len(hooks) == num_layers
|
||||
for l in layer_order:
|
||||
fetched = hooks[l]()
|
||||
data, fetched_sf = fetched if use_fp8 else (fetched, None)
|
||||
assert data.shape == (args.num_tokens, args.num_entries_per_token * args.hidden)
|
||||
assert torch.equal(ref_data[l], data.reshape(args.num_tokens, -1).contiguous().view(torch.uint8)), \
|
||||
f'data mismatch (layer {l}, fp8={use_fp8}, col_major={col_major})'
|
||||
if use_fp8:
|
||||
expected_stride = (1, args.num_tokens) if col_major else (fetched_sf.size(1), 1)
|
||||
assert fetched_sf.stride() == expected_stride, (fetched_sf.stride(), expected_stride)
|
||||
assert torch.equal(ref_sfs[l], fetched_sf), f'sf mismatch (layer {l})'
|
||||
buffer.destroy()
|
||||
|
||||
|
||||
def worker(local_rank: int, num_ranks: int, args):
|
||||
os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
|
||||
os.environ.setdefault('MASTER_PORT', str(args.port))
|
||||
if DEVICE == 'neuron':
|
||||
# Must be set before the Neuron runtime initializes (first device op) in this process.
|
||||
os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
|
||||
dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
|
||||
group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
|
||||
for use_fp8 in (False, True):
|
||||
run(group.rank(), num_ranks, group, args, use_fp8)
|
||||
if group.rank() == 0:
|
||||
print(f' ok: fp8={use_fp8}', flush=True)
|
||||
if group.rank() == 0:
|
||||
print(f'Engram tests passed (world_size={num_ranks})', flush=True)
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--num-processes', type=int, default=2)
|
||||
parser.add_argument('--num-entries-per-layer', type=str, default='37,64')
|
||||
parser.add_argument('--hidden', type=int, default=128)
|
||||
parser.add_argument('--num-tokens', type=int, default=11)
|
||||
parser.add_argument('--num-entries-per-token', type=int, default=3)
|
||||
parser.add_argument('--port', type=int, default=29541)
|
||||
args = parser.parse_args()
|
||||
if args.num_processes == 1:
|
||||
worker(0, 1, args)
|
||||
else:
|
||||
mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
|
||||
301
tests/neuron/deep_ep/test_ep.py
Normal file
301
tests/neuron/deep_ep/test_ep.py
Normal file
@ -0,0 +1,301 @@
|
||||
"""EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py.
|
||||
|
||||
CPU + gloo, no Neuron device. Run: python test_ep.py --num-processes {1,2,4}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
||||
import deep_ep # noqa: E402
|
||||
from deep_ep.utils.math import align, per_token_cast_to_fp8 # noqa: E402
|
||||
from refs import combine as ref_combine # noqa: E402
|
||||
from refs import dispatch as ref_dispatch # noqa: E402
|
||||
from refs import generate_pre_combine_data # noqa: E402
|
||||
|
||||
# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
|
||||
# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
|
||||
DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
|
||||
|
||||
|
||||
def enumerate_modes():
|
||||
for expert_alignment in (128, 1):
|
||||
for use_fp8 in (True, False):
|
||||
for num_bias in (0, 1, 2):
|
||||
for col_major in ((False, True) if use_fp8 else (False,)):
|
||||
yield expert_alignment, num_bias, use_fp8, col_major
|
||||
|
||||
|
||||
def launch(buffer, name, defer_epilogue, **params):
|
||||
result = getattr(buffer, name)(defer_epilogue=defer_epilogue, **params)
|
||||
if defer_epilogue:
|
||||
event = result
|
||||
return *event.current_stream_wait(release_handle=True), event
|
||||
result[-1].current_stream_wait(release_handle=True)
|
||||
return result
|
||||
|
||||
|
||||
def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync):
|
||||
expert_alignment, num_bias, use_fp8, col_major = mode
|
||||
tag = f'mode={mode} defer={defer_epilogue} cpu_sync={do_cpu_sync}'
|
||||
cap = args.num_tokens
|
||||
num_tokens = max(1, args.num_tokens - rank)
|
||||
hidden, num_topk, num_experts = args.hidden, args.num_topk, args.num_experts
|
||||
num_local_experts = num_experts // num_ranks
|
||||
|
||||
torch.manual_seed(1234 + rank)
|
||||
scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=DEVICE)
|
||||
topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
|
||||
topk_idx = topk_idx.to(torch.int64)
|
||||
if args.masked_ratio > 0:
|
||||
rand_mask = torch.rand_like(topk_idx, dtype=torch.float)
|
||||
topk_idx = topk_idx.masked_fill(rand_mask < args.masked_ratio, -1)
|
||||
topk_weights = topk_weights.masked_fill(topk_idx < 0, 0)
|
||||
topk_idx, topk_weights = topk_idx.contiguous(), topk_weights.contiguous()
|
||||
|
||||
x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE)
|
||||
if use_fp8:
|
||||
x_data, sf = per_token_cast_to_fp8(x)
|
||||
sf = sf.view(torch.int16)
|
||||
num_sf_packs = sf.size(1)
|
||||
if col_major:
|
||||
sf_storage = torch.empty((num_tokens * num_sf_packs,), dtype=torch.int16, device=DEVICE)
|
||||
input_sf = torch.as_strided(sf_storage, (num_tokens, num_sf_packs), (1, num_tokens))
|
||||
input_sf.copy_(sf)
|
||||
else:
|
||||
input_sf = sf
|
||||
x_in = (x_data, input_sf)
|
||||
else:
|
||||
x_in = x
|
||||
bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) if num_bias == 1 else None
|
||||
if num_bias == 2:
|
||||
bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) for _ in range(2))
|
||||
|
||||
stats = torch.zeros((num_local_experts,), dtype=torch.int32, device=DEVICE)
|
||||
dispatch_args = dict(
|
||||
x=x_in, topk_idx=topk_idx, topk_weights=topk_weights,
|
||||
num_sms=7, num_max_tokens_per_rank=cap, num_experts=num_experts,
|
||||
expert_alignment=expert_alignment, async_with_compute_stream=True,
|
||||
do_handle_copy=False, do_cpu_sync=do_cpu_sync, do_expand=True,
|
||||
do_zero_padding=True, use_tma_aligned_col_major_sf=col_major)
|
||||
|
||||
ref_x = (x_data.view(torch.uint8), input_sf) if use_fp8 else x
|
||||
(ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights,
|
||||
ref_recv_src_token_idx, ref_num_recv_tokens_per_rank) = ref_dispatch(
|
||||
ref_x, topk_idx, topk_weights, cap, num_experts)
|
||||
|
||||
recv_x, recv_topk_idx, recv_weights, handle, _ = launch(
|
||||
buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=stats, **dispatch_args)
|
||||
assert recv_topk_idx is None
|
||||
assert recv_weights is not None
|
||||
assert handle.num_sms == 7
|
||||
|
||||
recv_x_data, recv_sf = recv_x if use_fp8 else (recv_x, None)
|
||||
check_x = recv_x_data.view(torch.uint8) if use_fp8 else recv_x_data
|
||||
|
||||
num_recv_tokens = handle.psum_num_recv_tokens_per_rank[-1].item()
|
||||
num_expanded_tokens = align(handle.psum_num_recv_tokens_per_expert[-1].item(), expert_alignment)
|
||||
expected_num_recv = num_recv_tokens if do_cpu_sync else num_ranks * cap
|
||||
assert handle.num_recv_tokens == expected_num_recv, (handle.num_recv_tokens, expected_num_recv)
|
||||
assert handle.recv_src_metadata.size(0) == expected_num_recv
|
||||
if use_fp8:
|
||||
expected_stride = (1, recv_x_data.size(0)) if col_major else (num_sf_packs, 1)
|
||||
assert recv_sf.stride() == expected_stride, (recv_sf.stride(), expected_stride)
|
||||
|
||||
ref_recv_x_data, ref_recv_sf = ref_recv_x if use_fp8 else (ref_recv_x, None)
|
||||
assert num_recv_tokens == ref_recv_x_data.size(0), (num_recv_tokens, ref_recv_x_data.size(0))
|
||||
|
||||
metadata = handle.recv_src_metadata[:num_recv_tokens]
|
||||
recv_src_token_idx = metadata[:, -2]
|
||||
recv_expanded_slots = metadata[:, :-2]
|
||||
ref_order = torch.argsort(ref_recv_src_token_idx)
|
||||
metadata_order = torch.argsort(recv_src_token_idx)
|
||||
assert torch.equal(recv_src_token_idx[metadata_order], ref_recv_src_token_idx[ref_order])
|
||||
|
||||
def reorder(src):
|
||||
out = torch.empty_like(src)
|
||||
out[metadata_order] = src[ref_order]
|
||||
return out
|
||||
|
||||
ref_x_by_meta = reorder(ref_recv_x_data)
|
||||
ref_sf_by_meta = reorder(ref_recv_sf) if ref_recv_sf is not None else None
|
||||
ref_idx_by_meta = reorder(ref_recv_topk_idx)
|
||||
ref_w_by_meta = reorder(ref_recv_topk_weights)
|
||||
|
||||
recv_valid_mask = recv_expanded_slots >= 0
|
||||
assert torch.equal(recv_valid_mask, ref_idx_by_meta >= 0)
|
||||
|
||||
valid_slots = recv_expanded_slots[recv_valid_mask].long()
|
||||
assert valid_slots.numel() == handle.num_unaligned_recv_tokens_per_expert.sum().item()
|
||||
if valid_slots.numel() > 0:
|
||||
assert valid_slots.min().item() >= 0 and valid_slots.max().item() < num_expanded_tokens
|
||||
assert torch.unique(valid_slots).numel() == valid_slots.numel()
|
||||
expected_x = ref_x_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
|
||||
assert torch.equal(check_x[valid_slots], expected_x)
|
||||
if ref_sf_by_meta is not None:
|
||||
expected_sf = ref_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
|
||||
assert torch.equal(recv_sf[valid_slots], expected_sf)
|
||||
assert torch.equal(recv_weights[valid_slots], ref_w_by_meta[recv_valid_mask])
|
||||
|
||||
# Master top-k slot column
|
||||
topk_range = torch.arange(num_topk, device=DEVICE, dtype=metadata.dtype)
|
||||
master = torch.where(recv_valid_mask, topk_range.view(1, num_topk),
|
||||
torch.full_like(recv_expanded_slots, -1)).max(dim=1).values
|
||||
assert torch.equal(metadata[:, -1], master)
|
||||
|
||||
# Per-rank and per-expert counts
|
||||
psum_rank = [0] + handle.psum_num_recv_tokens_per_rank.tolist()
|
||||
for r in range(num_ranks):
|
||||
assert psum_rank[r + 1] - psum_rank[r] == ref_num_recv_tokens_per_rank[r].item()
|
||||
psum_expert = [0] + handle.psum_num_recv_tokens_per_expert.tolist()
|
||||
for e in range(num_local_experts):
|
||||
count = psum_expert[e + 1] - align(psum_expert[e], expert_alignment)
|
||||
ref_count = (ref_recv_topk_idx == e).sum().item()
|
||||
assert count == ref_count, (e, count, ref_count)
|
||||
assert stats[e].item() == ref_count
|
||||
assert handle.num_unaligned_recv_tokens_per_expert[e].item() == ref_count
|
||||
if do_cpu_sync:
|
||||
assert handle.num_recv_tokens_per_expert_list[e] == align(ref_count, expert_alignment)
|
||||
expert_start = align(psum_expert[e], expert_alignment)
|
||||
expert_end = psum_expert[e + 1]
|
||||
expert_slots = recv_expanded_slots[ref_idx_by_meta == e]
|
||||
assert expert_slots.numel() == ref_count
|
||||
if ref_count > 0:
|
||||
assert torch.equal(torch.sort(expert_slots).values,
|
||||
torch.arange(expert_start, expert_end, device=DEVICE,
|
||||
dtype=expert_slots.dtype))
|
||||
padding_end = align(expert_end, expert_alignment)
|
||||
assert torch.count_nonzero(recv_x_data[expert_end:padding_end].view(torch.uint8)).item() == 0
|
||||
if recv_sf is not None:
|
||||
assert torch.count_nonzero(recv_sf[expert_end:padding_end]).item() == 0
|
||||
|
||||
# --- Combine ---
|
||||
local_y = generate_pre_combine_data(metadata[:, -2], cap, num_topk, hidden)
|
||||
input_for_combine = torch.empty((recv_x_data.size(0) + 1, hidden), dtype=torch.bfloat16, device=DEVICE)
|
||||
input_for_combine[metadata[:, :-2].long().flatten()] = local_y.view(-1, hidden)
|
||||
input_for_combine = input_for_combine[:-1].contiguous()
|
||||
|
||||
combined_x, combined_weights, _ = launch(
|
||||
buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle,
|
||||
topk_weights=recv_weights, bias=bias, num_sms=7, async_with_compute_stream=True)
|
||||
|
||||
ref_y = generate_pre_combine_data(rank * cap + torch.arange(num_tokens, device=DEVICE), cap, num_topk, hidden)
|
||||
ref_y[topk_idx == -1] = 0
|
||||
ref_combined_x = ref_combine(ref_y, topk_idx, num_experts, bias, reduce_in_local=True)
|
||||
max_diff = (combined_x.float() - ref_combined_x.float()).abs().max().item()
|
||||
assert torch.equal(combined_x, ref_combined_x), f'{tag} {max_diff=}'
|
||||
assert torch.equal(combined_weights, topk_weights)
|
||||
|
||||
if num_bias == 0 and not col_major:
|
||||
combined_x2, combined_weights2, _ = launch(
|
||||
buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle,
|
||||
bias=bias, num_sms=7, async_with_compute_stream=True)
|
||||
assert torch.equal(combined_x2, ref_combined_x)
|
||||
assert combined_weights2 is None
|
||||
|
||||
# --- Cached dispatch replay ---
|
||||
cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=handle)
|
||||
del cached_args['topk_idx']
|
||||
cached_tensors = (handle.psum_num_recv_tokens_per_rank, handle.psum_num_recv_tokens_per_expert,
|
||||
handle.num_unaligned_recv_tokens_per_expert, handle.recv_src_metadata)
|
||||
cached_copies = [t.clone() for t in cached_tensors]
|
||||
padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=DEVICE)
|
||||
padding_mask[valid_slots] = False
|
||||
for send_weights in (True, False):
|
||||
updated_x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) * (8 if send_weights else 0.125)
|
||||
if use_fp8:
|
||||
updated_data, updated_sf = per_token_cast_to_fp8(updated_x)
|
||||
updated_sf = updated_sf.view(torch.int16)
|
||||
if col_major:
|
||||
updated_sf = updated_sf.t().contiguous().t()
|
||||
updated_in = (updated_data, updated_sf)
|
||||
else:
|
||||
updated_in = updated_x
|
||||
updated_weights = torch.randn_like(topk_weights).masked_fill(topk_idx < 0, 0)
|
||||
ref_updated = (updated_data.view(torch.uint8), updated_sf) if use_fp8 else updated_x
|
||||
ref_cached_x, _, _, _, _ = ref_dispatch(ref_updated, topk_idx, updated_weights, cap, num_experts)
|
||||
cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch(
|
||||
buffer, 'dispatch', defer_epilogue,
|
||||
**dict(cached_args, x=updated_in,
|
||||
topk_weights=updated_weights if send_weights else None))
|
||||
assert cached_handle is handle and cached_topk_idx is None
|
||||
for tensor, saved in zip(cached_tensors, cached_copies):
|
||||
assert torch.equal(tensor, saved)
|
||||
cached_data, cached_sf = cached_x if use_fp8 else (cached_x, None)
|
||||
ref_cached_data, ref_cached_sf = ref_cached_x if use_fp8 else (ref_cached_x, None)
|
||||
cached_check = cached_data.view(torch.uint8) if use_fp8 else cached_data
|
||||
ref_cached_by_meta = reorder(ref_cached_data)
|
||||
assert torch.equal(cached_check[valid_slots],
|
||||
ref_cached_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
|
||||
assert torch.count_nonzero(cached_check[:num_expanded_tokens][padding_mask]).item() == 0
|
||||
if cached_sf is not None:
|
||||
assert cached_sf.stride() == expected_stride
|
||||
ref_cached_sf_by_meta = reorder(ref_cached_sf)
|
||||
assert torch.equal(cached_sf[valid_slots],
|
||||
ref_cached_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
|
||||
if send_weights:
|
||||
assert torch.equal(cached_w[valid_slots], reorder(updated_weights_ref(
|
||||
updated_weights, topk_idx, cap, num_experts))[recv_valid_mask])
|
||||
else:
|
||||
assert cached_w is None
|
||||
|
||||
|
||||
def updated_weights_ref(updated_weights, topk_idx, cap, num_experts):
|
||||
_, _, ref_w, _, _ = ref_dispatch(
|
||||
torch.zeros((topk_idx.size(0), 8), dtype=torch.bfloat16, device=DEVICE),
|
||||
topk_idx, updated_weights, cap, num_experts)
|
||||
return ref_w
|
||||
|
||||
|
||||
def worker(local_rank: int, num_ranks: int, args: argparse.Namespace):
|
||||
os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
|
||||
os.environ.setdefault('MASTER_PORT', str(args.port))
|
||||
if DEVICE == 'neuron':
|
||||
# Must be set before the Neuron runtime initializes (first device op) in this process.
|
||||
os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
|
||||
dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
|
||||
group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
|
||||
rank = group.rank()
|
||||
|
||||
buffer = deep_ep.EPBuffer(
|
||||
group, num_max_tokens_per_rank=args.num_tokens, hidden=args.hidden,
|
||||
num_topk=args.num_topk, use_fp8_dispatch=True, allow_hybrid_mode=True,
|
||||
allow_multiple_reduction=True, explicitly_destroy=True)
|
||||
assert buffer.num_bytes > 0
|
||||
assert deep_ep.get_physical_domain_size(group) == (1, num_ranks)
|
||||
|
||||
for mode in enumerate_modes():
|
||||
for defer_epilogue in (True, False):
|
||||
for do_cpu_sync in (True, False):
|
||||
check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync)
|
||||
if rank == 0:
|
||||
print(f' ok: expert_alignment={mode[0]} num_bias={mode[1]} fp8={mode[2]} col_major={mode[3]}',
|
||||
flush=True)
|
||||
|
||||
buffer.destroy()
|
||||
if rank == 0:
|
||||
print(f'EP tests passed (world_size={num_ranks})', flush=True)
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--num-processes', type=int, default=2)
|
||||
parser.add_argument('--num-tokens', type=int, default=17)
|
||||
parser.add_argument('--hidden', type=int, default=256)
|
||||
parser.add_argument('--num-topk', type=int, default=4)
|
||||
parser.add_argument('--num-experts', type=int, default=16)
|
||||
parser.add_argument('--masked-ratio', type=float, default=0.2)
|
||||
parser.add_argument('--port', type=int, default=29531)
|
||||
args = parser.parse_args()
|
||||
assert args.num_experts % args.num_processes == 0
|
||||
if args.num_processes == 1:
|
||||
worker(0, 1, args)
|
||||
else:
|
||||
mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
|
||||
693
tests/neuron/deep_gemm/test_deep_gemm_neuron.py
Normal file
693
tests/neuron/deep_gemm/test_deep_gemm_neuron.py
Normal file
@ -0,0 +1,693 @@
|
||||
"""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")
|
||||
52
tests/neuron/deep_gemm/test_interop_tile_kernels.py
Normal file
52
tests/neuron/deep_gemm/test_interop_tile_kernels.py
Normal file
@ -0,0 +1,52 @@
|
||||
"""Cross-check: tile_kernels.quant.per_token_cast output feeds deep_gemm's FP8 GEMM.
|
||||
|
||||
This is the exact ReplicatedLinear.apply_weight_fp8 path
|
||||
(deepinfer/model_executor/layers/linear.py) with CPU tensors.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
||||
|
||||
import torch
|
||||
|
||||
import deep_gemm
|
||||
import tile_kernels
|
||||
|
||||
CAST_PARAMS = dict(fmt="e4m3", num_per_channels=32, round_sf=True,
|
||||
use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
||||
|
||||
|
||||
def calc_diff(x, y):
|
||||
x, y = x.double(), y.double()
|
||||
denom = (x * x + y * y).sum().item()
|
||||
return 0.0 if denom == 0 else 1 - 2 * (x * y).sum().item() / denom
|
||||
|
||||
|
||||
m, n, k = 48, 128, 256
|
||||
x = torch.randn(m, k, dtype=torch.bfloat16)
|
||||
w = torch.randn(n, k, dtype=torch.bfloat16)
|
||||
|
||||
xq, x_sf = tile_kernels.quant.per_token_cast(x, **CAST_PARAMS)
|
||||
wq, w_sf = tile_kernels.quant.per_token_cast(w, **CAST_PARAMS)
|
||||
print(f"{xq.dtype=} {x_sf.dtype=} {x_sf.shape=} {x_sf.stride()=}")
|
||||
print(f"{wq.dtype=} {w_sf.dtype=} {w_sf.shape=} {w_sf.stride()=}")
|
||||
|
||||
out = torch.empty(m, n, dtype=torch.bfloat16)
|
||||
deep_gemm.fp8_fp4_gemm_nt((xq, x_sf), (wq, w_sf), out, None, recipe=(1, 1, 32))
|
||||
|
||||
ref = x.float() @ w.float().T
|
||||
diff = calc_diff(out.float(), ref)
|
||||
print(f"diff vs bf16 reference: {diff:.3e}")
|
||||
assert diff < 2e-3, diff
|
||||
|
||||
# The dequantized operands must agree with the GEMM result exactly.
|
||||
from deep_gemm._quant import dequantize # noqa: E402
|
||||
|
||||
x_deq = dequantize(xq, deep_gemm.transform_sf_into_required_layout(x_sf, m, k, (1, 1, 32), is_sfa=True))
|
||||
w_deq = dequantize(wq, deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 1, 32), is_sfa=False))
|
||||
exact = (x_deq @ w_deq.T).to(torch.bfloat16)
|
||||
print("exact match with dequantized matmul:", torch.equal(out, exact))
|
||||
print(f"quantization error (dequant vs bf16): {calc_diff(exact.float(), ref):.3e}")
|
||||
print("OK")
|
||||
128
tests/neuron/device/test_device_primitives.py
Normal file
128
tests/neuron/device/test_device_primitives.py
Normal file
@ -0,0 +1,128 @@
|
||||
"""Native torch-neuronx behaviours the Neuron port depends on (or works around).
|
||||
|
||||
Run on a NeuronCore (no DeepInfer import, plain torch-neuronx):
|
||||
|
||||
NEURON_RT_VISIBLE_CORES=0 python tests/neuron/device/test_device_primitives.py
|
||||
|
||||
Each check compares the device result with the host. ``expect="ok"`` checks must pass; ``expect="known"``
|
||||
checks document an observed torch-neuronx 2.14 limitation the port works around -- if one of those starts
|
||||
passing, the workaround can go. Exit status is non-zero only when an ``ok`` check fails.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
DEV = torch.device("neuron", 0)
|
||||
RESULTS: list[tuple[str, str, str]] = []
|
||||
|
||||
|
||||
def check(name: str, expect: str, fn) -> None:
|
||||
try:
|
||||
ok = bool(fn())
|
||||
status = "ok" if ok else "DIFF"
|
||||
except Exception as e: # noqa: BLE001
|
||||
status = f"EXC {type(e).__name__}: {str(e).splitlines()[0][:100]}"
|
||||
RESULTS.append((name, expect, status))
|
||||
|
||||
|
||||
def same(dev: torch.Tensor, host: torch.Tensor) -> bool:
|
||||
d = dev.cpu()
|
||||
if d.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
|
||||
d, host = d.view(torch.uint8), host.view(torch.uint8)
|
||||
return torch.equal(d, host)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(11, 512, dtype=torch.bfloat16)
|
||||
xd = x.to(DEV)
|
||||
|
||||
# dtypes used by the kernels
|
||||
f8 = torch.randn(64, 64).clamp(-400, 400).to(torch.float8_e4m3fn)
|
||||
check("fp8 e4m3 H2D/D2H bit-exact", "ok", lambda: same(f8.to(DEV), f8))
|
||||
check("fp8 -> bf16 upcast on device", "ok", lambda: same(f8.to(DEV).to(torch.bfloat16), f8.to(torch.bfloat16)))
|
||||
e8 = torch.randint(100, 140, (8, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
|
||||
check("e8m0 H2D bit-exact", "ok", lambda: same(e8.to(DEV), e8))
|
||||
b = torch.randint(-128, 127, (4, 16), dtype=torch.int8)
|
||||
check("int8 nibble unpack", "ok", lambda: same(b.to(DEV) & 0xF, b & 0xF) and same((b.to(DEV) >> 4) & 0xF, (b >> 4) & 0xF))
|
||||
s = torch.randint(118, 130, (4, 16), dtype=torch.uint8)
|
||||
check("ue8m0 bits -> fp32 (<<23, view)", "ok",
|
||||
lambda: same((s.to(DEV).to(torch.int32) << 23).view(torch.float32), (s.to(torch.int32) << 23).view(torch.float32)))
|
||||
u = torch.tensor([1, 2**40 + 7, 2**63 + 5], dtype=torch.uint64)
|
||||
check("uint64 round trip", "known", lambda: same(u.to(DEV), u))
|
||||
check("uint64 add", "known", lambda: same(u.to(DEV) + 3, u + 3))
|
||||
|
||||
# views / in-place semantics
|
||||
check("column-slice view -> host", "ok", lambda: same(xd[:, 256:], x[:, 256:]))
|
||||
check("transpose view -> host", "ok", lambda: same(xd.t(), x.t()))
|
||||
|
||||
def inplace_slice():
|
||||
o, oc = torch.zeros(11, 1024, dtype=torch.bfloat16, device=DEV), torch.zeros(11, 1024, dtype=torch.bfloat16)
|
||||
o[:, 512:].copy_(xd)
|
||||
oc[:, 512:].copy_(x)
|
||||
o[:, :512] += xd
|
||||
oc[:, :512] += x
|
||||
return same(o, oc)
|
||||
|
||||
check("copy_ / += into column slice", "ok", inplace_slice)
|
||||
buf = torch.randint(0, 256, (2 * 902880,), dtype=torch.uint8)
|
||||
bufd = buf.to(DEV)
|
||||
blk, row = torch.tensor([0, 0, 1, 1, 0]), torch.tensor([0, 5, 7, 128, 3])
|
||||
|
||||
def paged(t: torch.Tensor) -> torch.Tensor: # KV-cache style strided record view
|
||||
return t.as_strided((2, 129, 512), (902880, 528, 1))
|
||||
|
||||
check("uint8 strided view advanced indexing", "ok",
|
||||
lambda: same(paged(bufd)[blk.to(DEV), row.to(DEV)], paged(buf)[blk, row]))
|
||||
check("fp8 strided view advanced indexing", "known",
|
||||
lambda: same(paged(bufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(buf).view(torch.float8_e4m3fn)[blk, row]))
|
||||
idx = torch.tensor([1, 1, 4])
|
||||
check("index_add_ (duplicates)", "ok", lambda: same(torch.zeros(11, 512, dtype=torch.bfloat16, device=DEV).index_add_(0, idx.to(DEV), xd[:3]),
|
||||
torch.zeros(11, 512, dtype=torch.bfloat16).index_add_(0, idx, x[:3])))
|
||||
|
||||
def inference_version_counter():
|
||||
with torch.inference_mode():
|
||||
t = torch.topk(xd.float(), 4).indices
|
||||
return t.is_inference() or t._version >= 0
|
||||
|
||||
check("inference-mode output has is_inference() or a version counter", "known", inference_version_counter)
|
||||
|
||||
# numerics
|
||||
a, w = torch.randn(64, 256), torch.randn(256, 128)
|
||||
ref = a.double() @ w.double()
|
||||
rel = lambda r: ((r.cpu().double() - ref).norm() / ref.norm()).item() # noqa: E731
|
||||
check("fp32 matmul rel err < 1e-5 (vs fp64)", "ok", lambda: rel(a.to(DEV) @ w.to(DEV)) < 1e-5)
|
||||
check("bf16 topk indices == host", "ok", lambda: torch.equal(
|
||||
torch.topk(xd[0], 64).indices.cpu().sort().values, torch.topk(x[0], 64).indices.sort().values))
|
||||
|
||||
# single-rank process group (EP=1 bring-up)
|
||||
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
||||
os.environ.setdefault("MASTER_PORT", "29761")
|
||||
dist.init_process_group("gloo", rank=0, world_size=1)
|
||||
g = dist.new_group([0], backend="neuron")
|
||||
check("neuron PG size 1: all_reduce", "ok", lambda: (lambda t: (dist.all_reduce(t, group=g), same(t, x))[1])(xd.clone()))
|
||||
check("neuron PG size 1: all_gather_into_tensor", "ok",
|
||||
lambda: (lambda o: (dist.all_gather_into_tensor(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
|
||||
check("neuron PG size 1: all_to_all_single", "known",
|
||||
lambda: (lambda o: (dist.all_to_all_single(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
|
||||
dist.destroy_process_group()
|
||||
|
||||
failed = 0
|
||||
for name, expect, status in RESULTS:
|
||||
if expect == "ok":
|
||||
tag = "PASS" if status == "ok" else "FAIL"
|
||||
failed += tag == "FAIL"
|
||||
else:
|
||||
tag = "KNOWN-ISSUE" if status != "ok" else "FIXED-UPSTREAM?"
|
||||
print(f"{tag:16s} {name:62s} {status}")
|
||||
print(f"\n{failed} unexpected failure(s)")
|
||||
return 1 if failed else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
89
tests/neuron/flash_mla/ref_attn.py
Normal file
89
tests/neuron/flash_mla/ref_attn.py
Normal file
@ -0,0 +1,89 @@
|
||||
"""Reference sparse attention ported from vendor/FlashMLA/tests/ref.py (no TestParam wrappers)."""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def ref_sparse_attn_fwd(q, kv, indices, sm_scale, topk_length=None, attn_sink=None, d_v=512, rms=None):
|
||||
"""q [s_q,h_q,d], kv [s_kv,1,d], indices [s_q,1,topk] -> (out bf16, max_logits, lse)."""
|
||||
s_q, h_q, d_qk = q.shape
|
||||
s_kv = kv.shape[0]
|
||||
topk = indices.shape[-1]
|
||||
idx = indices.clone().squeeze(1)
|
||||
if topk_length is not None:
|
||||
mask = torch.arange(topk).unsqueeze(0).broadcast_to(s_q, topk) >= topk_length.unsqueeze(1)
|
||||
idx.masked_fill_(mask, -1)
|
||||
invalid = (idx < 0) | (idx >= s_kv)
|
||||
idx = idx.masked_fill(invalid, 0)
|
||||
|
||||
qf = q.float()
|
||||
gathered = kv.view(s_kv, d_qk).index_select(0, idx.flatten().long()).reshape(s_q, topk, d_qk).float()
|
||||
P = qf @ gathered.transpose(1, 2)
|
||||
if rms is not None:
|
||||
P = P * (sm_scale * rms.unsqueeze(-1))
|
||||
else:
|
||||
P = P * sm_scale
|
||||
P.masked_fill_(invalid.unsqueeze(1).broadcast_to(P.shape), float("-inf"))
|
||||
|
||||
orig_lse = torch.logsumexp(P, dim=-1)
|
||||
max_logits = P.max(dim=-1).values
|
||||
|
||||
lse_for_o = orig_lse
|
||||
if attn_sink is not None:
|
||||
lse_for_o = torch.logsumexp(
|
||||
torch.stack([orig_lse.view(s_q, h_q), attn_sink.broadcast_to(s_q, h_q)], dim=0), dim=0
|
||||
)
|
||||
lse_for_o = lse_for_o.clone()
|
||||
lse_for_o.masked_fill_(lse_for_o == float("-inf"), float("+inf"))
|
||||
s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
|
||||
out = s_for_o @ gathered[..., :d_v]
|
||||
|
||||
orig_lse = orig_lse.clone()
|
||||
orig_lse.masked_fill_(orig_lse == float("-inf"), float("+inf"))
|
||||
return out.to(torch.bfloat16), max_logits, orig_lse
|
||||
|
||||
|
||||
def ref_sparse_attn_decode(q, scopes, sm_scale, attn_sink=None, d_v=512, rms=None):
|
||||
"""q [b,s_q,h_q,d]; scopes = list of (blocked_k_bf16 [n_tokens, d], indices [b,s_q,topk], topk_length|None).
|
||||
|
||||
Returns (out bf16 [b,s_q,h_q,d_v], lse [b,h_q,s_q]).
|
||||
"""
|
||||
b, s_q, h_q, d_qk = q.shape
|
||||
gathered_all, invalid_all = [], []
|
||||
for blocked_k, indices, topk_length in scopes:
|
||||
topk = indices.shape[-1]
|
||||
fixed = torch.clamp_min(indices, 0)
|
||||
g = blocked_k.view(-1, d_qk).index_select(0, fixed.reshape(-1).long()).view(b, s_q, topk, d_qk)
|
||||
invalid = indices == -1
|
||||
if topk_length is not None:
|
||||
invalid = invalid | (
|
||||
torch.arange(0, topk).view(1, 1, topk).broadcast_to(b, s_q, topk) >= topk_length.view(b, 1, 1)
|
||||
)
|
||||
gathered_all.append(g)
|
||||
invalid_all.append(invalid)
|
||||
gathered = torch.cat(gathered_all, dim=2)
|
||||
invalid = torch.cat(invalid_all, dim=2)
|
||||
|
||||
gathered = gathered.reshape(b * s_q, -1, d_qk).float()
|
||||
gathered.masked_fill_(gathered != gathered, 0.0)
|
||||
qf = q.float().view(b * s_q, h_q, d_qk)
|
||||
w = qf @ gathered.transpose(-1, -2)
|
||||
if rms is not None:
|
||||
w = w * (sm_scale * rms.view(b * s_q, h_q).unsqueeze(-1))
|
||||
else:
|
||||
w = w * sm_scale
|
||||
w[invalid.view(b * s_q, 1, -1).broadcast_to(b * s_q, h_q, invalid.size(-1))] = float("-inf")
|
||||
lse = w.logsumexp(dim=-1)
|
||||
w = torch.exp(w - lse.unsqueeze(-1))
|
||||
out = w @ gathered[..., :d_v]
|
||||
out = out.view(b, s_q, h_q, d_v)
|
||||
lse = lse.view(b, s_q, h_q)
|
||||
|
||||
if attn_sink is not None:
|
||||
out = out * (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(-1)
|
||||
|
||||
lonely = lse == float("-inf")
|
||||
out = out.clone()
|
||||
out[lonely.unsqueeze(-1).broadcast_to(b, s_q, h_q, d_v)] = 0.0
|
||||
lse = lse.clone()
|
||||
lse[lonely] = float("+inf")
|
||||
return out.to(torch.bfloat16), lse.transpose(1, 2)
|
||||
89
tests/neuron/flash_mla/ref_quant.py
Normal file
89
tests/neuron/flash_mla/ref_quant.py
Normal file
@ -0,0 +1,89 @@
|
||||
"""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)
|
||||
119
tests/neuron/flash_mla/test_deep_select.py
Normal file
119
tests/neuron/flash_mla/test_deep_select.py
Normal file
@ -0,0 +1,119 @@
|
||||
"""CPU correctness tests for the pure-torch deep_select vendor package.
|
||||
|
||||
The checks are ported from vendor/DeepSelect/tests/test.py (NaN guard, index range, uniqueness,
|
||||
gathered value, topk condition, sorted index / sorted value).
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from deep_select.interface import get_stride_requirement, topk as ds_topk
|
||||
|
||||
torch.manual_seed(0)
|
||||
FAILS = []
|
||||
NAN_SENTINEL = 0x3F3F3F3F
|
||||
|
||||
|
||||
def chk(name, cond):
|
||||
print(f"[{'ok ' if cond else 'FAIL'}] {name}")
|
||||
if not cond:
|
||||
FAILS.append(name)
|
||||
|
||||
|
||||
def run_case(b, vocab, k, dtype, idx_dtype, sorted_value, sorted_index, return_value,
|
||||
use_end, use_offset, with_nan=False, use_out=False):
|
||||
tag = (f"b{b} v{vocab} k{k} {str(dtype).split('.')[-1]} {str(idx_dtype).split('.')[-1]} "
|
||||
f"sv{int(sorted_value)} si{int(sorted_index)} rv{int(return_value)} "
|
||||
f"end{int(use_end)} off{int(use_offset)} nan{int(with_nan)} out{int(use_out)}")
|
||||
x = (torch.randn(b, vocab) * 3).to(dtype)
|
||||
end = torch.randint(1, vocab + 1, (b,), dtype=torch.int32) if use_end else None
|
||||
if use_end:
|
||||
end[0] = min(k // 2 + 1, vocab) # a short row, fewer valid entries than k
|
||||
offset = torch.randint(0, 1000, (b,), dtype=torch.int32) if use_offset else None
|
||||
nan_rows = torch.zeros(b, dtype=torch.bool)
|
||||
if with_nan:
|
||||
nan_rows[0] = True
|
||||
x[0, 0] = float("nan")
|
||||
out_buf = torch.empty((b, k), dtype=idx_dtype) if use_out else None
|
||||
|
||||
val, idx = ds_topk(
|
||||
x, k, sorted=sorted_value, end=end, indices_type=idx_dtype, sorted_index=sorted_index,
|
||||
output_idx=out_buf, output_idx_offset=offset, idx_oob_fill_value=-1,
|
||||
value_oob_fill_value=float("-inf"), return_value=return_value, abort_when_nan_found=False,
|
||||
)
|
||||
chk(f"{tag}: dtypes", idx.dtype is idx_dtype and (val is None) == (not return_value)
|
||||
and (val is None or val.dtype is dtype))
|
||||
chk(f"{tag}: shapes", tuple(idx.shape) == (b, k) and (val is None or tuple(val.shape) == (b, k)))
|
||||
if use_out:
|
||||
chk(f"{tag}: writes output_idx in place", idx.data_ptr() == out_buf.data_ptr())
|
||||
|
||||
valid_len = (torch.full((b,), vocab, dtype=torch.int64) if end is None else end.to(torch.int64))
|
||||
counts = torch.clamp(valid_len, max=k)
|
||||
sel = torch.arange(k).unsqueeze(0) < counts.unsqueeze(1)
|
||||
|
||||
# NaN guard: whole row filled with the sentinel (Ascend behaviour)
|
||||
if with_nan:
|
||||
chk(f"{tag}: NaN row sentinel", bool((idx[nan_rows] == NAN_SENTINEL).all()))
|
||||
sel = sel & ~nan_rows.unsqueeze(1)
|
||||
|
||||
raw = idx.to(torch.int64)
|
||||
if offset is not None:
|
||||
raw = raw - offset.to(torch.int64).unsqueeze(1)
|
||||
in_range = (raw >= 0) & (raw < valid_len.unsqueeze(1))
|
||||
chk(f"{tag}: index range", bool((in_range | ~sel).all()))
|
||||
chk(f"{tag}: padding fill", bool((idx[~sel & ~nan_rows.unsqueeze(1)] == -1).all()))
|
||||
|
||||
srt = torch.sort(torch.where(sel, raw, torch.iinfo(torch.int64).max), dim=1).values
|
||||
dup = (srt[:, 1:] == srt[:, :-1]) & sel[:, 1:]
|
||||
chk(f"{tag}: unique index", not bool(dup.any()))
|
||||
|
||||
safe = torch.where(sel & in_range, raw, torch.zeros_like(raw))
|
||||
gathered = x.gather(1, safe)
|
||||
rows = ~nan_rows
|
||||
if rows.any():
|
||||
if val is not None:
|
||||
exp = gathered[rows].masked_fill(~sel[rows], float("-inf"))
|
||||
chk(f"{tag}: gathered value bitwise", torch.equal(val[rows], exp))
|
||||
gv = gathered[rows].float()
|
||||
smin = gv.masked_fill(~sel[rows], float("inf")).amin(dim=1)
|
||||
unsel = x[rows].float().clone()
|
||||
unsel.scatter_(1, safe[rows], float("-inf"))
|
||||
col = torch.arange(vocab).unsqueeze(0)
|
||||
unsel = torch.where(col >= valid_len[rows].unsqueeze(1), torch.full_like(unsel, float("-inf")), unsel)
|
||||
unsel = torch.where(unsel != unsel, torch.full_like(unsel, float("-inf")), unsel)
|
||||
chk(f"{tag}: topk condition", bool((smin >= unsel.amax(dim=1)).all()))
|
||||
|
||||
ordered = sel[:, 1:]
|
||||
if sorted_index:
|
||||
ok = (idx[:, 1:].to(torch.int64) >= idx[:, :-1].to(torch.int64)) | ~ordered
|
||||
chk(f"{tag}: sorted index", bool(ok.all()))
|
||||
if sorted_value and val is not None:
|
||||
ok = (val[:, :-1].float() >= val[:, 1:].float()) | ~ordered
|
||||
chk(f"{tag}: sorted value", bool(ok.all()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
chk("stride requirement == (32, 32)", get_stride_requirement() == (32, 32))
|
||||
|
||||
combos = [(False, False, False), (False, False, True), (True, False, True),
|
||||
(False, True, True), (False, True, False)]
|
||||
for dtype in (torch.float32, torch.bfloat16):
|
||||
for idx_dtype in (torch.int32, torch.int64):
|
||||
for sv, si, rv in combos:
|
||||
if sv and dtype is torch.bfloat16:
|
||||
continue # sorted_value is fp32-only upstream
|
||||
run_case(5, 97, 8, dtype, idx_dtype, sv, si, rv, use_end=True, use_offset=True)
|
||||
for use_end, use_offset in itertools.product((False, True), repeat=2):
|
||||
run_case(3, 64, 16, torch.float32, torch.int32, False, True, False, use_end, use_offset)
|
||||
run_case(4, 33, 8, torch.float32, torch.int32, False, True, False, True, True, with_nan=True)
|
||||
run_case(4, 33, 8, torch.float32, torch.int32, True, False, True, False, False, with_nan=True)
|
||||
run_case(3, 16, 8, torch.float32, torch.int32, False, True, False, True, True, use_out=True)
|
||||
# k larger than the vocab / every row short
|
||||
run_case(2, 6, 8, torch.float32, torch.int32, False, True, False, True, False)
|
||||
# DeepInfer indexer call shape: big-ish k, end, offset, sorted_index, no values
|
||||
run_case(2, 512, 128, torch.float32, torch.int32, False, True, False, True, True)
|
||||
|
||||
print("\nFAILURES:", FAILS if FAILS else "none")
|
||||
sys.exit(1 if FAILS else 0)
|
||||
316
tests/neuron/flash_mla/test_flash_mla.py
Normal file
316
tests/neuron/flash_mla/test_flash_mla.py
Normal file
@ -0,0 +1,316 @@
|
||||
"""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)
|
||||
38
tests/neuron/flash_mla/test_tilelang.py
Normal file
38
tests/neuron/flash_mla/test_tilelang.py
Normal file
@ -0,0 +1,38 @@
|
||||
"""Checks the tilelang stub provides exactly what platform-independent DeepInfer code uses."""
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
FAILS = []
|
||||
|
||||
|
||||
def chk(name, cond):
|
||||
print(f"[{'ok ' if cond else 'FAIL'}] {name}")
|
||||
if not cond:
|
||||
FAILS.append(name)
|
||||
|
||||
|
||||
chk("PassConfigKey.TL_ENABLE_FAST_MATH", isinstance(tilelang.PassConfigKey.TL_ENABLE_FAST_MATH, str))
|
||||
chk("T.float8_e4m3fn", T.float8_e4m3fn == "float8_e4m3fn")
|
||||
chk("dtype map", tilelang.TL_DTYPE_TO_TORCH[T.float8_e4m3fn] is torch.float8_e4m3fn)
|
||||
chk("norm.py dtype strings", T.float32 == "float" and T.bfloat16 == "bfloat16")
|
||||
chk("max_value e4m3", T.max_value(T.float8_e4m3fn) == 448.0)
|
||||
chk("min_value fp32", T.min_value(T.float32) == torch.finfo(torch.float32).min)
|
||||
|
||||
# the pass-config dict pattern of tilelang_ops/compress.py
|
||||
cfg = {}
|
||||
cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = True
|
||||
chk("pass_configs assignment", cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] is True)
|
||||
|
||||
try:
|
||||
T.some_dsl_primitive
|
||||
FAILS.append("unknown attr should raise")
|
||||
except AttributeError as e:
|
||||
chk("unknown T attribute raises AttributeError", "Neuron stub" in str(e))
|
||||
|
||||
print("\nFAILURES:", FAILS if FAILS else "none")
|
||||
sys.exit(1 if FAILS else 0)
|
||||
17
tests/neuron/nie_skip.py
Normal file
17
tests/neuron/nie_skip.py
Normal file
@ -0,0 +1,17 @@
|
||||
"""pytest plugin: report kernels stubbed out in Neuron basic mode as skipped, not failed.
|
||||
|
||||
The Neuron backend exports every asc kernel factory, but those reachable only from DSpark/MTP, ZOS, CP or
|
||||
VL raise ``NotImplementedError("<name>: not supported on Neuron (basic mode)")``. The upstream interface
|
||||
tests still parametrize over them; with ``-p nie_skip`` they show up as skips.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.hookimpl(hookwrapper=True)
|
||||
def pytest_runtest_makereport(item, call):
|
||||
outcome = yield
|
||||
report = outcome.get_result()
|
||||
if call.excinfo is not None and call.excinfo.errisinstance(NotImplementedError) and "basic mode" in str(call.excinfo.value):
|
||||
report.outcome = "skipped"
|
||||
report.longrepr = (str(item.path), item.location[1] or 0, f"Neuron basic mode: {call.excinfo.value}")
|
||||
114
tests/neuron/run_all.sh
Executable file
114
tests/neuron/run_all.sh
Executable file
@ -0,0 +1,114 @@
|
||||
#!/bin/bash
|
||||
# Neuron port test suite. Run after tests/neuron/setup_trn2.sh.
|
||||
#
|
||||
# bash tests/neuron/run_all.sh cpu # host-only: kernel unit tests + end-to-end vs golden (~1 h)
|
||||
# CORES=0-7 bash tests/neuron/run_all.sh device # NeuronCore tests on the given logical cores
|
||||
# CORES=0-7 bash tests/neuron/run_all.sh all
|
||||
# bash tests/neuron/run_all.sh cpu deep_gemm flash_mla # only selected stages
|
||||
#
|
||||
# Device stages use the first core of CORES for single-core runs and CORES for multi-rank deep_ep
|
||||
# (world sizes 2/4/8 limited by the core count). They refuse to start if those cores are in use.
|
||||
# Logs: $WORK/test_logs/<timestamp>/<stage>.log; a PASS/FAIL table is printed at the end.
|
||||
set -uo pipefail
|
||||
|
||||
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
WORK="${WORK:-$HOME/di}"
|
||||
MODE="${1:-cpu}"; shift || true
|
||||
ONLY=("$@")
|
||||
CORES="${CORES:-0-7}"
|
||||
LOGDIR="$WORK/test_logs/$(date +%Y%m%d-%H%M%S)"
|
||||
mkdir -p "$LOGDIR"
|
||||
# shellcheck disable=SC1091
|
||||
source "$WORK/venv/bin/activate"
|
||||
cd "$REPO" || exit 1
|
||||
VENDOR="$REPO/deepinfer/neuron_ops/vendor"
|
||||
T="$REPO/tests/neuron"
|
||||
export DEEPINFER_ACCELERATOR=neuron
|
||||
declare -a SUMMARY=()
|
||||
|
||||
first_core() { echo "${CORES%%[-,]*}"; }
|
||||
core_count() { python3 - "$CORES" <<'EOF'
|
||||
import sys
|
||||
n = 0
|
||||
for part in sys.argv[1].split(","):
|
||||
a, _, b = part.partition("-")
|
||||
n += (int(b) - int(a) + 1) if b else 1
|
||||
print(n)
|
||||
EOF
|
||||
}
|
||||
selected() { [ ${#ONLY[@]} -eq 0 ] && return 0; for s in "${ONLY[@]}"; do [ "$s" = "$1" ] && return 0; done; return 1; }
|
||||
|
||||
run() { # run <stage> <command...>
|
||||
local stage="$1"; shift
|
||||
selected "$stage" || return 0
|
||||
echo "=== $stage"
|
||||
local t0=$SECONDS
|
||||
if "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi
|
||||
SUMMARY+=("$(printf '%-6s %-28s %5ss %s' "$r" "$stage" $((SECONDS - t0)) "$LOGDIR/$stage.log")")
|
||||
echo " $r ($((SECONDS - t0))s)"
|
||||
[ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/ | /'
|
||||
}
|
||||
|
||||
cores_free() {
|
||||
local busy
|
||||
busy=$(neuron-ls 2>/dev/null | grep -cE "python|vLLM|Worker|torchrun")
|
||||
[ "$busy" = "0" ] || { echo "NeuronCores busy ($busy processes); not starting device stages"; return 1; }
|
||||
}
|
||||
|
||||
cpu_stages() {
|
||||
export DEEPINFER_NEURON_DEVICE=cpu
|
||||
run deep_gemm env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_deep_gemm_neuron.py"
|
||||
run deep_gemm_tk env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_interop_tile_kernels.py"
|
||||
run flash_mla bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_flash_mla.py"
|
||||
run deep_select bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_deep_select.py"
|
||||
run tilelang_stub bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_tilelang.py"
|
||||
run tile_kernels python -m pytest -q -p no:cacheprovider "$T/tile_kernels/test_tile_kernels.py"
|
||||
for ws in 1 2 4; do
|
||||
run "deep_ep_ws$ws" python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29531 + ws))
|
||||
run "engram_ws$ws" python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29541 + ws))
|
||||
run "bucket_ws$ws" python "$T/deep_ep/test_bucket.py" --num-processes $ws --port $((29551 + ws))
|
||||
done
|
||||
run sampling_kat env PYTHONPATH="$REPO" python "$T/sampling/kat.py"
|
||||
run sampling_check env PYTHONPATH="$REPO" python "$T/sampling/check2.py"
|
||||
# Upstream interface tests of tilelang_ops (compare the Neuron backend against the *_ref implementations).
|
||||
for m in norm store_kv_cache load_paged_kv_cache fused_transform compress state_cache get_kv_index kv_index_cat \
|
||||
swa_index candidate_blocks flash_mla_utils gather prepare_v4_decode_metadata random_sampling fused_sample \
|
||||
sampler sum_smaller_probs structure_output; do
|
||||
run "tl_$m" env PYTHONPATH="$T" python -m pytest -q -p no:cacheprovider --noconftest -p nie_skip \
|
||||
"deepinfer/model_executor/layers/tilelang_ops/$m.py" -k "not benchmark and not 129280-4096"
|
||||
done
|
||||
run e2e_cpu python examples/neuron/run_deepinfer_tiny.py --device cpu --threads 32 --num-kv-blocks 32 \
|
||||
--out "$WORK/e2e_cpu" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
|
||||
}
|
||||
|
||||
device_stages() {
|
||||
cores_free || return 0
|
||||
local c0 n
|
||||
c0="$(first_core)"; n="$(core_count)"
|
||||
export DEEPINFER_NEURON_DEVICE=neuron
|
||||
run device_primitives env NEURON_RT_VISIBLE_CORES="$c0" python "$T/device/test_device_primitives.py"
|
||||
for ws in 2 4 8; do
|
||||
[ "$ws" -le "$n" ] || continue
|
||||
run "deep_ep_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
|
||||
python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29631 + ws))
|
||||
run "engram_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
|
||||
python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29641 + ws))
|
||||
done
|
||||
run e2e_neuron env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29733 \
|
||||
python examples/neuron/run_deepinfer_tiny.py --device neuron --num-kv-blocks 32 \
|
||||
--out "$WORK/e2e_neuron" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
|
||||
run device_diff env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29734 \
|
||||
python examples/neuron/device_diff.py --prompts short --num-kv-blocks 32 \
|
||||
--out "$WORK/device_diff" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt"
|
||||
}
|
||||
|
||||
case "$MODE" in
|
||||
cpu) cpu_stages ;;
|
||||
device) device_stages ;;
|
||||
all) cpu_stages; device_stages ;;
|
||||
*) echo "usage: $0 cpu|device|all [stage...]"; exit 2 ;;
|
||||
esac
|
||||
|
||||
echo; echo "===== summary ($LOGDIR)"
|
||||
printf '%s\n' "${SUMMARY[@]}"
|
||||
! printf '%s\n' "${SUMMARY[@]}" | grep -q '^FAIL'
|
||||
36
tests/neuron/sampling/check2.py
Normal file
36
tests/neuron/sampling/check2.py
Normal file
@ -0,0 +1,36 @@
|
||||
import os
|
||||
import torch
|
||||
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import (
|
||||
two_level_sample, quantize_weights, acc_slots_for, rng_first_two, rng_sequence)
|
||||
from deepinfer.model_executor.layers.tilelang_ops.random_sampling import random_sampling
|
||||
from deepinfer.model_executor.layers.tilelang_ops.sampler import prepare_verification_metadata, prepare_verification_metadata_ref
|
||||
|
||||
for V in [8, 100, 1024, 4032, 8192, 129280]:
|
||||
Q = 64
|
||||
p = torch.rand(Q, V) ** 8; p[:, ::3] = 0; p = p / p.sum(1, keepdim=True)
|
||||
seeds = torch.randint(-2**62, 2**62, (Q,), dtype=torch.int64)
|
||||
pos = torch.randint(-5, 10**9, (Q,), dtype=torch.int32)
|
||||
out = two_level_sample(p, seeds, pos, 1, V)
|
||||
# independent: inverse CDF over columns ordered by (col % AS, col // AS)
|
||||
AS = acc_slots_for(V)
|
||||
fx = quantize_weights(p).to(torch.int64)
|
||||
cols = torch.arange(V)
|
||||
order = torch.argsort((cols % AS) * (V + 1) + cols // AS)
|
||||
cdf = fx[:, order].cumsum(1)
|
||||
total = cdf[:, -1]
|
||||
lo, hi = rng_first_two(seeds, rng_sequence(pos, 1))
|
||||
r = torch.tensor([((int(h) << 32) | int(l)) % int(t) for h, l, t in zip(hi, lo, total)])
|
||||
exp = order[(cdf <= r[:, None]).sum(1)]
|
||||
assert torch.equal(out.long(), exp), V
|
||||
assert (p.gather(1, out.long()[:, None]) > 0).all()
|
||||
print("two-level == slot-ordered inverse CDF: OK")
|
||||
# strided probs rows
|
||||
p = torch.softmax(torch.randn(16, 2048), -1); big = torch.zeros(16, 4096); big[:, :2048] = p
|
||||
a = random_sampling(big[:, :2048], torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
|
||||
b = random_sampling(p.contiguous(), torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
|
||||
assert torch.equal(a, b); print("strided rows OK")
|
||||
# nextn=0 metadata path (num_steps=0) and seqlen=1 rows
|
||||
seqlen = torch.ones(37, dtype=torch.int32); ftp = torch.arange(37, dtype=torch.int32)
|
||||
for got, ref in zip(prepare_verification_metadata(ftp, seqlen, 37, 0), prepare_verification_metadata_ref(ftp, seqlen, 37)):
|
||||
assert torch.equal(got, ref)
|
||||
print("prepare_verification_metadata num_steps=0 OK")
|
||||
11
tests/neuron/sampling/kat.py
Normal file
11
tests/neuron/sampling/kat.py
Normal file
@ -0,0 +1,11 @@
|
||||
import torch
|
||||
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import philox4x32_10
|
||||
def t(v): return torch.tensor([v], dtype=torch.int64)
|
||||
for ctr, key, exp in [((0,0,0,0),(0,0),(0x6627e8d5,0xe169c58d,0xbc57ac4c,0x9b00dbd8)),
|
||||
((0xffffffff,)*4,(0xffffffff,)*2,(0x408f276d,0x41c83b0e,0xa20bc7c6,0x6d5451fd)),
|
||||
((0x243f6a88,0x85a308d3,0x13198a2e,0x03707344),(0xa4093822,0x299f31d0),(0xd16cfe09,0x94fdcceb,0x5001e420,0x24126ea1))]:
|
||||
out = philox4x32_10(*map(t, ctr), *map(t, key))
|
||||
got = tuple(int(o) for o in out)
|
||||
print([hex(g) for g in got], got == exp)
|
||||
from deepinfer.model_executor.layers.tilelang_ops import ops
|
||||
print(ops.random_sampling_tl, ops.fused_sample_tl, ops.map_lookup_tl)
|
||||
69
tests/neuron/setup_trn2.sh
Executable file
69
tests/neuron/setup_trn2.sh
Executable file
@ -0,0 +1,69 @@
|
||||
#!/bin/bash
|
||||
# One-time setup of a fresh Trn2 (Neuron DLAMI, SDK 2.33+) for the DeepInfer Neuron port.
|
||||
#
|
||||
# git clone -b neuron https://github.com/whn09/DeepInfer.git && cd DeepInfer
|
||||
# bash tests/neuron/setup_trn2.sh # WORK=~/di by default
|
||||
#
|
||||
# Creates $WORK/venv (layered on the DLAMI's read-only torch-neuronx venv), builds the C++ extension,
|
||||
# downloads the official reference implementation + tokenizer, generates the tiny checkpoint and the
|
||||
# CPU golden. Idempotent: finished steps are skipped. Host CPU only -- never touches a NeuronCore.
|
||||
set -euo pipefail
|
||||
|
||||
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
WORK="${WORK:-$HOME/di}"
|
||||
OPT_VENV="${OPT_VENV:-/opt/aws_neuronx_venv_torch_neuronx_2_14_0}"
|
||||
HF_REPO="deepseek-ai/DeepSeek-V4.1-Flash-1017"
|
||||
mkdir -p "$WORK"
|
||||
log() { echo "[setup $(date +%H:%M:%S)] $*"; }
|
||||
|
||||
# 1. venv layered on the read-only DLAMI torch-neuronx venv
|
||||
if [ ! -x "$WORK/venv/bin/python" ]; then
|
||||
log "venv -> $WORK/venv (layered on $OPT_VENV)"
|
||||
PY="$(readlink -f "$OPT_VENV/bin/python")"
|
||||
"$PY" -m venv "$WORK/venv"
|
||||
SP="$("$WORK/venv/bin/python" -c 'import site; print(site.getsitepackages()[0])')"
|
||||
echo "import site; site.addsitedir('$OPT_VENV/lib/python$("$PY" -c 'import sys; print(f"{sys.version_info[0]}.{sys.version_info[1]}")')/site-packages')" > "$SP/zz_opt_neuron.pth"
|
||||
# neuronx-cc and the neuron tools live in the DLAMI venv
|
||||
echo "export PATH=\"\$PATH:$OPT_VENV/bin\"" >> "$WORK/venv/bin/activate"
|
||||
fi
|
||||
# shellcheck disable=SC1091
|
||||
source "$WORK/venv/bin/activate"
|
||||
python -c "import torch, torch_neuronx; print('torch', torch.__version__, 'torch-neuronx', torch_neuronx.__version__)"
|
||||
which neuronx-cc >/dev/null || { echo "neuronx-cc not on PATH"; exit 1; }
|
||||
|
||||
# 2. Python deps: requirements.txt minus torch / tilelang / tile-kernels (pure-torch shims are vendored)
|
||||
log "pip deps"
|
||||
pip install -q einops "huggingface_hub<1" ninja orjson pillow posix_ipc pybind11 pyzmq requests scipy sympy tabulate \
|
||||
tokenizers tqdm "transformers==4.50.0" xgrammar pytest psutil safetensors
|
||||
|
||||
# 3. C++ extension (host-only, -DWITH_NEURON)
|
||||
if ! ls "$REPO"/deepinfer/_so/neuron-*/_C*.so >/dev/null 2>&1; then
|
||||
log "build C++ extension"
|
||||
(cd "$REPO" && DEEPINFER_ACCELERATOR=neuron MAX_JOBS=16 bash ./scripts/build.sh build > "$WORK/build.log" 2>&1) \
|
||||
|| { tail -30 "$WORK/build.log"; exit 1; }
|
||||
fi
|
||||
|
||||
# 4. official reference implementation + tokenizer + real config (small files only, no weights)
|
||||
REF="$WORK/ref_dsv41"
|
||||
if [ ! -f "$REF/inference/model.py" ]; then
|
||||
log "download reference implementation from $HF_REPO"
|
||||
mkdir -p "$REF/inference"
|
||||
for f in inference/model.py inference/kernel.py inference/convert.py inference/engram.py inference/generate.py \
|
||||
inference/vision.py inference/image_processor.py inference/config.json tokenizer.json tokenizer_config.json config.json; do
|
||||
curl -fsSL "https://huggingface.co/$HF_REPO/resolve/main/$f" -o "$REF/$f"
|
||||
done
|
||||
fi
|
||||
|
||||
# 5. tiny checkpoint (~1.6 GB) and the CPU golden from the official reference
|
||||
if [ ! -f "$WORK/tiny_ckpt/config.json" ]; then
|
||||
log "tiny checkpoint -> $WORK/tiny_ckpt"
|
||||
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python "$REPO/examples/neuron/make_tiny_checkpoint.py" \
|
||||
--out "$WORK/tiny_ckpt" --tokenizer-dir "$REF" --real-config "$REF/config.json"
|
||||
fi
|
||||
if [ ! -f "$WORK/golden/long.pt" ]; then
|
||||
log "golden (official reference on CPU) -> $WORK/golden"
|
||||
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I "$REPO/examples/neuron/reference_golden.py" \
|
||||
--hf-ckpt "$WORK/tiny_ckpt" --ref-dir "$REF/inference" --converted "$WORK/tiny_ckpt_ref_mp1" \
|
||||
--out "$WORK/golden" --threads 32
|
||||
fi
|
||||
log "done. Next: bash tests/neuron/run_all.sh"
|
||||
20
tests/neuron/tile_kernels/integ.py
Normal file
20
tests/neuron/tile_kernels/integ.py
Normal file
@ -0,0 +1,20 @@
|
||||
import torch
|
||||
import deepinfer # noqa: sets vendor path
|
||||
from deepinfer.model_executor.models.v4_modules import mhc_ops
|
||||
import tile_kernels.torch as ref
|
||||
torch.manual_seed(0)
|
||||
t, h, M = 6, 256, 4
|
||||
m3 = M * (M + 2)
|
||||
res = torch.randn(1, t, M, h, dtype=torch.bfloat16)
|
||||
fn = torch.randn(m3, M * h) * 0.02
|
||||
w = torch.rand(M * h) + 0.5
|
||||
print('n_splits', mhc_ops.default_n_splits())
|
||||
for ns in (1, 16):
|
||||
try:
|
||||
out = mhc_ops.mhc_pre_norm_fn(res, fn, w, 1e-6, n_splits=ns)
|
||||
r = ref.mhc_pre_norm_fn_ref(res, fn, w, 1e-6)
|
||||
print(ns, 'max diff', (out - r).abs().max().item())
|
||||
except Exception as e:
|
||||
print(ns, 'ERR', type(e).__name__, str(e)[:300])
|
||||
x = torch.randn(t, 2, 128, dtype=torch.bfloat16)
|
||||
o = mhc_ops.expand_to_mhc(x, 4); print('expand ok', o.shape)
|
||||
15
tests/neuron/tile_kernels/plugin/conftest.py
Normal file
15
tests/neuron/tile_kernels/plugin/conftest.py
Normal file
@ -0,0 +1,15 @@
|
||||
import os, sys
|
||||
sys.path.insert(0, os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "..")), "deepinfer", "neuron_ops", "vendor"))
|
||||
_D = os.path.dirname(os.path.abspath(__file__))
|
||||
os.environ.setdefault("TK_TEST_MERGE_MODE", "overwrite")
|
||||
os.environ.setdefault("TK_BENCHMARK_TESTS_DIR", _D)
|
||||
os.environ.setdefault("TK_BENCHMARK_BASELINES_DIR", os.path.join(_D, "bb"))
|
||||
os.environ.setdefault("TK_BENCHMARK_FAIL_MODE", "none")
|
||||
os.environ.setdefault("TK_GPU_MEM_PROFILES_DIR", os.path.join(_D, "gm"))
|
||||
pytest_plugins = [
|
||||
"tile_kernels.testing.pytest.benchmark",
|
||||
"tile_kernels.testing.pytest.gpu_mem",
|
||||
"tile_kernels.testing.pytest.random",
|
||||
"tile_kernels.testing.pytest.tilelang_warning",
|
||||
"tile_kernels.testing.pytest.xdist_failfast",
|
||||
]
|
||||
6
tests/neuron/tile_kernels/plugin/test_plugin.py
Normal file
6
tests/neuron/tile_kernels/plugin/test_plugin.py
Normal file
@ -0,0 +1,6 @@
|
||||
import torch
|
||||
|
||||
|
||||
def test_timer(benchmark_timer, seed):
|
||||
us = benchmark_timer(lambda: torch.ones(8).sum())
|
||||
assert us > 0 and isinstance(seed, int)
|
||||
557
tests/neuron/tile_kernels/test_tile_kernels.py
Normal file
557
tests/neuron/tile_kernels/test_tile_kernels.py
Normal file
@ -0,0 +1,557 @@
|
||||
"""CPU correctness tests: Neuron pure-torch tile_kernels vs the upstream tile_kernels.torch references."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
VENDOR = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "neuron_ops", "vendor")
|
||||
sys.path.insert(0, VENDOR)
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import tile_kernels
|
||||
from tile_kernels import quant
|
||||
from tile_kernels import torch as ref
|
||||
from tile_kernels.quant.common import get_packed_ue8m0_torch_dtype
|
||||
|
||||
torch.set_num_threads(4)
|
||||
|
||||
|
||||
def bits(t):
|
||||
if t.dtype in (torch.float8_e4m3fn,):
|
||||
return t.view(torch.uint8)
|
||||
return t
|
||||
|
||||
|
||||
def assert_same(a, b, check_stride=True):
|
||||
assert a.dtype == b.dtype, (a.dtype, b.dtype)
|
||||
assert a.shape == b.shape, (a.shape, b.shape)
|
||||
if check_stride and a.numel() > 1:
|
||||
assert a.stride() == b.stride(), (a.stride(), b.stride())
|
||||
if a.dtype.is_floating_point and a.dtype != torch.float8_e4m3fn:
|
||||
assert torch.equal(torch.nan_to_num(a), torch.nan_to_num(b)), (a - b).abs().max()
|
||||
else:
|
||||
assert torch.equal(bits(a), bits(b))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# config / surface
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_surface():
|
||||
assert tile_kernels.config.is_ascend() is False
|
||||
assert tile_kernels.get_num_vec_cores() == 2 * tile_kernels.get_num_sms()
|
||||
assert tile_kernels.get_max_ub_per_vector_core() > 0 and tile_kernels.config.get_max_ub_per_vector_core() > 0
|
||||
prev = tile_kernels.get_num_sms()
|
||||
tile_kernels.set_num_sms(3)
|
||||
assert tile_kernels.get_num_sms() == 3 and tile_kernels.get_num_vec_cores() == 6
|
||||
tile_kernels.set_num_sms(prev)
|
||||
tile_kernels.set_pdl(True)
|
||||
tile_kernels.set_pdl(False)
|
||||
tile_kernels.use_deterministic_algorithms(True)
|
||||
tile_kernels.set_token_alignment(64)
|
||||
assert tile_kernels.get_token_alignment() == 64
|
||||
tile_kernels.set_token_alignment(128)
|
||||
assert tile_kernels.utils.align(5, 4) == 8
|
||||
assert tile_kernels.utils.dtype_to_str(torch.bfloat16) == 'bf16'
|
||||
assert get_packed_ue8m0_torch_dtype() == torch.int16
|
||||
from tile_kernels.config import is_ascend # noqa
|
||||
from tile_kernels.mhc.expand_kernel import expand_to_mhc_fwd # noqa
|
||||
from tile_kernels.mhc.head_compute_mix_kernel import mhc_head_compute_mix_fwd # noqa
|
||||
from tile_kernels.mhc.norm_fn_kernel import mhc_fn_normw_merge_fwd, mhc_reduce_partials_and_rmsnorm_fwd # noqa
|
||||
from tile_kernels.mhc.post_kernel import mhc_post_fwd # noqa
|
||||
from tile_kernels.mhc.pre_apply_mix_kernel import mhc_pre_apply_mix_fwd # noqa
|
||||
from tile_kernels.mhc.pre_split_mixes_kernel import mhc_pre_split_mixes_fwd # noqa
|
||||
from tile_kernels.mhc.sinkhorn_kernel import mhc_sinkhorn_fwd # noqa
|
||||
from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd # noqa
|
||||
from tile_kernels.moe import moe_topk_gate # noqa
|
||||
from tile_kernels.quant import swiglu_forward_and_per_token_cast, swiglu_forward, per_token_cast, batched_transpose_weight_sf # noqa
|
||||
from tile_kernels.transform import apply_rotary # noqa
|
||||
from tile_kernels.torch import cast, cast_back # noqa
|
||||
from tile_kernels import get_max_ub_per_vector_core, get_num_vec_cores # noqa
|
||||
for name in ('engram_hash', 'engram_gate_fwd', 'fused_weight'):
|
||||
assert callable(getattr(tile_kernels.engram, name))
|
||||
assert callable(tile_kernels.torch.engram.make_offsets)
|
||||
for name in ('per_token_cast', 'per_token_cast_back', 'per_token_cast_and_cast_back'):
|
||||
assert callable(getattr(quant, name))
|
||||
import importlib
|
||||
|
||||
for mod in ('benchmark', 'gpu_mem', 'random', 'tilelang_warning', 'xdist_failfast'):
|
||||
importlib.import_module(f'tile_kernels.testing.pytest.{mod}')
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# per_token_cast
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
LAYOUTS = [
|
||||
# round_sf, packed, col_major, e4m3_sf
|
||||
(False, False, False, False),
|
||||
(True, False, False, False),
|
||||
(True, True, False, False),
|
||||
(True, True, True, False),
|
||||
(True, False, True, False),
|
||||
(False, False, True, False),
|
||||
(False, False, False, True),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('layout', LAYOUTS)
|
||||
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
||||
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
|
||||
@pytest.mark.parametrize('m,hidden', [(1, 128), (7, 256), (33, 512), (64, 192)])
|
||||
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
||||
def test_per_token_cast(layout, fmt, npc, m, hidden, dtype):
|
||||
round_sf, packed, col, e4m3_sf = layout
|
||||
if hidden % npc:
|
||||
pytest.skip('hidden % npc')
|
||||
torch.manual_seed(m * 1000 + hidden + npc)
|
||||
x = (torch.randn(m, hidden) * torch.logspace(-3, 3, hidden)).to(dtype)
|
||||
x[0, :npc] = 0 # all-zero group exercises the amax clamp
|
||||
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
|
||||
q, sf = quant.per_token_cast(x, fmt, npc, **kw)
|
||||
rq, rsf = ref.cast(x, fmt, (1, npc), **kw)
|
||||
assert_same(q, rq)
|
||||
assert_same(sf, rsf)
|
||||
# sf_only / precomputed sf variants
|
||||
sf2 = quant.per_token_cast_with_sf_only(x, fmt, npc, use_tma_aligned_col_major_sf=col, round_sf=round_sf, use_packed_ue8m0=packed, use_e4m3_sf=e4m3_sf)
|
||||
assert_same(sf2, rsf)
|
||||
if not packed and not col:
|
||||
q3 = quant.per_token_cast_with_precomputed_sf(x, fmt, npc, rsf)
|
||||
rq3 = ref.cast(x, fmt, (1, npc), sf=rsf.float())
|
||||
assert_same(q3, rq3)
|
||||
# out= in place
|
||||
if not col:
|
||||
out_x = torch.empty_like(q)
|
||||
out_sf = torch.empty_like(sf)
|
||||
rx, rs = quant.per_token_cast(x, fmt, npc, out=(out_x, out_sf), **kw)
|
||||
assert rx.data_ptr() == out_x.data_ptr() and rs.data_ptr() == out_sf.data_ptr()
|
||||
assert_same(out_x, rq)
|
||||
assert_same(out_sf, rsf)
|
||||
|
||||
|
||||
def test_per_token_cast_ascend_layout_bytes():
|
||||
"""Packed col-major sf: byte (k % 2) of storage element [k // 2, m] is exponent of group k (Ascend get_sf_index)."""
|
||||
torch.manual_seed(0)
|
||||
m, hidden, npc = 5, 256, 32
|
||||
x = torch.randn(m, hidden, dtype=torch.bfloat16) * 10
|
||||
_, sf = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
||||
assert sf.dtype == torch.int16 and sf.shape == (m, hidden // npc // 2) and sf.stride() == (1, m)
|
||||
storage = sf.T.contiguous().view(torch.uint8) # (P, 2m) bytes
|
||||
_, sf_rm = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=False)
|
||||
exp = (sf_rm.view(torch.int32) >> 23)
|
||||
for mi in range(m):
|
||||
for k in range(hidden // npc):
|
||||
assert storage[k // 2, 2 * mi + k % 2].item() == exp[mi, k].item()
|
||||
# row-major packed: byte 2p + j of row m = group 2p + j
|
||||
_, sfp = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True)
|
||||
assert sfp.shape == (m, hidden // npc // 2) and sfp.is_contiguous()
|
||||
assert torch.equal(sfp.view(torch.uint8).to(torch.int32), exp)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('in_layout', ['fp32_128', 'packed_col_32', 'packed_row_32', 'fp32_1x32'])
|
||||
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
||||
def test_per_token_cast_with_input_sf(in_layout, fmt):
|
||||
torch.manual_seed(1)
|
||||
m, hidden = 256, 256
|
||||
w = torch.randn(m, hidden) * 3
|
||||
if in_layout == 'fp32_128':
|
||||
xq, xsf = ref.cast(w, 'e4m3', (128, 128), round_sf=True)
|
||||
blk = (128, 128)
|
||||
elif in_layout == 'fp32_1x32':
|
||||
xq, xsf = ref.cast(w, 'e4m3', (1, 32))
|
||||
blk = (1, 32)
|
||||
elif in_layout == 'packed_col_32':
|
||||
xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
||||
blk = (1, 32)
|
||||
else:
|
||||
xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True)
|
||||
blk = (1, 32)
|
||||
kw = dict(round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True)
|
||||
q, sf = quant.per_token_cast((xq, xsf), fmt, 32, x_block_size=blk, **kw)
|
||||
rq, rsf = ref.cast((xq, xsf), fmt, (1, 32), x_block_size=blk, **kw)
|
||||
assert_same(q, rq)
|
||||
assert_same(sf, rsf)
|
||||
|
||||
|
||||
def test_per_token_cast_empty():
|
||||
x = torch.empty(0, 128, dtype=torch.bfloat16)
|
||||
q, sf = quant.per_token_cast(x, 'e4m3', 32, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
||||
assert q.shape == (0, 128) and sf.shape == (0, 2) and sf.dtype == torch.int16
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# cast_back / cast_and_cast_back / per_block_cast
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize('layout', LAYOUTS)
|
||||
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
||||
@pytest.mark.parametrize('npc', [16, 32, 128])
|
||||
@pytest.mark.parametrize('out_fmt', ['bf16', 'fp32'])
|
||||
def test_per_token_cast_back(layout, fmt, npc, out_fmt):
|
||||
round_sf, packed, col, e4m3_sf = layout
|
||||
torch.manual_seed(npc)
|
||||
x = torch.randn(19, 256, dtype=torch.bfloat16) * 5
|
||||
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
|
||||
q, sf = ref.cast(x, fmt, (1, npc), **kw)
|
||||
out = quant.per_token_cast_back((q, sf), out_fmt, npc)
|
||||
rout = ref.cast_back((q, sf), out_fmt, (1, npc))
|
||||
assert_same(out, rout)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
|
||||
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
|
||||
@pytest.mark.parametrize('sf_mode', ['plain', 'round', 'e4m3'])
|
||||
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
||||
def test_per_token_cast_and_cast_back(fmt, npc, sf_mode, dtype):
|
||||
torch.manual_seed(npc)
|
||||
base = torch.randn(23, 384, dtype=dtype) * 7
|
||||
big = torch.zeros(23, 512, dtype=dtype)
|
||||
big[:, :384] = base
|
||||
x = big[:, :384] # strided rows
|
||||
round_sf, e4m3_sf = sf_mode == 'round', sf_mode == 'e4m3'
|
||||
quant.per_token_cast_and_cast_back(x, fmt, npc, round_sf, e4m3_sf)
|
||||
q, sf = ref.cast(base, fmt, (1, npc), round_sf=round_sf, use_e4m3_sf=e4m3_sf)
|
||||
expected = ref.cast_back((q, sf), 'fp32', (1, npc)).to(dtype)
|
||||
assert torch.equal(x, expected)
|
||||
assert torch.equal(big[:, 384:], torch.zeros(23, 128, dtype=dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize('block', [(128, 128), (32, 32), (1, 128)])
|
||||
@pytest.mark.parametrize('layout', LAYOUTS[:6])
|
||||
def test_per_block_cast(block, layout):
|
||||
round_sf, packed, col, _ = layout
|
||||
torch.manual_seed(2)
|
||||
x = torch.randn(256, 384, dtype=torch.bfloat16)
|
||||
kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col)
|
||||
q, sf = quant.per_block_cast(x, 'e4m3', block, **kw)
|
||||
rq, rsf = ref.cast(x, 'e4m3', block, **kw)
|
||||
assert_same(q, rq)
|
||||
assert_same(sf, rsf)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# swiglu
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize('with_weights', [False, True])
|
||||
@pytest.mark.parametrize('clamp', [None, 3.0])
|
||||
@pytest.mark.parametrize('with_psum', [False, True])
|
||||
def test_swiglu_forward(dtype, with_weights, clamp, with_psum):
|
||||
torch.manual_seed(3)
|
||||
tile_kernels.set_token_alignment(8)
|
||||
n, hidden = 64, 256
|
||||
x = (torch.randn(n, hidden * 2) * 3).to(dtype)
|
||||
w = torch.rand(n) if with_weights else None
|
||||
rsf = 2.5 if with_weights else None
|
||||
psum = torch.tensor([5, 8, 8, 20, 37, 44], dtype=torch.int32) if with_psum else None # aligned-prefix layout
|
||||
for fmt in ('fp32', 'bf16'):
|
||||
cnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
|
||||
rcnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
|
||||
out = quant.swiglu_forward(x, fmt, psum, w, rsf, clamp, cnt)
|
||||
rout = ref.swiglu_forward(x, psum, w, rsf, clamp, rcnt)
|
||||
assert torch.equal(out, rout.to(out.dtype))
|
||||
if cnt is not None:
|
||||
assert torch.equal(cnt, rcnt), (cnt, rcnt)
|
||||
# fused fp8 cast (deepinfer get_fp8_cast_params)
|
||||
kw = dict(round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
||||
for npc in (32, 128):
|
||||
q, sf = quant.swiglu_forward_and_per_token_cast(x, 'e4m3', npc, psum, w, rsf, clamp_value=clamp, **kw)
|
||||
rq, rsf_ = ref.cast(ref.swiglu_forward(x, psum, w, rsf, clamp), 'e4m3', (1, npc), **kw)
|
||||
assert_same(q, rq)
|
||||
assert_same(sf, rsf_)
|
||||
tile_kernels.set_token_alignment(128)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# batched_transpose_weight_sf
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize('block', [32, 128])
|
||||
def test_batched_transpose_weight_sf(block):
|
||||
torch.manual_seed(4)
|
||||
|
||||
def make(*shape):
|
||||
e = torch.randint(110, 140, shape, dtype=torch.int32)
|
||||
return (e << 23).view(torch.float32)
|
||||
|
||||
sfs = [make(384 // block, 640 // block), make(3, 256 // block, 896 // block), make(128 // block, 128 // block)]
|
||||
flags = [False, True, True]
|
||||
outs = quant.batched_transpose_weight_sf(sfs, block, flags, outs=[None] * 3)
|
||||
refs = ref.batched_transpose_weight_sf(sfs, block, flags)
|
||||
for (a0, a1), (r0, r1) in zip(outs, refs):
|
||||
assert_same(a0, r0)
|
||||
assert_same(a1, r1)
|
||||
# in-place rewrite with new values
|
||||
sfs2 = [make(*s.shape) for s in sfs]
|
||||
outs2 = quant.batched_transpose_weight_sf(sfs2, block, flags, outs=outs)
|
||||
refs2 = ref.batched_transpose_weight_sf(sfs2, block, flags)
|
||||
for (a0, a1), (o0, o1), (r0, r1) in zip(outs2, outs, refs2):
|
||||
assert a0.data_ptr() == o0.data_ptr() and a1.data_ptr() == o1.data_ptr()
|
||||
assert_same(a0, r0)
|
||||
assert_same(a1, r1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# rope
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize('rot_dim', [64, 128])
|
||||
@pytest.mark.parametrize('ndim', [3, 4])
|
||||
@pytest.mark.parametrize('interleaved', [False, True])
|
||||
@pytest.mark.parametrize('conjugate', [False, True])
|
||||
@pytest.mark.parametrize('use_positions', [False, True])
|
||||
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
|
||||
def test_apply_rotary(rot_dim, ndim, interleaved, conjugate, use_positions, dtype):
|
||||
torch.manual_seed(5)
|
||||
seqlen_ro, s, h, hk, b = 64, 9, 4, 2, 2
|
||||
cache = torch.randn(seqlen_ro, rot_dim)
|
||||
cos, sin = cache[:, : rot_dim // 2], cache[:, rot_dim // 2 :]
|
||||
lead = (s,) if ndim == 3 else (b, s)
|
||||
full_q = torch.randn(*lead, h, rot_dim + 32).to(dtype)
|
||||
q = full_q[..., 32:] # tail partial rope, strided
|
||||
k = torch.randn(*lead, hk, rot_dim).to(dtype)
|
||||
pos = torch.randint(0, seqlen_ro - 4, lead, dtype=torch.int32) if use_positions else None
|
||||
off = 3
|
||||
rq = ref.rotary_embedding_ref(q.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
|
||||
rk = ref.rotary_embedding_ref(k.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
|
||||
head = full_q[..., :32].clone()
|
||||
tile_kernels.transform.apply_rotary(q, cache, k, pos, interleaved, conjugate, off)
|
||||
tol = dict(atol=0, rtol=0) if dtype == torch.float32 else dict(atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(q.float(), rq.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
|
||||
torch.testing.assert_close(k.float(), rk.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
|
||||
assert torch.equal(full_q[..., :32], head)
|
||||
# bitwise vs the Ascend op order: out1 = fma(first, cos, second * -sin), out2 = fma(first, sin, second * cos)
|
||||
def fma(a, b, c):
|
||||
return (a.double() * b.double() + c.double()).float()
|
||||
|
||||
q2 = full_q[..., 32:].clone()
|
||||
q0 = torch.randn_like(q2.float()).to(dtype)
|
||||
q2.copy_(q0)
|
||||
tile_kernels.transform.apply_rotary(q2, cache, None, pos, interleaved, conjugate, off)
|
||||
idx = (pos.long() + off) if pos is not None else (torch.arange(s) + off).expand(*lead[:-1], s)
|
||||
c_, s_ = cos[idx].unsqueeze(-2), sin[idx].unsqueeze(-2)
|
||||
s_ = -s_ if conjugate else s_
|
||||
xf = q0.float()
|
||||
f, sc = (xf[..., 0::2], xf[..., 1::2]) if interleaved else (xf[..., : rot_dim // 2], xf[..., rot_dim // 2 :])
|
||||
c_, s_ = c_.expand_as(f), s_.expand_as(f)
|
||||
o1, o2 = fma(f, c_, sc * -s_), fma(f, s_, sc * c_)
|
||||
exp = torch.stack([o1, o2], -1).flatten(-2) if interleaved else torch.cat([o1, o2], -1)
|
||||
assert torch.equal(q2, exp.to(dtype))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# moe
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize('num_shared', [0, 1, 2])
|
||||
@pytest.mark.parametrize('with_map', [False, True])
|
||||
@pytest.mark.parametrize('with_mask', [False, True])
|
||||
@pytest.mark.parametrize('with_image', [False, True])
|
||||
@pytest.mark.parametrize('with_fix', [False, True])
|
||||
def test_moe_topk_gate(num_shared, with_map, with_mask, with_image, with_fix):
|
||||
torch.manual_seed(6)
|
||||
t, e, k = 37, 64, 6
|
||||
logits = torch.randn(t, e) * 2
|
||||
logits[3, 10] = logits[3, 11] # tie
|
||||
bias = torch.randn(e) * 0.1
|
||||
kwargs = dict(num_topk=k, use_shared_as_routed=num_shared > 0, num_shared_experts=num_shared, routed_scaling_factor=2.5, ep_rank=3, scoring_func='sqrtsoftplus', bias=bias)
|
||||
if with_mask:
|
||||
kwargs['mask'] = torch.rand(t) > 0.3
|
||||
if with_image:
|
||||
kwargs['image_bias'] = torch.randn(e) * 0.1
|
||||
kwargs['image_token_mask'] = torch.rand(t) > 0.5
|
||||
if with_map:
|
||||
nl = e + num_shared
|
||||
dup = 3
|
||||
count = torch.randint(1, dup + 1, (nl,), dtype=torch.int32)
|
||||
kwargs['to_physical_map'] = torch.randint(0, 500, (nl, dup), dtype=torch.int32)
|
||||
kwargs['logical_count'] = count
|
||||
unm = unm_ref = None
|
||||
if with_fix:
|
||||
unm = torch.stack([torch.randperm(e)[:k] for _ in range(t)]).to(torch.int64)
|
||||
unm_ref = unm.clone()
|
||||
kwargs['fix_routing_mask'] = torch.rand(t) > 0.5
|
||||
idx, w = tile_kernels.moe.moe_topk_gate(logits, unmapped_topk_idx=unm, **kwargs)
|
||||
ridx, rw = ref.moe_topk_gate_forward(logits, unmapped_topk_idx=unm_ref, **kwargs)
|
||||
assert torch.equal(idx, ridx)
|
||||
torch.testing.assert_close(w, rw, atol=1e-6, rtol=1e-6)
|
||||
if with_fix:
|
||||
assert torch.equal(unm, unm_ref)
|
||||
|
||||
|
||||
def test_moe_topk_gate_out_scores_force_random():
|
||||
torch.manual_seed(7)
|
||||
t, e, k = 16, 32, 4
|
||||
logits = torch.randn(t, e)
|
||||
mask = torch.ones(t, dtype=torch.bool)
|
||||
mask[0] = False
|
||||
fr = torch.zeros(t, dtype=torch.bool)
|
||||
fr[1:4] = True
|
||||
out = (torch.empty(t, k, dtype=torch.int64), torch.empty(t, k))
|
||||
scores = torch.empty(t, e)
|
||||
unm = torch.empty(t, k, dtype=torch.int64)
|
||||
idx, w = tile_kernels.moe.moe_topk_gate(logits, k, False, 0, 1.0, 0, mask=mask, force_random=fr, out=out, scores=scores, unmapped_topk_idx=unm)
|
||||
assert idx.data_ptr() == out[0].data_ptr()
|
||||
assert (idx[0] == -1).all() and (w[0] == 0).all() and (scores[0] == 0).all() and (unm[0] == -1).all()
|
||||
assert (unm[1:4] == -1).all() and (scores[1:4] == 0).all()
|
||||
for r in range(1, 4):
|
||||
assert len(set(idx[r].tolist())) == k and idx[r].min() >= 0 and idx[r].max() < e
|
||||
ridx, rw = ref.moe_topk_gate_forward(logits[4:], k, False, 0, 1.0, 0, 'sqrtsoftplus')
|
||||
assert torch.equal(idx[4:], ridx)
|
||||
torch.testing.assert_close(scores[4:], torch.nn.functional.softplus(logits[4:]).sqrt())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# engram
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_engram_hash():
|
||||
torch.manual_seed(8)
|
||||
t, n, layers, tables = 29, 4, 2, 3
|
||||
ids = torch.randint(0, 100000, (t, n), dtype=torch.int32)
|
||||
mult = torch.randint(1, 2**40, (layers, n), dtype=torch.int64)
|
||||
vs = torch.randint(1000, 50000, (layers, n - 1, tables), dtype=torch.int32)
|
||||
off = tile_kernels.torch.engram.make_offsets(vs)
|
||||
mask = torch.rand(t) > 0.7
|
||||
for m in (None, mask):
|
||||
out = tile_kernels.engram.engram_hash(ids, mult, vs, off, m)
|
||||
rout = ref.engram_hash_ref(ids, mult, vs, off, m)
|
||||
assert out.dtype == torch.int32 and torch.equal(out, rout)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('inplace', [False, True])
|
||||
@pytest.mark.parametrize('with_mask', [False, True])
|
||||
def test_engram_gate(inplace, with_mask):
|
||||
torch.manual_seed(9)
|
||||
t, hc, d = 13, 4, 256
|
||||
hs = torch.randn(t, hc, d, dtype=torch.bfloat16)
|
||||
kv = torch.randn(t, hc + 1, d, dtype=torch.bfloat16)
|
||||
wh = torch.randn(hc, d, dtype=torch.bfloat16)
|
||||
we = torch.randn(hc, d, dtype=torch.bfloat16)
|
||||
mask = (torch.rand(t) > 0.5) if with_mask else None
|
||||
wf = tile_kernels.engram.fused_weight(wh, we)
|
||||
assert wf.dtype == torch.float32 and torch.equal(wf, wh.float() * we.float())
|
||||
r_out, r_dot, r_g, r_rx, r_rk = ref.engram_gate_ref(hs, kv, wh, we, 1e-4, 1e-6, save_for_backward=True, image_token_mask=mask)
|
||||
hs_in = hs.clone()
|
||||
out, dot, g, rx, rk = tile_kernels.engram.engram_gate_fwd(hs_in, kv, wf, 1e-6, 1e-4, save_for_backward=True, image_token_mask=mask, inplace=inplace)
|
||||
if inplace:
|
||||
assert out.data_ptr() == hs_in.data_ptr()
|
||||
sel = ~mask if mask is not None else slice(None)
|
||||
assert (out.float() - r_out.float()).abs().max() <= 0.0625
|
||||
assert torch.equal(out[mask], hs[mask]) if mask is not None else True
|
||||
torch.testing.assert_close(dot[sel], r_dot[sel], atol=1e-3, rtol=1e-4)
|
||||
torch.testing.assert_close(g[sel], r_g[sel], atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(rx[sel], r_rx[sel], atol=1e-6, rtol=1e-5)
|
||||
torch.testing.assert_close(rk[sel], r_rk[sel], atol=1e-6, rtol=1e-5)
|
||||
out2, *rest = tile_kernels.engram.engram_gate_fwd(hs.clone(), kv, wf, 1e-6, 1e-4, save_for_backward=False, image_token_mask=mask)
|
||||
assert all(r is None for r in rest)
|
||||
assert torch.equal(out2, out)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mhc
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MHC = 4
|
||||
|
||||
|
||||
def test_mhc_expand_and_mix():
|
||||
torch.manual_seed(10)
|
||||
n, h = 11, 256
|
||||
x = torch.randn(n, h, dtype=torch.bfloat16)
|
||||
o = torch.empty(n, MHC, h, dtype=torch.bfloat16)
|
||||
tile_kernels.mhc.expand_to_mhc_fwd(x, o)
|
||||
assert torch.equal(o, ref.expand_to_mhc_ref(x, MHC))
|
||||
|
||||
inp = torch.randn(n, MHC)
|
||||
scale, base = torch.randn(1), torch.randn(MHC)
|
||||
out = torch.empty_like(inp)
|
||||
tile_kernels.mhc.mhc_head_compute_mix_fwd(inp, scale, base, out, 1e-6)
|
||||
torch.testing.assert_close(out, ref.mhc_head_compute_mix_ref(inp, scale, base, 1e-6), atol=0, rtol=0)
|
||||
|
||||
m3 = MHC * (MHC + 2)
|
||||
mixes = torch.randn(1, n, m3)
|
||||
scale3, base3 = torch.randn(3), torch.randn(m3)
|
||||
pre, post, comb = torch.empty(n, MHC), torch.empty(n, MHC), torch.empty(n, MHC * MHC)
|
||||
tile_kernels.mhc.mhc_pre_split_mixes_fwd(mixes.view(n, m3), scale3, base3, pre, post, comb, 2.0, 1e-6)
|
||||
rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale3, base3, MHC, 2.0, 1e-6)
|
||||
torch.testing.assert_close(pre, rpre.view(n, MHC), atol=1e-7, rtol=1e-6)
|
||||
torch.testing.assert_close(post, rpost.view(n, MHC), atol=1e-7, rtol=1e-6)
|
||||
torch.testing.assert_close(comb, rcomb.reshape(n, MHC * MHC), atol=1e-7, rtol=1e-6)
|
||||
|
||||
xr = torch.randn(n, MHC, h, dtype=torch.bfloat16)
|
||||
mix = torch.rand(n, MHC)
|
||||
o2 = torch.empty(n, h, dtype=torch.bfloat16)
|
||||
tile_kernels.mhc.mhc_pre_apply_mix_fwd(xr, mix, o2)
|
||||
torch.testing.assert_close(o2.float(), ref.mhc_pre_apply_mix_ref(xr.float(), mix.unsqueeze(-1)).float(), atol=1e-2, rtol=1e-2)
|
||||
|
||||
cm = torch.randn(n, MHC, MHC)
|
||||
cmo = torch.empty_like(cm)
|
||||
tile_kernels.mhc.mhc_sinkhorn_fwd(cm, cmo, 10, 1e-6)
|
||||
torch.testing.assert_close(cmo, ref.sinkhorn_normalize_ref(cm, 10, 1e-6), atol=1e-6, rtol=1e-5)
|
||||
|
||||
|
||||
def test_mhc_post():
|
||||
torch.manual_seed(11)
|
||||
a, b, h = 2, 5, 256
|
||||
x = torch.randn(a, b, h, dtype=torch.bfloat16)
|
||||
res = torch.randn(a, b, MHC, h, dtype=torch.bfloat16)
|
||||
post = torch.rand(a, b, MHC, 1)
|
||||
comb = torch.rand(a, b, MHC, MHC)
|
||||
out = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb)
|
||||
rout = ref.mhc_post_ref(x, res, post, comb)
|
||||
assert out.dtype == torch.bfloat16 and out.shape == res.shape
|
||||
assert (out.float() - rout.float()).abs().max() <= 0.0625
|
||||
out2 = torch.empty_like(res)
|
||||
r = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb, out=out2)
|
||||
assert r.data_ptr() == out2.data_ptr() and torch.equal(out2, out)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('n_splits', [1, 4, 16])
|
||||
def test_mhc_norm_fn(n_splits):
|
||||
torch.manual_seed(12)
|
||||
t, h = 9, 256
|
||||
m3 = MHC * (MHC + 2)
|
||||
residual = torch.randn(1, t, MHC, h, dtype=torch.bfloat16)
|
||||
fn = torch.randn(m3, MHC * h) * 0.02
|
||||
normw = torch.rand(MHC * h) + 0.5
|
||||
merged = torch.empty_like(fn)
|
||||
tile_kernels.mhc.mhc_fn_normw_merge_fwd(fn, normw, merged)
|
||||
assert torch.equal(merged, fn * normw)
|
||||
om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual, merged, n_splits) # (S, T, 1, M3), (S, T, 1)
|
||||
om, sq, out = torch.empty(t, 1, m3), torch.empty(t, 1), torch.empty(t, m3)
|
||||
tile_kernels.mhc.mhc_reduce_partials_and_rmsnorm_fwd(om_s, sq_s, om, sq, out, MHC * h, 1e-6, n_splits)
|
||||
rout = ref.mhc_pre_norm_fn_ref(residual, fn, normw, 1e-6).view(t, m3)
|
||||
torch.testing.assert_close(out, rout, atol=1e-4, rtol=1e-4)
|
||||
torch.testing.assert_close(om, om_s.sum(0))
|
||||
torch.testing.assert_close(sq, sq_s.sum(0))
|
||||
|
||||
|
||||
@pytest.mark.parametrize('n_splits', [1, 16])
|
||||
def test_mhc_pre_big_fuse(n_splits):
|
||||
torch.manual_seed(13)
|
||||
t, h = 7, 256
|
||||
m3 = MHC * (MHC + 2)
|
||||
residual = torch.randn(t, MHC, h, dtype=torch.bfloat16)
|
||||
fn = torch.randn(m3, MHC * h) * 0.02
|
||||
om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual.unsqueeze(0), fn, n_splits)
|
||||
gemm_mul, gemm_sq = om_s.view(n_splits, t, m3), sq_s.view(n_splits, t)
|
||||
scale, base = torch.rand(3), torch.randn(m3) * 0.1
|
||||
post, comb, layer_in = torch.empty(t, MHC), torch.empty(t, MHC * MHC), torch.empty(t, h, dtype=torch.bfloat16)
|
||||
tile_kernels.mhc.mhc_pre_big_fuse_fwd(gemm_mul, gemm_sq, scale, base, residual, post, comb, layer_in, 1e-6, 1e-6, 1e-6, 2.0, 20, n_splits=n_splits)
|
||||
mixes = ref.mhc_pre_norm_fn_ref(residual.unsqueeze(0), fn, None, 1e-6) # (1, T, M3)
|
||||
rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale, base, MHC, 2.0, 1e-6)
|
||||
rcomb = ref.sinkhorn_normalize_ref(rcomb, 20, 1e-6)
|
||||
rli = ref.mhc_pre_apply_mix_ref(residual.unsqueeze(0).float(), rpre)
|
||||
torch.testing.assert_close(post, rpost.view(t, MHC), atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(comb, rcomb.reshape(t, MHC * MHC), atol=1e-5, rtol=1e-4)
|
||||
torch.testing.assert_close(layer_in.float(), rli.view(t, h).float(), atol=2e-2, rtol=1e-2)
|
||||
Loading…
x
Reference in New Issue
Block a user