mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
"""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)
|