mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
c6a65fe6ad
The remaining blocker for a TileLang-free path was sparse MLA: `h20_ops/sparse_mla_sm90.py` is a
231-line TileLang program with no Triton counterpart in this repo. FlagGems already has one, and it
turns out to be drop-in.
`flagos-ai/FlagGems` -> `src/flag_gems/fused/flashmla_sparse.py` -> `flash_mla_sparse_fwd`. Pure
Triton. At V4.1's real dims (h_q=64, h_kv=1, d_qk=512, topk=512) it compiles and runs on H20, and
on 7 shapes it agrees with the TileLang kernel within 2 bf16 ULP -- including `topk_length` and
`attn_sink`, the two optional arguments.
The pairing is closer than expected, which is why it needed no adaptation:
prefill_attn_varlen_sparse(query, kv, index_topk, softmax_scale, attn_sink, index_length)
flash_mla_sparse_fwd(q, kv, indices, sm_scale, attn_sink, topk_length)
Causality matches too, and for the same reason. `attn_sm90.py` passes `is_causal=False` because the
indexer only scores positions before each query, so the indices already carry causality -- bounding
on the row index instead collapses the layer to NaN, which already happened once and is documented
at that call site. The FlagGems kernel has no causal parameter either.
Recorded as *not* the answer: `DSA/sparse_mla.py` in the same repo. It hardcodes causal and needs
the V3.2 576-wide split layout; V4.1 folds RoPE inside a single 512-wide latent, so TD = 0 and it
fails to compile on `tl.arange(0, 0)`. Worth writing down because the name looks like the better
match.
What remains, all much smaller than writing a kernel: wire it in behind the existing
`patch_imports` hook; test the decode path separately (it gathers to bf16 with identity indices, a
different call shape than the 7 cases); and convert the return, since the kernel gives
(out, max_logits, lse) with lse in natural log where DeepInfer expects (out, lse) with lse in log2
-- a 1.4427x scale error if dropped.
56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
"""Can FlagGems' Triton sparse-MLA kernel run on H20 (SM90) at DeepSeek-V4.1's shape?
|
|
|
|
Deliberately narrow: first prove the kernel compiles and produces finite output, then (in
|
|
sparse_mla_ab.py) compare it against the TileLang kernel DeepInfer uses today. Mixing the two into
|
|
one script makes a compile failure look like a numerical mismatch.
|
|
|
|
Shape comes from the real config: 64 query heads, 1 kv head, head_dim 512, index_topk 512.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
import traceback
|
|
|
|
import torch
|
|
|
|
# The vendored copy imports `flag_gems.utils...`; its stub sits next to it, so put that on the
|
|
# path here rather than making the caller set PYTHONPATH for a layout detail.
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent / "vendor"))
|
|
|
|
from vendor.fg_flashmla_sparse import flash_mla_sparse_fwd
|
|
|
|
torch.manual_seed(0)
|
|
DEV = "cuda"
|
|
|
|
B, SQ, HQ, HKV, D, TOPK, SKV = 1, 8, 64, 1, 512, 512, 64
|
|
|
|
q = torch.randn((B, SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1
|
|
kv = torch.randn((B, SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1
|
|
|
|
# Valid indices into SKV, no duplicates within a row is not required by the kernel.
|
|
idx = torch.randint(0, SKV, (B, SQ, HKV, TOPK), device=DEV, dtype=torch.int32)
|
|
|
|
# The kernel takes varlen shapes: fold the batch dim away (B is 1 here).
|
|
q2 = q.view(B * SQ, HQ, D).contiguous()
|
|
kv2 = kv.view(B * SKV, HKV, D).contiguous()
|
|
idx2 = idx.view(B * SQ, HKV, TOPK).contiguous()
|
|
|
|
print(f"q {tuple(q2.shape)} kv {tuple(kv2.shape)} idx {tuple(idx2.shape)}", flush=True)
|
|
print("h_q=64, h_kv=1, d=512 -- all inside the kernel's asserts", flush=True)
|
|
|
|
try:
|
|
out, max_logits, lse = flash_mla_sparse_fwd(
|
|
q2, kv2, idx2, sm_scale=D**-0.5, d_v=D
|
|
)
|
|
torch.cuda.synchronize()
|
|
print(f"OK out {tuple(out.shape)} {out.dtype}")
|
|
print(f" max_logits {tuple(max_logits.shape)} {max_logits.dtype}")
|
|
print(f" lse {tuple(lse.shape)} {lse.dtype}")
|
|
print(f" finite: out={torch.isfinite(out).all().item()} "
|
|
f"lse={torch.isfinite(lse).all().item()}")
|
|
print(f" mean |out| {out.float().abs().mean().item():.6f}")
|
|
except Exception:
|
|
print("FAIL")
|
|
traceback.print_exc()
|
|
sys.exit(1)
|