mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +08:00
f25e5d8b6c
The last thing that kept the SM90 path on TileLang was the sparse MLA attention kernel. FlagGems
has a pure-Triton implementation of the same operator, and it turns out to be drop-in once two
differences are handled. With TileLang unimportable, the engine now runs end to end.
`h20_ops/sparse_mla.py` picks the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
means TileLang when it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
two entry points were the only call sites touched. Auto resolution catches `ModuleNotFoundError`
rather than testing for the package by name, because a missing TileLang fails inside
`sparse_mla_sm90`, not at the module lookup -- and it only swallows a *missing* dependency, so a
present-but-broken module still surfaces instead of being swapped out.
Vendored rather than depended upon: `flag_gems/__init__.py` registers ATen implementations
process-wide (`torch.library.Library("aten", "IMPL")` then `replace_customized_ops`). Importing the
package to reach one kernel would swap operators under the whole engine. The PyPI 5.0.2 file has no
`flag_gems` internal imports at all, so there was nothing to gain; it is 329 lines, accumulates in
fp32, and already guards the zero-width rope tail that V4.1 needs (RoPE folded inside the 512-wide
latent). The file header carries the source md5 and both deviations.
Two things it needed, both found by comparison rather than by reading:
* 5.0.2 guards each top-k block by its **position** in the index array (`if ck * BK <= max_col`).
Non-causal that is `SKV-1`, so when `topk > SKV` every block past the first `SKV` entries is
skipped. At SQ=16/SKV=128/topk=512 the output is off by 2.8e-2 -- and any prompt shorter than the
512-entry top-k hits it. Restored the data-dependent guard FlagGems master uses, marked in place.
* Its lse is a natural log; the TileLang one is log2. Unconverted, `_apply_sink` would be off by
1.4427x: plausible output, wrong numbers. The adapter converts, and the test asserts the ratio is
1.0 rather than 0.693 so it cannot regress quietly.
The kernel's operands have no batch axis, while decode hands it a different KV window per query.
Rather than loop, the adapter concatenates the windows and shifts query b's indices by b*seq_kv --
one launch, no extra memory, since the concatenated buffer is the tensor decode already built.
Verification ran in order, each step gating the next: 7 shapes against TileLang (`ALL_MATCH`, <=2
bf16 ULP, including `topk_length` and `attn_sink`), the decode call shape incl. the SWA+full
two-cache case (`ALL_MATCH`), the adapter contract (5 tests), the dispatcher and both `attn_sm90`
entry points (4 tests), then end to end with TileLang blocked: exit=0, 8/8 DP ranks agreeing, zero
kernel fallbacks, correct answers. A control run with TileLang present is unchanged.
Not token-identical, and the doc says so: prompt A shares 15 tokens then diverges, prompt B shares
1. A control run of the TileLang backend twice is bit-identical, so this is accumulation order, not
noise -- consistent with the 2.7% logits drift already recorded for the SM90 quantisation floor.
One TileLang dependency remains, and it is the deliberate one: `fused_sample`'s rejection-sampling
variant, which needs TileLang's Philox stream to reproduce the same tokens. So the pure-Triton run
disables speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`).
87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""A/B the vendored FlagGems 5.0.2 kernel against the TileLang kernel, 7 shapes.
|
|
|
|
Same cases as scripts/sparse_mla_port/ab_vs_tilelang.py, but importing the in-tree copy instead of
|
|
the downloaded one, so this is what actually runs in the engine.
|
|
"""
|
|
|
|
import math
|
|
import sys
|
|
|
|
import torch
|
|
|
|
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
|
|
from deepinfer.h20_ops.vendor.flashmla_sparse import flash_mla_sparse_fwd
|
|
|
|
DEV = "cuda"
|
|
HQ, HKV, D, TOPK = 64, 1, 512, 512
|
|
FAILED = []
|
|
|
|
|
|
def tl_call(q3: torch.Tensor, kv3: torch.Tensor, idx3: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]:
|
|
s_q, s_kv = q3.shape[0], kv3.shape[0]
|
|
out, lse = sparse_mla_fwd_attn(
|
|
q3.view(1, s_q, HQ, D),
|
|
kv3.view(1, s_kv, 1, D),
|
|
idx3.view(1, s_q, 1, -1),
|
|
sm_scale=scale,
|
|
is_causal=False,
|
|
)
|
|
return out.view(s_q, HQ, D), lse.view(s_q, HQ)
|
|
|
|
|
|
def compare(tag: str, a: torch.Tensor, b: torch.Tensor) -> None:
|
|
af, bf = a.float(), b.float()
|
|
if af.shape != bf.shape:
|
|
print(f" FAIL {tag}: shape {tuple(af.shape)} vs {tuple(bf.shape)}")
|
|
FAILED.append(tag)
|
|
return
|
|
diff = (af - bf).abs()
|
|
scale = af.abs().max().item()
|
|
tol = 2.0**-8 * max(scale, 1e-9) * 2
|
|
ok = diff.max().item() <= tol and bool(torch.isfinite(bf).all())
|
|
print(f" {'ok ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f} max diff {diff.max().item():.3e} tol {tol:.3e}")
|
|
if not ok:
|
|
FAILED.append(tag)
|
|
|
|
|
|
def case(tag: str, SQ: int, SKV: int, seed: int, distinct: bool = True) -> None:
|
|
torch.manual_seed(seed)
|
|
q = (torch.randn((SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
|
kv = (torch.randn((SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
|
if distinct and SKV > TOPK:
|
|
idx = torch.stack([torch.randperm(SKV, device=DEV)[:TOPK] for _ in range(SQ)])
|
|
idx = idx.view(SQ, 1, TOPK).to(torch.int32)
|
|
else:
|
|
idx = torch.randint(0, SKV, (SQ, 1, TOPK), device=DEV, dtype=torch.int32)
|
|
idx = idx.contiguous()
|
|
|
|
print(f"\n[{tag}] SQ={SQ} SKV={SKV} distinct={distinct}")
|
|
scale = D**-0.5
|
|
tl_out, tl_lse = tl_call(q, kv, idx, scale)
|
|
torch.cuda.synchronize()
|
|
fg_out, fg_max, fg_lse = flash_mla_sparse_fwd(q, kv, idx, sm_scale=scale, d_v=D)
|
|
torch.cuda.synchronize()
|
|
compare("out", tl_out, fg_out)
|
|
# lse units differ: TileLang returns log2, FlagGems natural log.
|
|
compare("lse (after log2->ln)", tl_lse.float() * math.log(2), fg_lse)
|
|
|
|
|
|
CASES = [
|
|
("prefill, distinct top-k", 64, 2048, 0, True),
|
|
("prefill, duplicate top-k", 32, 512, 1, False),
|
|
("decode-like (SQ=1)", 1, 4096, 2, True),
|
|
("odd SQ", 7, 1024, 3, True),
|
|
("short KV < topk", 16, 128, 4, False),
|
|
("larger prefill", 256, 4096, 5, True),
|
|
("equal SQ==SKV", 512, 512, 6, False),
|
|
]
|
|
|
|
if __name__ == "__main__":
|
|
for c in CASES:
|
|
case(*c)
|
|
print()
|
|
if FAILED:
|
|
print(f"MISMATCH {len(FAILED)}: {FAILED}")
|
|
sys.exit(1)
|
|
print(f"ALL_MATCH ({len(CASES)} cases)")
|