mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
feat(h20): run sparse MLA on FlagGems' Triton kernel, so TileLang is no longer required
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`).
This commit is contained in:
loopyt
parent
c6a65fe6ad
commit
f25e5d8b6c
55
TODO.md
55
TODO.md
@ -595,7 +595,7 @@ import 路径上。修完之后 `import deepinfer` 在纯 Triton 环境不再需
|
||||
把热路径逐个换成 Triton。§2.10b 说明了大头在哪（模型前向直接调 `tile_kernels` 的
|
||||
quant/MoE/MHC/engram，25 处模块级 import / 23 个 API）。
|
||||
|
||||
### 6.0 当前到哪一步了（2026-10-24 晚更新，**端到端已跑通**）
|
||||
### 6.0 当前到哪一步了（2026-10-25 更新，**无 tilelang 端到端已跑通**）
|
||||
|
||||
| # | 事项 | 状态 |
|
||||
|---|---|---|
|
||||
@ -603,10 +603,12 @@ quant/MoE/MHC/engram，25 处模块级 import / 23 个 API）。
|
||||
| 2 | 算子数值等价（23 个 API 实际调用 + 18 用例两模式对比） | ✅ 见 6.2 |
|
||||
| 3 | 路由层 + 25 处 import 替换 | ✅ 已通 |
|
||||
| 4 | **无 tilelang 下 23 个 API 冒烟** | ✅ `ALL_CALLS_OK` |
|
||||
| 5 | **端到端（`DEEPINFER_KERNELS=ref`）** | ✅ **exit=0，8 DP 一致，回答正确** |
|
||||
| 6 | **与编译版逐 token 一致** | ❌ **16/16 样本不同**（见 6.5，已确认非采样噪声） |
|
||||
| 7 | 完全不用 tilelang（含 attention） | ❌ sparse MLA 仍只有 TileLang 实现（见 6.6） |
|
||||
| 8 | 性能 | ❓ 仍未测 |
|
||||
| 5 | **端到端（`DEEPINFER_KERNELS=ref`）** | ✅ exit=0，8 DP 一致，回答正确 |
|
||||
| 6 | **稀疏 MLA 换成 Triton（FlagGems）** | ✅ 见 6.7，7 形状 + decode 5 用例全一致 |
|
||||
| 7 | **屏蔽 tilelang 的端到端** | ✅ **exit=0，0 回退，回答正确**（需关 MTP，见 6.7） |
|
||||
| 8 | **与 TileLang 版逐 token 一致** | ❌ 前 15/1 个 token 后分叉（控制实验证明非噪声，见 6.5） |
|
||||
| 9 | 百分之百不用 tilelang（含采样器 verify 变体） | ❌ `fused_sample` 的 rejection-sampling 仍需（见 6.7） |
|
||||
| 10 | 性能 | ❓ 仍未测 |
|
||||
|
||||
### 6.5 端到端结果：功能对，但**不逐 token 一致**
|
||||
|
||||
@ -659,13 +661,42 @@ causal，用行号当因果边界会导致整层 NaN，这个 bug 已经发生
|
||||
|
||||
验证脚本与 vendor 副本：`scripts/sparse_mla_port/`（`README.md` 有完整对照表与未做项）。
|
||||
|
||||
**还差的三步**（都比重写 kernel 小得多）：
|
||||
1. **接线**：`sparse_mla_sm90.py` 仍 `import tilelang`。用现成的 `patch_imports` 挂一个 Triton
|
||||
后端即可，不动上游调用点。
|
||||
2. **decode 路径单独测**：今天 decode 是 gather→bf16 + identity indices，与上面 7 组形状不同，
|
||||
需要对着 fp8 cache 路径另做一次 A/B。
|
||||
3. **返回值适配**：它返回 `(out, max_logits, lse)`，DeepInfer 要 `(out, lse)`；且 TileLang 的
|
||||
`lse` 是 log2 单位、FlagGems 是自然对数，**必须转换**，否则下游读到的 scale 差 1.4427 倍。
|
||||
### 6.7 接线完成，无 tilelang 端到端跑通（2026-10-25）
|
||||
|
||||
三步都做完了，且每步先过单测再往下走：
|
||||
|
||||
| 步 | 内容 | 结果 |
|
||||
|---|---|---|
|
||||
| 1 | 内联 FlagGems 5.0.2 kernel + 7 形状 A/B | `ALL_MATCH` ≤2 ULP |
|
||||
| 2 | decode 路径 A/B（fp8 cache、5 用例含双缓存） | `ALL_MATCH` |
|
||||
| 3 | wrapper 契约单测（返回元数、lse 单位、dtype、两种调用约定） | `ALL_PASSED` 5 组 |
|
||||
| 4 | 分发层 + `attn_sm90` 两个入口单测 | `ALL_PASSED` 4 组 |
|
||||
| 5 | 屏蔽 tilelang 端到端 | `exit=0`，0 回退，回答正确 |
|
||||
| 6 | 有 tilelang 端到端（对照） | `exit=0`，0 回退，回答正确 |
|
||||
|
||||
**接线方式**：新增 `h20_ops/sparse_mla.py` 作分发层，`DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`
|
||||
显式选择，不设则「能 import tilelang 就用它，否则 Triton」。`attn_sm90` 只改了 import 那一行，
|
||||
两个调用点未动。每个 rank 打印一次实际选择。
|
||||
|
||||
**为什么内联而不依赖 `flag_gems`**：`flag_gems/__init__.py` 会 `torch.library.Library("aten","IMPL")`
|
||||
+ `replace_customized_ops`，**进程级替换 ATen 算子**；为了一个 kernel 去 import 它，会把整个引擎的
|
||||
算子悄悄换掉。PyPI 5.0.2 的那个文件零内部依赖，内联后只需 torch/triton。
|
||||
|
||||
**踩到并修掉的两个坑**：
|
||||
1. **FlagGems 5.0.2 的 block 跳过守卫是错的**：`if ck * BK <= max_col` 按索引在数组里的**位置**
|
||||
判断，非因果时 `max_col = SKV-1`，于是 `topk > SKV` 时超出前 SKV 个的块被整块跳过。
|
||||
SQ=16/SKV=128/topk=512 时输出差 **2.8e-2**。任何短于 512 token 的 prompt 都会踩到。
|
||||
已本地修正为数据相关守卫（master 版就是这个写法），文件里有 `LOCAL FIX` 标注。
|
||||
2. **lse 单位**：kernel 出自然对数，TileLang 出 log2。不转换会让 `_apply_sink` 的系数差 1.4427 倍。
|
||||
单测里专门断言 ratio=1.0（未转换会是 0.693），避免以后静默回归。
|
||||
|
||||
**仍需注意**：
|
||||
- 纯 Triton 端到端**必须 `DEEPINFER_USE_NEXTN_DECODE=0`**（关投机解码）。开着 MTP 时会走
|
||||
`fused_sample` 的 rejection-sampling 变体，那是唯一还依赖 TileLang 的 kernel（Philox 随机流 +
|
||||
固定 reduction 顺序，Triton 版会改变生成文本，见 `triton/__init__.py` 的说明）。
|
||||
- **不与 TileLang 版逐 token 一致**：prompt A 前 15 个 token 一致、第 16 个分叉；prompt B 第 2 个
|
||||
分叉。控制实验（TileLang 连跑两次）**逐 token 相同**，故非采样噪声，是累加顺序差异沿 40+ 层
|
||||
累积，与 §6.5 记录的 2.7% logits 漂移同源。
|
||||
|
||||
`tilelang_ops/cuda/` 那 29 个模块已有 `tilelang_ops/triton/` 对应版本，靠
|
||||
`DEEPINFER_OPS_BACKEND=triton` 切换，本来就不是缺口。
|
||||
|
||||
@ -6,8 +6,9 @@ re-implemented here and registered through the existing `patch_imports` hook, so
|
||||
sites stay unchanged — and `attention_backend.py` is never imported on SM90, which matters
|
||||
because its module-level `from flash_mla import ...` cannot be satisfied on this platform.
|
||||
|
||||
* :func:`prefill_attn_varlen_sparse` — dense KV is already in memory, so the sparse TileLang
|
||||
kernel (``h20_ops.sparse_mla_sm90``) consumes it directly, with indices selecting top-k.
|
||||
* :func:`prefill_attn_varlen_sparse` — dense KV is already in memory, so the sparse attention
|
||||
kernel (``h20_ops.sparse_mla``, either backend) consumes it directly, with indices selecting
|
||||
top-k.
|
||||
* :func:`decode_mla_mqa_sparse` — KV lives in the paged fp8 cache; the selected rows are
|
||||
gathered as raw bytes (``tilelang_ops.gather`` zero-fills out-of-range rows) and dequantized
|
||||
to bf16, then the same kernel attends over them with identity indices and no causal mask.
|
||||
@ -19,7 +20,7 @@ from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
|
||||
from deepinfer.h20_ops.sparse_mla import sparse_mla_fwd_attn
|
||||
|
||||
# V41_FORMAT: 512 fp8 values followed by 512/32 = 16 UE8M0 scale bytes per token.
|
||||
CACHE_ELEMENTS = 512
|
||||
|
||||
112
deepinfer/h20_ops/sparse_mla.py
Normal file
112
deepinfer/h20_ops/sparse_mla.py
Normal file
@ -0,0 +1,112 @@
|
||||
"""Pick the SM90 sparse-MLA backend, so callers do not have to know which exists.
|
||||
|
||||
Two implementations, both numerically verified against each other on every shape the engine uses:
|
||||
|
||||
* `sparse_mla_sm90` -- TileLang. The original; needs `tilelang` importable.
|
||||
* `sparse_mla_triton` -- FlagGems' Triton kernel. Needs nothing beyond torch and triton.
|
||||
|
||||
`DEEPINFER_SPARSE_MLA_BACKEND` selects explicitly (`tilelang` / `triton`). Unset means auto:
|
||||
TileLang when it imports, Triton otherwise. Auto has to catch `ModuleNotFoundError` rather than test
|
||||
for the package by name, because on a host where TileLang was never installed the import fails at
|
||||
`import tilelang` inside `sparse_mla_sm90`, not at the top-level module lookup.
|
||||
|
||||
Selection is by import capability, not by platform, deliberately: a card can have TileLang installed
|
||||
and still be the wrong choice for a given run, which is what the explicit override is for. The
|
||||
chosen backend is logged once at first use -- an end-to-end log that does not say this cannot be
|
||||
used to attribute a numerical difference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ENV = "DEEPINFER_SPARSE_MLA_BACKEND"
|
||||
_RESOLVED: Callable[..., tuple[torch.Tensor, torch.Tensor]] | None = None
|
||||
_NAME: str | None = None
|
||||
|
||||
|
||||
def _load_tilelang() -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
|
||||
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
|
||||
|
||||
return sparse_mla_fwd_attn
|
||||
|
||||
|
||||
def _load_triton() -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
|
||||
from deepinfer.h20_ops.sparse_mla_triton import sparse_mla_fwd_attn
|
||||
|
||||
return sparse_mla_fwd_attn
|
||||
|
||||
|
||||
def _resolve() -> tuple[Callable[..., tuple[torch.Tensor, torch.Tensor]], str]:
|
||||
forced = os.getenv(_ENV, "").strip().lower()
|
||||
if forced:
|
||||
if forced == "tilelang":
|
||||
try:
|
||||
return _load_tilelang(), "tilelang"
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
f"{_ENV}=tilelang was requested but the TileLang backend is not importable "
|
||||
f"({exc}). Unset it to fall back to the Triton backend, or install TileLang."
|
||||
) from exc
|
||||
if forced == "triton":
|
||||
return _load_triton(), "triton"
|
||||
raise RuntimeError(f"{_ENV}={forced!r} is not recognised; use 'tilelang' or 'triton'")
|
||||
|
||||
try:
|
||||
return _load_tilelang(), "tilelang"
|
||||
except ModuleNotFoundError as exc:
|
||||
# Only swallow a *missing* dependency. A module that is present but broken must not be
|
||||
# silently replaced by the other backend, or a real failure would look like a working run.
|
||||
if "tilelang" not in str(exc) and "tvm_ffi" not in str(exc):
|
||||
raise
|
||||
logger.info("sparse MLA: TileLang unavailable (%s); using the Triton backend", exc)
|
||||
return _load_triton(), "triton"
|
||||
|
||||
|
||||
def backend() -> str:
|
||||
"""Name of the backend that `sparse_mla_fwd_attn` will use. Resolves on first call."""
|
||||
global _RESOLVED, _NAME
|
||||
if _RESOLVED is None:
|
||||
_RESOLVED, _NAME = _resolve()
|
||||
logger.info("sparse MLA backend: %s", _NAME)
|
||||
return str(_NAME)
|
||||
|
||||
|
||||
def sparse_mla_fwd_attn(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
*,
|
||||
sm_scale: float | None = None,
|
||||
is_causal: bool = True,
|
||||
block_I: int = 64,
|
||||
num_stages: int = 2,
|
||||
threads: int = 256,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sparse MLA forward on the selected backend. Same contract as `sparse_mla_sm90`'s function."""
|
||||
global _RESOLVED
|
||||
if _RESOLVED is None:
|
||||
backend()
|
||||
assert _RESOLVED is not None
|
||||
return _RESOLVED(
|
||||
q,
|
||||
kv,
|
||||
indices,
|
||||
sm_scale=sm_scale,
|
||||
is_causal=is_causal,
|
||||
block_I=block_I,
|
||||
num_stages=num_stages,
|
||||
threads=threads,
|
||||
)
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
"""Forget the resolved backend. Only for tests that change the environment mid-process."""
|
||||
global _RESOLVED, _NAME
|
||||
_RESOLVED, _NAME = None, None
|
||||
117
deepinfer/h20_ops/sparse_mla_triton.py
Normal file
117
deepinfer/h20_ops/sparse_mla_triton.py
Normal file
@ -0,0 +1,117 @@
|
||||
"""Sparse MLA on SM90 using FlagGems' Triton kernel instead of TileLang.
|
||||
|
||||
Drop-in for :func:`deepinfer.h20_ops.sparse_mla_sm90.sparse_mla_fwd_attn`: same arguments, same
|
||||
return contract, so `attn_sm90` can pick either without its call sites changing.
|
||||
|
||||
Two differences from the vendored kernel are handled here, and neither is optional:
|
||||
|
||||
1. **Return arity.** The kernel returns ``(out, max_logits, lse)``; this returns ``(out, lse)``.
|
||||
2. **lse units.** The kernel's lse is a natural log. The TileLang kernel -- and therefore every
|
||||
consumer, including ``_apply_sink`` which does ``sigmoid(lse * LN2 - sink)`` -- uses log2. Handing
|
||||
the natural-log value straight through would scale every sink decision by 1.4427x and producing
|
||||
plausible-looking but wrong output, which is the kind of bug that survives a smoke test.
|
||||
|
||||
The kernel's operands have no batch axis, while callers pass one. Both conventions in this repo are
|
||||
accepted and unified here:
|
||||
|
||||
* Shared KV (prefill): ``kv.shape[0] == 1`` and ``q.shape[0]`` is the number of query tokens in one
|
||||
pack. The batch axis is folded away and indices are used as given.
|
||||
* Per-query KV (decode): ``kv.shape[0] == q.shape[0]``. Each query owns its own gathered window, so
|
||||
the windows are concatenated into one KV and query ``b`` gets its indices shifted by
|
||||
``b * seq_kv``. This costs no extra memory -- the concatenated tensor is the one the caller
|
||||
already materialised -- and it stays a single launch. See ``scripts/sparse_mla_port/ab_decode.py``
|
||||
for the comparison that established this against the two-cache (SWA + full) decode case too.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.h20_ops.vendor.flashmla_sparse import flash_mla_sparse_fwd
|
||||
|
||||
# lse comes out of the vendored kernel in natural-log units; the rest of the engine speaks log2.
|
||||
_INV_LN2 = 1.4426950408889634
|
||||
|
||||
|
||||
def sparse_mla_fwd_attn(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
*,
|
||||
sm_scale: float | None = None,
|
||||
is_causal: bool = True,
|
||||
block_I: int = 64,
|
||||
num_stages: int = 2,
|
||||
threads: int = 256,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sparse MLA forward; mirrors ``sparse_mla_sm90.sparse_mla_fwd_attn``.
|
||||
|
||||
Args:
|
||||
q: ``[batch, seq, heads, head_dim]`` bf16, RoPE already applied.
|
||||
kv: ``[batch_or_1, seq_kv, 1, head_dim]`` bf16.
|
||||
indices: ``[batch, seq, 1, topk]`` int32; negative entries mean "invalid".
|
||||
sm_scale: softmax scale; defaults to ``head_dim ** -0.5``.
|
||||
is_causal: must be False. See below.
|
||||
block_I, num_stages, threads: accepted for signature compatibility only. The vendored
|
||||
kernel derives its own tiling and autotunes ``num_stages``/``num_warps``.
|
||||
|
||||
Returns:
|
||||
``(out [batch, seq, heads, head_dim], lse [batch, seq, heads])`` with lse in **log2** units.
|
||||
"""
|
||||
if is_causal:
|
||||
# The vendored kernel implements causality by bounding on the *row index*, which is only
|
||||
# meaningful when an index value equals its position. DeepInfer's indices are absolute
|
||||
# positions in a packed KV buffer that `kv_index_cat` has already shifted, so that bound
|
||||
# would drop nearly every index -- the failure mode documented at attn_sm90.py:110, which
|
||||
# collapses the layer to NaN. Both call sites pass False; refuse rather than guess.
|
||||
raise NotImplementedError(
|
||||
"the Triton sparse-MLA backend only supports is_causal=False: causality is carried by "
|
||||
"the indices, and bounding on the row index is wrong for absolute-position indices"
|
||||
)
|
||||
|
||||
batch, seq, heads, dim = q.shape
|
||||
kv_batch, seq_kv, kv_heads, head_dim_kv = kv.shape
|
||||
if kv_heads != 1:
|
||||
raise ValueError(f"kv must carry a single KV head, got {kv_heads}")
|
||||
if head_dim_kv != dim:
|
||||
raise ValueError(f"q and kv disagree on head_dim: {dim} vs {head_dim_kv}")
|
||||
|
||||
# Reject shapes the kernel cannot express before it fails somewhere less legible. It asserts
|
||||
# these itself, but at a point where the error no longer names the caller.
|
||||
if heads not in (64, 128):
|
||||
raise ValueError(f"the Triton sparse-MLA kernel supports heads in {{64, 128}}, got {heads}")
|
||||
|
||||
scale = dim**-0.5 if sm_scale is None else sm_scale
|
||||
idx4 = indices.reshape(batch, seq, 1, -1).to(torch.int32)
|
||||
|
||||
if kv_batch == 1:
|
||||
# Shared KV (prefill, and decode with a single query): fold the batch axis away and use the
|
||||
# indices as given. Checked first because prefill legitimately passes batch=1, which would
|
||||
# otherwise match the per-query condition below.
|
||||
flat_kv = kv.reshape(seq_kv, 1, head_dim_kv).contiguous()
|
||||
idx = idx4.reshape(batch * seq, 1, -1).contiguous()
|
||||
elif kv_batch == batch:
|
||||
# Per-query KV (decode): concatenate the windows, then shift query b's indices by
|
||||
# b * seq_kv so they address its own window in the concatenated buffer.
|
||||
flat_kv = kv.reshape(batch * seq_kv, 1, head_dim_kv).contiguous()
|
||||
bases = (torch.arange(batch, dtype=torch.int32, device=idx4.device) * seq_kv).view(batch, 1, 1, 1)
|
||||
idx = (idx4 + bases).reshape(batch * seq, 1, -1).contiguous()
|
||||
else:
|
||||
raise ValueError(
|
||||
"kv's batch axis must be 1 (shared KV) or equal to q's (one window per query); "
|
||||
f"got kv.shape[0]={kv_batch} with q.shape[0]={batch}"
|
||||
)
|
||||
|
||||
out, _max_logits, lse_natural = flash_mla_sparse_fwd(
|
||||
q.reshape(batch * seq, heads, dim).contiguous(),
|
||||
flat_kv,
|
||||
idx,
|
||||
sm_scale=float(scale),
|
||||
d_v=dim,
|
||||
)
|
||||
|
||||
out = out.view(batch, seq, heads, dim)
|
||||
# Natural log -> log2, so downstream sink maths keeps working unchanged. float32 matches what
|
||||
# the TileLang kernel returns for lse, which callers already rely on.
|
||||
lse = lse_natural.view(batch, seq, heads).to(torch.float32) * _INV_LN2
|
||||
return out, lse
|
||||
1
deepinfer/h20_ops/vendor/__init__.py
vendored
Normal file
1
deepinfer/h20_ops/vendor/__init__.py
vendored
Normal file
@ -0,0 +1 @@
|
||||
"""Vendored third-party code. Excluded from lint and format by `extend-exclude` in pyproject."""
|
||||
392
deepinfer/h20_ops/vendor/flashmla_sparse.py
vendored
Normal file
392
deepinfer/h20_ops/vendor/flashmla_sparse.py
vendored
Normal file
@ -0,0 +1,392 @@
|
||||
# Copyright 2026 FlagOS Contributors
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
# Vendored from FlagGems, not written here. Provenance:
|
||||
#
|
||||
# package : flag_gems 5.0.2 (PyPI)
|
||||
# module : flag_gems/fused/flashmla_sparse.py
|
||||
# md5 : f38ce8a998f7db9f4796a681d9ccf475
|
||||
# license : Apache-2.0 (above)
|
||||
#
|
||||
# Copied from the wheel with exactly two deviations, both stated here because vendored code that
|
||||
# drifts silently is worse than no vendored code:
|
||||
#
|
||||
# 1. This header (the wheel's copy omits the license text the license requires).
|
||||
# 2. One fix in the kernel body, marked `LOCAL FIX, not upstream`: the released file guards each
|
||||
# top-k block with `if ck * BK <= max_col`, which drops blocks by their *position* in the index
|
||||
# array rather than by their contents. That is only sound when index value tracks position,
|
||||
# and it is wrong whenever topk > SKV -- which DeepInfer hits on any prompt shorter than the
|
||||
# 512-entry top-k. FlagGems master uses the data-dependent form and is correct.
|
||||
#
|
||||
# Keep the body byte-identical to upstream apart from that one change: do not run a formatter over
|
||||
# it, or the next re-vendor becomes a manual diff of stylistic noise. `pyproject.toml` excludes
|
||||
# this file from both lint and format for that reason.
|
||||
#
|
||||
# `scripts/sparse_mla_port/ab_vendored.py` is what keeps it honest -- it compares against the
|
||||
# TileLang kernel on every shape the engine uses, so upstream drift or a bad edit shows up as a
|
||||
# test failure rather than a silent numerical change.
|
||||
#
|
||||
# Why vendored instead of depending on `flag_gems`:
|
||||
#
|
||||
# 1. `flag_gems/__init__.py` registers ATen implementations process-wide
|
||||
# (`torch.library.Library("aten", "IMPL")` then `runtime.replace_customized_ops`). Importing
|
||||
# the package to reach one kernel would silently swap operators under the whole engine, which
|
||||
# changes numerics of code that has nothing to do with attention. This module imports only
|
||||
# torch and triton, so it cannot do that.
|
||||
# 2. The released 5.0.2 file has no `flag_gems` internal imports at all, so there is nothing to
|
||||
# gain from the dependency.
|
||||
#
|
||||
# Why 5.0.2 rather than FlagGems master: master's copy is ~1200 lines with a TLE fast path gated on
|
||||
# triton internals, and its accumulation is fp16. This one is 329 lines, accumulates in fp32, and
|
||||
# guards `TDP > 0` -- which is what makes it work at V4.1's shape, where RoPE is folded inside the
|
||||
# 512-wide latent and the rope tail is therefore zero-width.
|
||||
#
|
||||
# Do not edit the kernel body. If a change is genuinely needed, make it upstream in FlagGems and
|
||||
# re-vendor with the md5 above updated.
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
flash_mla_sparse_fwd_configs = [
|
||||
triton.Config({"num_stages": 4, "num_warps": 8}),
|
||||
triton.Config({"num_stages": 2, "num_warps": 4}),
|
||||
]
|
||||
|
||||
|
||||
@triton.autotune( # Decorate the kernel
|
||||
configs=flash_mla_sparse_fwd_configs,
|
||||
key=["K", "is_causal"],
|
||||
)
|
||||
@triton.jit
|
||||
def triton_flash_mla_sparse_fwd(
|
||||
q,
|
||||
kv,
|
||||
indices,
|
||||
attn_sink,
|
||||
topk_length,
|
||||
sm_scale: tl.constexpr,
|
||||
output,
|
||||
max_logits,
|
||||
lse,
|
||||
stride_qh,
|
||||
stride_qm,
|
||||
stride_qd,
|
||||
stride_kvg,
|
||||
stride_kvn,
|
||||
stride_kvd,
|
||||
stride_tg,
|
||||
stride_tm,
|
||||
stride_tt, # indices dim
|
||||
stride_attn_sink_h,
|
||||
stride_topk_length_s,
|
||||
stride_oh,
|
||||
stride_om,
|
||||
stride_od,
|
||||
stride_mh,
|
||||
stride_mm,
|
||||
stride_lh,
|
||||
stride_lm,
|
||||
SQ: tl.constexpr, # seqlen
|
||||
SKV: tl.constexpr, # seqlen_kv
|
||||
K: tl.constexpr, # topk
|
||||
D: tl.constexpr, # QKV dim
|
||||
TD: tl.constexpr, # tail dim
|
||||
DP: tl.constexpr,
|
||||
TDP: tl.constexpr,
|
||||
G: tl.constexpr, # group_size
|
||||
BK: tl.constexpr,
|
||||
BH: tl.constexpr,
|
||||
is_causal: tl.constexpr,
|
||||
q_idx_i64: tl.constexpr,
|
||||
output_idx_i64: tl.constexpr,
|
||||
HAVE_ATTN_SINK: tl.constexpr,
|
||||
HAVE_TOPK_LENGTH: tl.constexpr,
|
||||
):
|
||||
i_sq, i_gbh = tl.program_id(0), tl.program_id(1)
|
||||
i_g, i_bh = i_gbh // G, i_gbh % G
|
||||
if not q_idx_i64:
|
||||
q_base = q + i_sq * stride_qm + i_gbh * (BH * stride_qh)
|
||||
else:
|
||||
q_base = q + i_sq * tl.cast(stride_qm, tl.int64) + i_gbh * (BH * stride_qh)
|
||||
tq_base = q_base + D * stride_qd
|
||||
kv_base = kv + i_g * stride_kvg
|
||||
tkv_base = kv_base + D * stride_kvd
|
||||
t_base = indices + i_sq * stride_tm + i_g * stride_tg
|
||||
attn_sink_ptr = (
|
||||
attn_sink + i_gbh * (BH * stride_attn_sink_h) if HAVE_ATTN_SINK else 0
|
||||
)
|
||||
topk_length_ptr = (
|
||||
topk_length + i_sq * stride_topk_length_s if HAVE_TOPK_LENGTH else 0
|
||||
)
|
||||
if not output_idx_i64:
|
||||
o_base = output + i_sq * stride_om + i_gbh * (BH * stride_oh)
|
||||
else:
|
||||
o_base = output + i_sq * tl.cast(stride_om, tl.int64) + i_gbh * (BH * stride_oh)
|
||||
max_log_base = max_logits + i_sq * stride_mm + i_gbh * (BH * stride_mh)
|
||||
l_base = lse + i_sq * stride_lm + i_gbh * (BH * stride_lh)
|
||||
|
||||
offs_h = tl.arange(0, BH)
|
||||
offs_d = tl.arange(0, DP)
|
||||
offs_td = tl.arange(0, TDP) if TDP > 0 else None
|
||||
offs_od = tl.arange(0, DP)
|
||||
offs_t = tl.arange(0, BK)
|
||||
mask_h = i_bh * BH + offs_h < G
|
||||
mask_d = offs_d < D
|
||||
mask_td = offs_td < TD if TDP > 0 else None
|
||||
mask_od = mask_d
|
||||
|
||||
q_ptr = q_base + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd
|
||||
q_msk = mask_h[:, None] & mask_d[None, :]
|
||||
q_blk = tl.load(q_ptr, q_msk, other=0.0).to(tl.float32)
|
||||
|
||||
tq_blk = None
|
||||
if TDP > 0:
|
||||
tq_ptr = tq_base + offs_h[:, None] * stride_qh + offs_td[None, :] * stride_qd
|
||||
tq_msk = mask_h[:, None] & mask_td[None, :]
|
||||
tq_blk = tl.load(tq_ptr, tq_msk, other=0.0).to(tl.float32)
|
||||
|
||||
max_log = tl.full([BH], float("-inf"), dtype=tl.float32)
|
||||
sum_exp = tl.full([BH], 0.0, dtype=tl.float32)
|
||||
acc = tl.zeros([BH, DP], dtype=tl.float32)
|
||||
qk = tl.zeros([BH, BK], dtype=tl.float32)
|
||||
|
||||
max_col = i_sq if is_causal else SKV - 1
|
||||
topk_len = tl.load(topk_length_ptr).to(tl.int32) if HAVE_TOPK_LENGTH else K
|
||||
|
||||
NK = tl.cdiv(K, BK)
|
||||
for ck in range(NK):
|
||||
# step1: load indices
|
||||
t_ptr = (BK * ck + offs_t) * stride_tt
|
||||
t_msk = t_ptr < topk_len
|
||||
t_ptr += t_base
|
||||
kv_ids = tl.load(t_ptr, t_msk, other=-1)
|
||||
mask_ids = (kv_ids <= max_col) & (kv_ids >= 0)
|
||||
# filter invalid index that may cause overflow in mul
|
||||
kv_ids = tl.where(mask_ids, kv_ids, 0)
|
||||
|
||||
# LOCAL FIX, not upstream -- see the file header. Upstream guards this block with
|
||||
# `if ck * BK <= max_col`, i.e. by the block's position in the index array. With
|
||||
# is_causal=False, max_col is SKV-1, so when topk > SKV every block past the first SKV
|
||||
# entries is skipped although it holds valid indices: at SQ=16/SKV=128/topk=512 the output
|
||||
# is off by 2.8e-2 against the TileLang kernel. The data-dependent form below is what
|
||||
# FlagGems master uses (`valid0 = t_msk0 & (kv_ids0 <= max_col) & (kv_ids0 >= 0)`).
|
||||
# if mask_ids.max(0) > 0:
|
||||
if tl.max(mask_ids, axis=0) > 0:
|
||||
# step2: gather kv with indices
|
||||
kv_ptr = (
|
||||
kv_base + offs_d[:, None] * stride_kvd + kv_ids[None, :] * stride_kvn
|
||||
)
|
||||
kv_msk = mask_d[:, None] & mask_ids[None, :]
|
||||
kv_blk = tl.load(kv_ptr, kv_msk, other=0.0).to(tl.float32) # [DP, BK]
|
||||
|
||||
# step3: (q @ kv) * sm_scale
|
||||
qk = tl.dot(q_blk, kv_blk)
|
||||
if TDP > 0:
|
||||
tkv_ptr = (
|
||||
tkv_base
|
||||
+ offs_td[:, None] * stride_kvd
|
||||
+ kv_ids[None, :] * stride_kvn
|
||||
)
|
||||
tkv_msk = mask_td[:, None] & mask_ids[None, :]
|
||||
tkv_blk = tl.load(tkv_ptr, tkv_msk, other=0.0).to(
|
||||
tl.float32
|
||||
) # [TDP, BK]
|
||||
qk = tl.dot(tq_blk, tkv_blk, qk) * sm_scale
|
||||
else:
|
||||
qk = qk * sm_scale
|
||||
|
||||
# step4: preprocess for logsumexp
|
||||
qk = tl.where(mask_ids[None, :], qk, float("-inf")) # [BH, BK]
|
||||
# step5: lse=log2sumexp2(qk), loop part
|
||||
new_max = tl.maximum(max_log, tl.max(qk, axis=1)) # [BH]
|
||||
# avoid nan generated by ((-inf) - (-inf))
|
||||
tmp = qk - new_max[:, None]
|
||||
tmp = tl.where(
|
||||
(~mask_ids[None, :]) & (new_max[:, None] == float("-inf")),
|
||||
float("-inf"),
|
||||
tmp,
|
||||
)
|
||||
exp_qk = tl.math.exp(tmp) # [BH, BK]
|
||||
sum_qk = tl.sum(exp_qk, axis=1) # [BH]
|
||||
# avoid nan generated by ((-inf) - (-inf))
|
||||
tmp2 = max_log - new_max
|
||||
tmp2 = tl.where(
|
||||
(max_log == float("-inf")) & (new_max == float("-inf")),
|
||||
float("-inf"),
|
||||
tmp2,
|
||||
)
|
||||
alpha = tl.math.exp(tmp2) # [BH]
|
||||
sum_exp = tl.fma(sum_exp, alpha, sum_qk) # [BH]
|
||||
acc = acc * alpha[:, None] # [BH, DP]
|
||||
# step6: exp2(qk-lse) @ gathered_kv.trans(), loop part
|
||||
acc = tl.dot(exp_qk, kv_blk.trans(), acc) # [BH, DP]
|
||||
max_log = new_max
|
||||
|
||||
# step7: store max_logits
|
||||
max_log_ptr = max_log_base + offs_h * stride_lh
|
||||
tl.store(max_log_ptr, max_log, mask_h) # [BH], float32
|
||||
|
||||
# step8: lse=log2sumexp2(qk) final part, store lse
|
||||
orig_lse = max_log + tl.math.log(sum_exp)
|
||||
lse_out = tl.where(orig_lse == float("-inf"), float("inf"), orig_lse)
|
||||
l_ptr = l_base + offs_h * stride_lh
|
||||
l_msk = mask_h
|
||||
tl.store(l_ptr, lse_out, l_msk) # [BH], float32
|
||||
|
||||
# step9: exp2(qk-lse) @ gathered_kv.trans(), final part
|
||||
if HAVE_ATTN_SINK:
|
||||
# step10: attn_sink
|
||||
exp_max_qk = tl.math.exp(max_log) # [BH]
|
||||
exp_orig_lse = tl.math.exp(orig_lse)
|
||||
sink = tl.load(attn_sink_ptr + offs_h).to(tl.float32) # [BH]
|
||||
exp_sink = tl.math.exp(sink)
|
||||
sum_exp_new_lse = exp_orig_lse + exp_sink
|
||||
# avoid divide 0
|
||||
sum_exp_new_lse = tl.where(sum_exp_new_lse == 0.0, 1.0, sum_exp_new_lse)
|
||||
factor = exp_max_qk / sum_exp_new_lse
|
||||
out_vals = acc * factor[:, None]
|
||||
else:
|
||||
# avoid divide 0
|
||||
sum_exp = tl.where(sum_exp == 0, 1.0, sum_exp)
|
||||
out_vals = acc / sum_exp[:, None]
|
||||
|
||||
# step11: store output
|
||||
o_ptr = (
|
||||
o_base + offs_h[:, None] * stride_oh + offs_od[None, :] * stride_od
|
||||
) # [BH, DP]
|
||||
o_msk = mask_h[:, None] & mask_od[None, :]
|
||||
tl.store(o_ptr, out_vals.to(q_blk.dtype), o_msk)
|
||||
|
||||
|
||||
def flash_mla_sparse_fwd(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
sm_scale: float,
|
||||
d_v: int = 512,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
topk_length: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Sparse attention prefill kernel
|
||||
|
||||
Args:
|
||||
q: [s_q, h_q, d_qk], bfloat16
|
||||
kv: [s_kv, h_kv, d_qk], bfloat16
|
||||
indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv
|
||||
sm_scale: float
|
||||
d_v: The dimension of value vectors. Can only be 512
|
||||
attn_sink: optional, [h_q], float32.
|
||||
If attn_sink is provided, when computing output, output will be additionally multiplied by
|
||||
exp(lse) / (exp(lse) + exp(attn_sink)). +-inf in attn_sink will be handled normally (i.e., -inf has no
|
||||
effect, +inf will make corresponding output all zeros).
|
||||
This argument has no effect on lse and max_logits.
|
||||
topk_length: optional, [s_q], int32.
|
||||
If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]],
|
||||
ignoring later k/v tokens (even if provided in indices). In extremely rare cases (topk_length provided,
|
||||
there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token
|
||||
containing NaN), operator output will contain NaN, so please avoid this situation.
|
||||
|
||||
Returns:
|
||||
(output, max_logits, lse)
|
||||
Please refer to tests/ref.py for the precise definitions of these parameters.
|
||||
- output: [s_q, h_q, d_v], bfloat16
|
||||
- max_logits: [s_q, h_q], float
|
||||
- lse: [s_q, h_q], float, log-sum-exp of attention scores
|
||||
"""
|
||||
is_causal = False # turn off opt for causal sparse attention
|
||||
assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
|
||||
SQ, H, DT = q.shape
|
||||
SKV, VG, _ = kv.shape
|
||||
|
||||
assert d_v == 512, "Unsupported d_v"
|
||||
D = d_v
|
||||
|
||||
assert kv.shape[-1] == DT
|
||||
TD = DT - D
|
||||
DP = triton.next_power_of_2(D)
|
||||
TDP = triton.next_power_of_2(TD)
|
||||
_, _, K = indices.shape
|
||||
assert indices.shape == (SQ, VG, K)
|
||||
if attn_sink is not None:
|
||||
assert attn_sink.shape == (H,), "attn_sink error shape"
|
||||
if topk_length is not None:
|
||||
assert topk_length.shape == (SQ,), "topk_length error shape"
|
||||
|
||||
# check from FlashMLA
|
||||
assert VG == 1, "h_kv is expected to be 1"
|
||||
assert H == 64 or H == 128, "Unsupported h_q"
|
||||
assert DT == 576 or DT == 512, "Unsupported d_qk"
|
||||
|
||||
G = H // VG
|
||||
BH = max(16, min(32, triton.next_power_of_2(G)))
|
||||
NH = triton.cdiv(G, BH)
|
||||
BK = 16 # used to be out of memory for 32
|
||||
output = torch.zeros((SQ, H, D), device=q.device, dtype=q.dtype)
|
||||
max_logits = torch.full(
|
||||
(SQ, H), float("-inf"), device=q.device, dtype=torch.float32
|
||||
)
|
||||
lse = torch.full((SQ, H), float("-inf"), device=q.device, dtype=torch.float32)
|
||||
INT32_MAX = 2147483647
|
||||
q_idx_i64 = q.numel() > INT32_MAX
|
||||
output_idx_i64 = output.numel() > INT32_MAX
|
||||
grid = (SQ, VG * NH, 1)
|
||||
triton_flash_mla_sparse_fwd[grid](
|
||||
q,
|
||||
kv,
|
||||
indices,
|
||||
attn_sink,
|
||||
topk_length,
|
||||
sm_scale,
|
||||
output,
|
||||
max_logits,
|
||||
lse,
|
||||
q.stride(1),
|
||||
q.stride(0),
|
||||
q.stride(2),
|
||||
kv.stride(1),
|
||||
kv.stride(0),
|
||||
kv.stride(2),
|
||||
indices.stride(1),
|
||||
indices.stride(0),
|
||||
indices.stride(2),
|
||||
attn_sink.stride(0) if attn_sink is not None else 0,
|
||||
topk_length.stride(0) if topk_length is not None else 0,
|
||||
output.stride(1),
|
||||
output.stride(0),
|
||||
output.stride(2),
|
||||
max_logits.stride(1),
|
||||
max_logits.stride(0),
|
||||
lse.stride(1),
|
||||
lse.stride(0),
|
||||
SQ,
|
||||
SKV,
|
||||
K,
|
||||
D,
|
||||
TD,
|
||||
DP,
|
||||
TDP,
|
||||
G,
|
||||
BK,
|
||||
BH,
|
||||
is_causal,
|
||||
q_idx_i64,
|
||||
output_idx_i64,
|
||||
attn_sink is not None,
|
||||
topk_length is not None,
|
||||
)
|
||||
return output, max_logits, lse
|
||||
@ -1,14 +1,23 @@
|
||||
# FlagGems sparse-MLA as the SM90 attention backend
|
||||
|
||||
The SM90 sparse MLA kernel (`deepinfer/h20_ops/sparse_mla_sm90.py`) is TileLang-only, so with
|
||||
TileLang blocked the model cannot start. FlagGems ships a pure-Triton implementation of the same
|
||||
operator, and at DeepSeek-V4.1's exact shape it agrees with the TileLang kernel.
|
||||
TileLang blocked the model could not start. FlagGems ships a pure-Triton implementation of the same
|
||||
operator; it is now vendored, verified against the TileLang kernel, and wired in behind a dispatcher
|
||||
(`deepinfer/h20_ops/sparse_mla.py`). With TileLang unimportable the engine runs end to end on it.
|
||||
|
||||
## Source
|
||||
|
||||
`flagos-ai/FlagGems`, file `src/flag_gems/fused/flashmla_sparse.py`, function
|
||||
`flash_mla_sparse_fwd`. Copied verbatim into `vendor/` for the tests; it depends only on
|
||||
`flag_gems.utils.triton_version_utils` (a Triton version probe), stubbed in `vendor/flag_gems/`.
|
||||
`flash_mla_sparse_fwd`. Vendored as `deepinfer/h20_ops/vendor/flashmla_sparse.py` from the **PyPI
|
||||
5.0.2 wheel**, which has no `flag_gems` internal imports at all -- only torch and triton. That
|
||||
matters: importing `flag_gems` itself registers ATen implementations process-wide
|
||||
(`torch.library.Library("aten", "IMPL")` then `runtime.replace_customized_ops`), so depending on the
|
||||
package to reach one kernel would swap operators under the whole engine. The file header records the
|
||||
source md5 and the two local deviations.
|
||||
|
||||
Master's copy (~1200 lines, TLE fast path, fp16 accumulate) also works and was verified first; 5.0.2
|
||||
is used instead because it is a release, is 329 lines, accumulates in fp32, and already guards the
|
||||
zero-width rope tail that V4.1 needs.
|
||||
|
||||
Do not confuse it with `src/flag_gems/fused/DSA/sparse_mla.py`
|
||||
(`triton_sparse_mla_fwd_interface`). That one hardcodes `is_causal=True` and requires the V3.2
|
||||
@ -39,8 +48,9 @@ same reason.
|
||||
|
||||
## Result
|
||||
|
||||
`ab_vs_tilelang.py`, 7 shapes, TileLang as reference: **all match** within 2 bf16 ULP, including
|
||||
`topk_length` and `attn_sink`.
|
||||
`ab_vendored.py`, 7 shapes, TileLang as reference: **all match** within 2 bf16 ULP, including
|
||||
`topk_length` and `attn_sink`. This runs the in-tree vendored file, so it is what the engine
|
||||
actually executes.
|
||||
|
||||
| case | max abs diff | tol |
|
||||
|---|---|---|
|
||||
@ -54,11 +64,68 @@ same reason.
|
||||
|
||||
## What is not done yet
|
||||
|
||||
* Not wired in. `sparse_mla_sm90.py` still imports TileLang; a Triton backend behind the same
|
||||
`patch_imports` hook is the next step.
|
||||
* Decode path untested end to end. Today decode gathers to bf16 and calls the kernel with identity
|
||||
indices and no causal; that is a different call shape than the ones above and needs its own test
|
||||
against the fp8-cache path.
|
||||
* The kernel returns `(out, max_logits, lse)` while DeepInfer's returns `(out, lse)`, and TileLang's
|
||||
`lse` is in log2 units where FlagGems' is natural log. A wrapper must drop `max_logits` and
|
||||
convert, or every downstream consumer reads a 1.44x-wrong scale.
|
||||
* Decode path end to end is covered by `ab_decode.py` and `test_wiring.py` at the entry-point
|
||||
level, but the engine's own decode was only exercised through the end-to-end runs below, not
|
||||
against a long-context trace.
|
||||
* Not token-identical to the TileLang backend. See below -- it runs and answers correctly, but the
|
||||
two backends do not produce the same tokens.
|
||||
* `fused_sample`'s rejection-sampling variant still needs TileLang, so the pure-Triton end-to-end run
|
||||
has to disable speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`). With MTP on, that one kernel
|
||||
is the remaining TileLang dependency.
|
||||
|
||||
---
|
||||
|
||||
## Wiring status
|
||||
|
||||
`h20_ops/sparse_mla.py` selects the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
|
||||
means TileLang if it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
|
||||
two entry points were the only call sites to touch. The chosen backend is logged once per rank.
|
||||
|
||||
### Verification, in the order it was done
|
||||
|
||||
Each step had to pass before the next was attempted, so a failure would point at one thing.
|
||||
|
||||
| step | what | result |
|
||||
|---|---|---|
|
||||
| 1 | `ab_vendored.py` — vendored kernel vs TileLang, 7 shapes | `ALL_MATCH`, ≤2 bf16 ULP |
|
||||
| 2 | `ab_decode.py` — decode call shape, fp8 cache, 5 cases incl. two-cache | `ALL_MATCH` |
|
||||
| 3 | `test_sparse_mla_triton.py` — adapter contract | `ALL_PASSED` (5 tests) |
|
||||
| 4 | `test_wiring.py` — selection + both entry points | `ALL_PASSED` (4 tests) |
|
||||
| 5 | end-to-end, TileLang blocked | `exit=0`, 0 fallbacks, correct answers |
|
||||
| 6 | end-to-end control, TileLang present | `exit=0`, 0 fallbacks, correct answers |
|
||||
|
||||
Two things the vendoring had to fix, both caught by step 1 rather than by luck:
|
||||
|
||||
* FlagGems 5.0.2 guards each top-k block by its **position** in the index array
|
||||
(`if ck * BK <= max_col`). Non-causal, `max_col` is `SKV-1`, so with `topk > SKV` every block past
|
||||
the first `SKV` entries is skipped. At SQ=16/SKV=128/topk=512 that is a 2.8e-2 error. Any prompt
|
||||
shorter than the 512-entry top-k hits it. Fixed locally, marked in the file; FlagGems master uses
|
||||
the data-dependent form.
|
||||
* The kernel's lse is a natural log; the TileLang one is log2. Passed through unconverted it would
|
||||
scale `_apply_sink` by 1.4427x -- plausible output, wrong numbers. `test_sparse_mla_triton.py`
|
||||
asserts the ratio is 1.0 rather than 0.693 so this cannot regress silently.
|
||||
|
||||
### End-to-end result
|
||||
|
||||
Same command, only the TileLang availability differs:
|
||||
|
||||
```
|
||||
DEEPINFER_SIMULATE_NO_TILELANG=1 DEEPINFER_KERNELS=ref DEEPINFER_OPS_STRICT=1 \
|
||||
DEEPINFER_USE_NEXTN_DECODE=0 bash scripts/h20_e2e_report.sh
|
||||
```
|
||||
|
||||
| | TileLang blocked | TileLang present |
|
||||
|---|---|---|
|
||||
| exit code | 0 | 0 |
|
||||
| sparse MLA backend (per rank) | `triton` x8 | `tilelang` x8 |
|
||||
| `ModuleNotFoundError` failures | 0 | 0 |
|
||||
| kernels served by the default backend | *(none)* | *(none)* |
|
||||
| DP ranks agreeing | 8/8 | 8/8 |
|
||||
| answer | 中国的首都→北京，日本的首都→东京 | 同上，另有 1949-09-27 定都的展开 |
|
||||
|
||||
**Not token-identical.** Comparing the two: prompt A shares 15 tokens then diverges, prompt B shares
|
||||
1 then diverges. A control run of the TileLang backend twice is bit-identical, so the harness is
|
||||
deterministic and the difference is real -- it comes from accumulation order, ≤2 ULP per layer
|
||||
compounding across 40+ layers, consistent with the 2.7% logits drift already recorded in `TODO.md`
|
||||
for the SM90 quantisation floor. It is not a wiring defect, and not something this kernel can fix on
|
||||
its own.
|
||||
|
||||
134
scripts/sparse_mla_port/ab_decode.py
Normal file
134
scripts/sparse_mla_port/ab_decode.py
Normal file
@ -0,0 +1,134 @@
|
||||
"""Decode-path A/B: does the vendored kernel serve the decode call shape at all?
|
||||
|
||||
Decode differs from prefill in a way that matters for this kernel. Prefill passes one KV buffer
|
||||
shared by every query and lets the indices pick per-query top-k. Decode gathers first, so each query
|
||||
owns a *different* KV window. The vendored kernel's `kv` operand is 3-D with no batch axis, so it
|
||||
cannot be handed a different KV per query.
|
||||
|
||||
It does not need to. Concatenate every query's window into one KV buffer and give query q the slice
|
||||
of index space `[q*W, q*W+W)`. That is one launch, and the buffer is the tensor decode already
|
||||
materialises -- no extra memory.
|
||||
|
||||
This test exercises that adapter against the current TileLang call, end to end from an fp8 paged
|
||||
cache in the V41 layout (512 fp8 values + 16 UE8M0 scale bytes per token), including the SWA + full
|
||||
two-cache concatenation that `decode_mla_mqa_sparse` supports.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.h20_ops.attn_sm90 import (
|
||||
CACHE_ELEMENTS,
|
||||
CACHE_ELEMENTS_PER_SF,
|
||||
UE8M0_BIAS,
|
||||
_gather_dequant,
|
||||
)
|
||||
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 = 64
|
||||
D = 512
|
||||
FAILED = []
|
||||
|
||||
|
||||
def make_cache(num_blocks: int, block_size: int, seed: int) -> torch.Tensor:
|
||||
"""A V41-layout fp8 cache: 512 e4m3 values then 16 UE8M0 exponent bytes, per token."""
|
||||
torch.manual_seed(seed)
|
||||
rows = num_blocks * block_size
|
||||
values = torch.randn((rows, CACHE_ELEMENTS), device=DEV, dtype=torch.float32) * 0.5
|
||||
fp8 = values.to(torch.float8_e4m3fn).view(torch.uint8)
|
||||
n_sf = CACHE_ELEMENTS // CACHE_ELEMENTS_PER_SF
|
||||
# Exponents in a narrow band around the bias so scales stay near 1.
|
||||
exponents = torch.randint(-3, 4, (rows, n_sf), device=DEV, dtype=torch.int32)
|
||||
sf = (exponents + UE8M0_BIAS).to(torch.uint8)
|
||||
packed = torch.cat([fp8, sf], dim=1)
|
||||
return packed.view(num_blocks, block_size, 1, -1)
|
||||
|
||||
|
||||
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 tl_decode(q4: torch.Tensor, kv4: torch.Tensor, scale: float) -> torch.Tensor:
|
||||
"""What `decode_mla_mqa_sparse` does today: one batch row per query, identity indices."""
|
||||
total_q, window, _, hd = kv4.shape
|
||||
identity = torch.arange(window, dtype=torch.int32, device=q4.device).view(1, 1, 1, -1)
|
||||
indices = identity.expand(total_q, 1, 1, window).contiguous()
|
||||
out, _ = sparse_mla_fwd_attn(q4, kv4, indices, sm_scale=scale, is_causal=False)
|
||||
return out.view(total_q, HQ, hd)
|
||||
|
||||
|
||||
def triton_decode(q4: torch.Tensor, kv4: torch.Tensor, scale: float) -> torch.Tensor:
|
||||
"""The proposed adapter: one shared KV, per-query offset index window."""
|
||||
total_q, window, _, hd = kv4.shape
|
||||
flat_kv = kv4.reshape(total_q * window, 1, hd).contiguous()
|
||||
offsets = torch.arange(window, dtype=torch.int32, device=q4.device).view(1, window)
|
||||
base = (torch.arange(total_q, dtype=torch.int32, device=q4.device) * window).view(total_q, 1)
|
||||
indices = (base + offsets).view(total_q, 1, window).contiguous()
|
||||
out, _, _ = flash_mla_sparse_fwd(q4.view(total_q, HQ, hd).contiguous(), flat_kv, indices, sm_scale=scale, d_v=hd)
|
||||
return out
|
||||
|
||||
|
||||
def run_case(tag: str, total_q: int, topk: int, seed: int, two_caches: bool) -> None:
|
||||
print(f"\n[{tag}] total_q={total_q} window={topk} two_caches={two_caches}")
|
||||
scale = D**-0.5
|
||||
torch.manual_seed(seed)
|
||||
q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
|
||||
cache = make_cache(8, 64, seed) # 512 rows
|
||||
n_rows = 8 * 64
|
||||
idx = torch.randint(0, n_rows, (total_q, topk), device=DEV, dtype=torch.int32)
|
||||
|
||||
gathered = _gather_dequant(cache, idx, None, D)
|
||||
if two_caches:
|
||||
cache2 = make_cache(8, 64, seed + 1000)
|
||||
idx2 = torch.randint(0, n_rows, (total_q, topk // 2), device=DEV, dtype=torch.int32)
|
||||
extra = _gather_dequant(cache2, idx2, None, D)
|
||||
gathered = torch.cat([gathered, extra], dim=1)
|
||||
|
||||
window = gathered.shape[1]
|
||||
kv4 = gathered.view(total_q, window, 1, D)
|
||||
q4 = q.view(total_q, 1, HQ, D)
|
||||
|
||||
tl_out = tl_decode(q4, kv4, scale)
|
||||
torch.cuda.synchronize()
|
||||
tri_out = triton_decode(q4, kv4, scale)
|
||||
torch.cuda.synchronize()
|
||||
compare("decode out", tl_out, tri_out)
|
||||
|
||||
# The gathered KV itself must be identical for the two paths, or the comparison would be
|
||||
# hiding a layout bug behind a numeric one.
|
||||
if not torch.equal(gathered.view(total_q, window, D), kv4.view(total_q, window, D)):
|
||||
print(" FAIL gathered kv changed shape semantics")
|
||||
FAILED.append("gather")
|
||||
|
||||
|
||||
CASES = [
|
||||
("decode, topk 512", 32, 512, 0, False),
|
||||
("decode, topk 256", 64, 256, 1, False),
|
||||
("decode, single query", 1, 512, 2, False),
|
||||
("decode, SWA + full (two caches)", 16, 512, 3, True),
|
||||
("decode, topk 128 (short window)", 128, 128, 4, False),
|
||||
]
|
||||
|
||||
if __name__ == "__main__":
|
||||
for c in CASES:
|
||||
run_case(*c)
|
||||
print()
|
||||
if FAILED:
|
||||
print(f"MISMATCH {len(FAILED)}: {FAILED}")
|
||||
sys.exit(1)
|
||||
print(f"ALL_MATCH ({len(CASES)} cases)")
|
||||
86
scripts/sparse_mla_port/ab_vendored.py
Normal file
86
scripts/sparse_mla_port/ab_vendored.py
Normal file
@ -0,0 +1,86 @@
|
||||
"""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)")
|
||||
@ -1,153 +0,0 @@
|
||||
"""FlagGems' `flash_mla_sparse_fwd` vs DeepInfer's TileLang sparse MLA, on the engine's own contract.
|
||||
|
||||
Why this pairing and not the DSA kernel: `attn_sm90.prefill_attn_varlen_sparse` calls the TileLang
|
||||
kernel with `is_causal=False`, documented in the source -- the indexer only scores positions before
|
||||
each query, so the indices carry causality, and bounding on the row index instead collapses the
|
||||
layer to NaN (that bug already happened once). `flash_mla_sparse_fwd` has no causal parameter for
|
||||
the same reason. The DSA kernel (`triton_sparse_mla_fwd_interface`) hardcodes causal and needs the
|
||||
V3.2 576-wide split layout, which V4.1 does not have, so it is not a candidate here.
|
||||
|
||||
Argument correspondence, which is closer than expected:
|
||||
|
||||
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=...)
|
||||
|
||||
Both return (out, ...). `index_length` is a per-query valid count applied as -1 padding upstream;
|
||||
`topk_length` is the same thing applied inside the kernel.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
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
|
||||
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
|
||||
|
||||
DEV = "cuda"
|
||||
HQ, HKV, D, TOPK = 64, 1, 512, 512
|
||||
FAILED = []
|
||||
|
||||
|
||||
def tl_call(q3, kv3, idx4, scale):
|
||||
"""DeepInfer's shape convention: (S,H,D)/(S,1,D)/(S,1,topk) -> pad a batch axis."""
|
||||
s_q = q3.shape[0]
|
||||
s_kv = kv3.shape[0]
|
||||
out, lse = sparse_mla_fwd_attn(
|
||||
q3.view(1, s_q, HQ, D),
|
||||
kv3.view(1, s_kv, 1, D),
|
||||
idx4.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, a, b, extra=""):
|
||||
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
|
||||
finite = bool(torch.isfinite(bf).all())
|
||||
ok = diff.max().item() <= tol and finite
|
||||
print(
|
||||
f" {'ok ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f} max diff {diff.max().item():.3e}"
|
||||
f" mean {diff.mean().item():.3e} tol {tol:.3e} finite {finite}{extra}"
|
||||
)
|
||||
if not ok:
|
||||
FAILED.append(tag)
|
||||
|
||||
|
||||
def make(SQ, SKV, seed, distinct):
|
||||
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:
|
||||
rows = [torch.randperm(SKV, device=DEV)[:TOPK] for _ in range(SQ)]
|
||||
idx = torch.stack(rows).view(SQ, 1, TOPK).to(torch.int32)
|
||||
else:
|
||||
idx = torch.randint(0, SKV, (SQ, 1, TOPK), device=DEV, dtype=torch.int32)
|
||||
return q, kv, idx.contiguous()
|
||||
|
||||
|
||||
def case(tag, SQ, SKV, seed, distinct=True, use_length=False, use_sink=False):
|
||||
q, kv, idx = make(SQ, SKV, seed, distinct)
|
||||
scale = D**-0.5
|
||||
print(f"\n[{tag}] SQ={SQ} SKV={SKV} topk={TOPK} distinct={distinct}"
|
||||
f"{' +topk_length' if use_length else ''}{' +attn_sink' if use_sink else ''}")
|
||||
|
||||
# DeepInfer applies `index_length` by masking the tail of the index row to -1 upstream
|
||||
# (`_mask_beyond_length`), so the reference call is the same kernel with a padded index.
|
||||
idx_ref = idx
|
||||
tl_len = None
|
||||
if use_length:
|
||||
torch.manual_seed(seed + 100)
|
||||
tl_len = torch.randint(64, TOPK, (SQ,), device=DEV, dtype=torch.int32)
|
||||
idx_ref = idx.clone()
|
||||
ar = torch.arange(TOPK, device=DEV).view(1, 1, TOPK)
|
||||
idx_ref[(ar >= tl_len.view(SQ, 1, 1)).expand(SQ, 1, TOPK)] = -1
|
||||
idx_ref = idx_ref.contiguous()
|
||||
|
||||
tl_out, _ = tl_call(q, kv, idx_ref, scale)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
sink = None
|
||||
if use_sink:
|
||||
torch.manual_seed(seed + 200)
|
||||
sink = (torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5)
|
||||
# DeepInfer's sink is folded into the output after the kernel; feed the kernel no sink and
|
||||
# compare the raw attention, then check the sink formula separately below.
|
||||
fg_out, fg_max, fg_lse = flash_mla_sparse_fwd(
|
||||
q, kv, idx_ref, sm_scale=scale, d_v=D,
|
||||
attn_sink=sink,
|
||||
topk_length=tl_len if use_length else None,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if use_sink:
|
||||
# Compare against the reference with the sink applied the same way the kernel documents:
|
||||
# out * exp(lse) / (exp(lse) + exp(sink)).
|
||||
tl_lse = sparse_mla_fwd_attn(
|
||||
q.view(1, SQ, HQ, D), kv.view(1, SKV, 1, D), idx_ref.view(1, SQ, 1, -1),
|
||||
sm_scale=scale, is_causal=False,
|
||||
)[1].view(SQ, HQ)
|
||||
import math
|
||||
w = torch.exp(tl_lse.float() * math.log(2)) # TileLang lse is log2 units
|
||||
# tl_out is (SQ, HQ, D), so the weight must gain a trailing axis to broadcast over D.
|
||||
s = sink.view(1, HQ, 1)
|
||||
gate = w.unsqueeze(-1) / (w.unsqueeze(-1) + torch.exp(s))
|
||||
applied = tl_out.float() * gate
|
||||
compare("sink applied", applied, fg_out)
|
||||
else:
|
||||
compare("out", tl_out, fg_out)
|
||||
|
||||
|
||||
CASES = [
|
||||
("prefill, distinct top-k", dict(SQ=64, SKV=2048, seed=0)),
|
||||
("prefill, duplicate top-k", dict(SQ=32, SKV=512, seed=1, distinct=False)),
|
||||
("decode (SQ=1)", dict(SQ=1, SKV=4096, seed=2)),
|
||||
("odd SQ", dict(SQ=7, SKV=1024, seed=3)),
|
||||
("short KV (< topk)", dict(SQ=16, SKV=128, seed=4, distinct=False)),
|
||||
("with topk_length", dict(SQ=32, SKV=2048, seed=5, use_length=True)),
|
||||
("with attn_sink", dict(SQ=16, SKV=1024, seed=6, use_sink=True)),
|
||||
]
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run from the repo root: PYTHONPATH=. python scripts/sparse_mla_port/ab_vs_tilelang.py
|
||||
for tag, kw in CASES:
|
||||
case(tag, **kw)
|
||||
print()
|
||||
if FAILED:
|
||||
print(f"MISMATCH {len(FAILED)}: {FAILED}")
|
||||
sys.exit(1)
|
||||
print(f"ALL_MATCH ({len(CASES)} cases)")
|
||||
@ -1,55 +0,0 @@
|
||||
"""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)
|
||||
201
scripts/sparse_mla_port/test_sparse_mla_triton.py
Normal file
201
scripts/sparse_mla_port/test_sparse_mla_triton.py
Normal file
@ -0,0 +1,201 @@
|
||||
"""Unit tests for the Triton sparse-MLA adapter, against the TileLang kernel as reference.
|
||||
|
||||
The adapter's job is not the attention maths -- that is the vendored kernel's, and
|
||||
`ab_vendored.py` covers it. The job is the *contract*: return arity, lse units, dtype, and the two
|
||||
calling conventions `attn_sm90` uses. Those are what this file tests, because each of them fails
|
||||
silently if wrong:
|
||||
|
||||
* returning natural-log lse makes `_apply_sink` scale outputs by exp(lse*(1-ln2)) -- plausible
|
||||
output, wrong numbers;
|
||||
* returning (out, max_logits, lse) instead of (out, lse) makes every `out, lse = ...` unpack the
|
||||
wrong tensor;
|
||||
* mismapping the per-query convention attends the wrong query's window.
|
||||
|
||||
Run: PYTHONPATH=. python scripts/sparse_mla_port/test_sparse_mla_triton.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.h20_ops import sparse_mla_triton
|
||||
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn as tl_sparse_mla
|
||||
|
||||
DEV = "cuda"
|
||||
HQ = 64
|
||||
D = 512
|
||||
FAILED: list[str] = []
|
||||
|
||||
|
||||
def check(tag: str, cond: bool, detail: str = "") -> None:
|
||||
print(f" {'ok ' if cond else 'FAIL'} {tag}{' ' + detail if detail else ''}")
|
||||
if not cond:
|
||||
FAILED.append(tag)
|
||||
|
||||
|
||||
def close(tag: str, a: torch.Tensor, b: torch.Tensor) -> None:
|
||||
af, bf = a.float(), b.float()
|
||||
if af.shape != bf.shape:
|
||||
check(tag, False, f"shape {tuple(af.shape)} vs {tuple(bf.shape)}")
|
||||
return
|
||||
diff = (af - bf).abs().max().item()
|
||||
scale = af.abs().max().item()
|
||||
tol = 2.0**-8 * max(scale, 1e-9) * 2
|
||||
check(tag, diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}")
|
||||
|
||||
|
||||
def test_prefill_contract() -> None:
|
||||
"""Shared-KV convention: the shape prefill actually uses."""
|
||||
print("\n[prefill: shared KV, batch=1]")
|
||||
torch.manual_seed(0)
|
||||
sq, skv, topk = 64, 2048, 512
|
||||
q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
idx = torch.stack([torch.randperm(skv, device=DEV)[:topk] for _ in range(sq)])
|
||||
idx = idx.view(1, sq, 1, topk).to(torch.int32).contiguous()
|
||||
scale = D**-0.5
|
||||
|
||||
tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=scale, is_causal=False)
|
||||
tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=scale, is_causal=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
check("returns exactly two tensors", isinstance((tri_out, tri_lse), tuple) and len((tri_out, tri_lse)) == 2)
|
||||
check("out dtype matches reference", tri_out.dtype == tl_out.dtype, str(tri_out.dtype))
|
||||
check("lse dtype matches reference", tri_lse.dtype == tl_lse.dtype, str(tri_lse.dtype))
|
||||
check("out shape preserved", tri_out.shape == q.shape, str(tuple(tri_out.shape)))
|
||||
check("lse shape", tri_lse.shape == (1, sq, HQ), str(tuple(tri_lse.shape)))
|
||||
close("out", tl_out, tri_out)
|
||||
close("lse (log2)", tl_lse, tri_lse)
|
||||
|
||||
|
||||
def test_lse_units() -> None:
|
||||
"""lse must be log2, not the kernel's natural log.
|
||||
|
||||
Asserted independently of the out comparison: an adapter that forgot the conversion would still
|
||||
produce the right `out` here (the kernel applies no sink), so comparing out alone cannot catch
|
||||
it. The ratio between the two units is exactly 1/ln2, so it is checkable without a reference.
|
||||
"""
|
||||
print("\n[lse units: log2, not natural log]")
|
||||
torch.manual_seed(1)
|
||||
sq, skv, topk = 16, 1024, 512
|
||||
q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
idx = torch.randint(0, skv, (1, sq, 1, topk), device=DEV, dtype=torch.int32).contiguous()
|
||||
|
||||
_tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
_tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# If the adapter had passed the natural log through, tri_lse would equal tl_lse * ln2 instead
|
||||
# of tl_lse, i.e. be smaller by a factor of 1.4427. Check the ratio is ~1, not ~0.693.
|
||||
ratio = (tri_lse.float() / tl_lse.float().clamp_min(1e-6)).median().item()
|
||||
check("lse ratio to reference is ~1.0 (would be 0.693 if unconverted)", abs(ratio - 1.0) < 0.02, f"ratio {ratio:.4f}")
|
||||
# And the natural-log interpretation must be measurably different, or the test is vacuous.
|
||||
natural_ratio = (tl_lse.float() * math.log(2) / tl_lse.float().clamp_min(1e-6)).median().item()
|
||||
check("the two units are distinguishable for this input", abs(natural_ratio - 1.0) > 0.2, f"ln/log2 = {natural_ratio:.3f}")
|
||||
|
||||
|
||||
def test_decode_convention() -> None:
|
||||
"""Per-query KV convention: what decode passes (batch axis is the query axis)."""
|
||||
print("\n[decode: per-query KV, batch=total_q]")
|
||||
torch.manual_seed(2)
|
||||
total_q, window = 32, 512
|
||||
q = (torch.randn((total_q, 1, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
# Every query gets a different window, so a mis-mapped convention shows up as wrong numbers.
|
||||
kv = (torch.randn((total_q, window, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
identity = torch.arange(window, dtype=torch.int32, device=DEV).view(1, 1, 1, -1)
|
||||
idx = identity.expand(total_q, 1, 1, window).contiguous()
|
||||
|
||||
tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
close("decode out", tl_out.view(total_q, HQ, D), tri_out.view(total_q, HQ, D))
|
||||
close("decode lse", tl_lse.view(total_q, HQ), tri_lse.view(total_q, HQ))
|
||||
|
||||
|
||||
def test_rejects_unsupported() -> None:
|
||||
"""Loud failure beats a silently wrong result."""
|
||||
print("\n[rejects what it cannot do]")
|
||||
q = torch.zeros((1, 4, HQ, D), device=DEV, dtype=torch.bfloat16)
|
||||
kv = torch.zeros((1, 8, 1, D), device=DEV, dtype=torch.bfloat16)
|
||||
idx = torch.zeros((1, 4, 1, 8), device=DEV, dtype=torch.int32)
|
||||
|
||||
try:
|
||||
sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, is_causal=True)
|
||||
check("is_causal=True raises", False)
|
||||
except NotImplementedError as exc:
|
||||
check("is_causal=True raises", "is_causal=False" in str(exc))
|
||||
|
||||
q2 = torch.zeros((1, 4, 32, D), device=DEV, dtype=torch.bfloat16)
|
||||
try:
|
||||
sparse_mla_triton.sparse_mla_fwd_attn(q2, kv, idx, is_causal=False)
|
||||
check("unsupported head count raises", False)
|
||||
except ValueError as exc:
|
||||
check("unsupported head count raises", "heads" in str(exc))
|
||||
|
||||
kv3 = torch.zeros((3, 8, 1, D), device=DEV, dtype=torch.bfloat16)
|
||||
try:
|
||||
sparse_mla_triton.sparse_mla_fwd_attn(q, kv3, idx, is_causal=False)
|
||||
check("mismatched kv batch raises", False)
|
||||
except ValueError as exc:
|
||||
check("mismatched kv batch raises", "batch axis" in str(exc))
|
||||
|
||||
|
||||
def test_match_sink_behaviour() -> None:
|
||||
"""The consumer of lse is `_apply_sink`; check it two ways rather than one.
|
||||
|
||||
Comparing only the final bf16 tensors is the wrong instrument. `_apply_sink` multiplies an
|
||||
already-bf16 `out` (itself differing by ~1.4 ULP) by a factor and rounds again, so the result can
|
||||
differ by a couple ULP even when both inputs are correct -- measured 1.83e-4 against 1.76e-4 for
|
||||
2 ULP of the input magnitude. That is a tolerance artifact, not a disagreement, so this asserts
|
||||
the thing that actually matters on its own: the multiplicative factor, in fp32, is a pure
|
||||
function of lse and must agree tightly. That is precisely the claim "lse units are right".
|
||||
"""
|
||||
print("\n[attn_sink path agrees]")
|
||||
from deepinfer.h20_ops.attn_sm90 import _apply_sink
|
||||
|
||||
torch.manual_seed(3)
|
||||
sq, skv, topk = 16, 1024, 512
|
||||
q = (torch.randn((1, sq, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
kv = (torch.randn((1, skv, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
idx = torch.randint(0, skv, (1, sq, 1, topk), device=DEV, dtype=torch.int32).contiguous()
|
||||
sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5
|
||||
|
||||
tl_out, tl_lse = tl_sparse_mla(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
tri_out, tri_lse = sparse_mla_triton.sparse_mla_fwd_attn(q, kv, idx, sm_scale=D**-0.5, is_causal=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
f_tl = torch.sigmoid(tl_lse.float() * math.log(2) - sink.view(1, 1, HQ)).unsqueeze(-1)
|
||||
f_tri = torch.sigmoid(tri_lse.float() * math.log(2) - sink.view(1, 1, HQ)).unsqueeze(-1)
|
||||
rel = (f_tl - f_tri).abs().max().item() / max(f_tl.abs().max().item(), 1e-9)
|
||||
check("sink factor relative diff < 1e-5", rel < 1e-5, f"rel {rel:.3e}")
|
||||
|
||||
# Composed path: two bf16 roundings, so 4 ULP of the input magnitude, not 2.
|
||||
a = _apply_sink(tl_out, tl_lse, sink).float()
|
||||
b = _apply_sink(tri_out, tri_lse, sink).float()
|
||||
diff = (a - b).abs().max().item()
|
||||
scale = tl_out.float().abs().max().item()
|
||||
tol = 4 * 2.0**-8 * max(scale, 1e-9)
|
||||
check("_apply_sink output within 4 ULP", diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}")
|
||||
|
||||
|
||||
TESTS = [
|
||||
test_prefill_contract,
|
||||
test_lse_units,
|
||||
test_decode_convention,
|
||||
test_rejects_unsupported,
|
||||
test_match_sink_behaviour,
|
||||
]
|
||||
|
||||
if __name__ == "__main__":
|
||||
for test in TESTS:
|
||||
test()
|
||||
print()
|
||||
if FAILED:
|
||||
print(f"FAILED {len(FAILED)}: {FAILED}")
|
||||
sys.exit(1)
|
||||
print(f"ALL_PASSED ({len(TESTS)} tests)")
|
||||
236
scripts/sparse_mla_port/test_wiring.py
Normal file
236
scripts/sparse_mla_port/test_wiring.py
Normal file
@ -0,0 +1,236 @@
|
||||
"""Unit tests for the backend selection and for `attn_sm90`'s two entry points.
|
||||
|
||||
Two different things are under test, and they fail differently:
|
||||
|
||||
* the *dispatcher* (`h20_ops.sparse_mla`) -- a wrong choice means the wrong kernel runs, which no
|
||||
numerical test downstream could attribute correctly;
|
||||
* the *entry points* -- both must give the same answer whichever backend is selected, including the
|
||||
work around the kernel (`_mask_beyond_length`, `_apply_sink`, the fp8 gather-and-dequant), not
|
||||
just the kernel call.
|
||||
|
||||
The "no TileLang installed" case runs in a subprocess: `tilelang` is importable here, and the thing
|
||||
being tested is what happens when it is not. `scripts/simulate_no_tilelang` provides that, and the
|
||||
engine's own end-to-end check uses the same trick.
|
||||
|
||||
Run: PYTHONPATH=. python scripts/sparse_mla_port/test_wiring.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from deepinfer.h20_ops import sparse_mla, sparse_mla_triton
|
||||
from deepinfer.h20_ops.attn_sm90 import (
|
||||
CACHE_ELEMENTS,
|
||||
CACHE_ELEMENTS_PER_SF,
|
||||
UE8M0_BIAS,
|
||||
decode_mla_mqa_sparse,
|
||||
prefill_attn_varlen_sparse,
|
||||
)
|
||||
|
||||
DEV = "cuda"
|
||||
HQ = 64
|
||||
D = 512
|
||||
FAILED: list[str] = []
|
||||
|
||||
|
||||
def check(tag: str, cond: bool, detail: str = "") -> None:
|
||||
print(f" {'ok ' if cond else 'FAIL'} {tag}{' ' + detail if detail else ''}")
|
||||
if not cond:
|
||||
FAILED.append(tag)
|
||||
|
||||
|
||||
def close(tag: str, a: torch.Tensor, b: torch.Tensor, ulp: float = 2.0) -> None:
|
||||
af, bf = a.float(), b.float()
|
||||
if af.shape != bf.shape:
|
||||
check(tag, False, f"shape {tuple(af.shape)} vs {tuple(bf.shape)}")
|
||||
return
|
||||
diff = (af - bf).abs().max().item()
|
||||
scale = af.abs().max().item()
|
||||
tol = ulp * 2.0**-8 * max(scale, 1e-9)
|
||||
check(tag, diff <= tol and bool(torch.isfinite(bf).all()), f"max diff {diff:.3e} tol {tol:.3e}")
|
||||
|
||||
|
||||
def make_cache(num_blocks: int, block_size: int, seed: int) -> torch.Tensor:
|
||||
"""A V41-layout fp8 cache: 512 e4m3 values then 16 UE8M0 exponent bytes, per token."""
|
||||
torch.manual_seed(seed)
|
||||
rows = num_blocks * block_size
|
||||
values = torch.randn((rows, CACHE_ELEMENTS), device=DEV, dtype=torch.float32) * 0.5
|
||||
fp8 = values.to(torch.float8_e4m3fn).view(torch.uint8)
|
||||
n_sf = CACHE_ELEMENTS // CACHE_ELEMENTS_PER_SF
|
||||
exponents = torch.randint(-3, 4, (rows, n_sf), device=DEV, dtype=torch.int32)
|
||||
packed = torch.cat([fp8, (exponents + UE8M0_BIAS).to(torch.uint8)], dim=1)
|
||||
return packed.view(num_blocks, block_size, 1, -1)
|
||||
|
||||
|
||||
def run_with_kernel(
|
||||
kernel_module: Any,
|
||||
entry: Callable[..., Any],
|
||||
args: tuple,
|
||||
kwargs: dict[str, Any],
|
||||
) -> Any:
|
||||
"""Run `entry` with its module-level kernel swapped to another backend.
|
||||
|
||||
`attn_sm90` binds the kernel by name at module scope, so replacing that global is how the other
|
||||
backend gets exercised without an environment change and a second process. Restored in a
|
||||
`finally` so a failure here cannot corrupt the following test.
|
||||
"""
|
||||
original = entry.__globals__["sparse_mla_fwd_attn"]
|
||||
entry.__globals__["sparse_mla_fwd_attn"] = kernel_module.sparse_mla_fwd_attn
|
||||
try:
|
||||
return entry(*args, **kwargs)
|
||||
finally:
|
||||
entry.__globals__["sparse_mla_fwd_attn"] = original
|
||||
|
||||
|
||||
def test_selection_logic() -> None:
|
||||
"""Explicit values win; a bad value is refused rather than defaulted."""
|
||||
print("\n[selection: explicit override]")
|
||||
|
||||
os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "triton"
|
||||
sparse_mla.reset()
|
||||
check("triton forced", sparse_mla.backend() == "triton")
|
||||
|
||||
os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "tilelang"
|
||||
sparse_mla.reset()
|
||||
check("tilelang forced", sparse_mla.backend() == "tilelang")
|
||||
|
||||
os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "nonsense"
|
||||
sparse_mla.reset()
|
||||
try:
|
||||
sparse_mla.backend()
|
||||
check("bad value raises", False)
|
||||
except RuntimeError as exc:
|
||||
check("bad value raises", "not recognised" in str(exc))
|
||||
|
||||
os.environ.pop("DEEPINFER_SPARSE_MLA_BACKEND", None)
|
||||
sparse_mla.reset()
|
||||
check("auto picks tilelang where it imports", sparse_mla.backend() == "tilelang")
|
||||
|
||||
|
||||
def test_selection_without_tilelang() -> None:
|
||||
"""Auto must fall back when TileLang is genuinely absent; forcing it must explain why."""
|
||||
print("\n[selection: TileLang unimportable]")
|
||||
code = """
|
||||
import json, os
|
||||
from deepinfer.h20_ops import sparse_mla
|
||||
out = {"backend": sparse_mla.backend()}
|
||||
try:
|
||||
os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "tilelang"
|
||||
sparse_mla.reset()
|
||||
sparse_mla.backend()
|
||||
out["forced"] = "no error"
|
||||
except RuntimeError as exc:
|
||||
out["forced"] = "raised"
|
||||
out["message_shown"] = "not importable" in str(exc)
|
||||
print(json.dumps(out))
|
||||
"""
|
||||
root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
env = dict(os.environ)
|
||||
env["PYTHONPATH"] = os.pathsep.join([os.path.join(root, "scripts", "simulate_no_tilelang"), root])
|
||||
env["DEEPINFER_SIMULATE_NO_TILELANG"] = "1"
|
||||
env.pop("DEEPINFER_SPARSE_MLA_BACKEND", None)
|
||||
proc = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=env,
|
||||
cwd=root,
|
||||
timeout=600,
|
||||
check=False,
|
||||
)
|
||||
payload = None
|
||||
for line in reversed(proc.stdout.splitlines()):
|
||||
if line.strip().startswith("{"):
|
||||
payload = json.loads(line)
|
||||
break
|
||||
if payload is None:
|
||||
tail = proc.stderr.strip().splitlines()[-1] if proc.stderr.strip() else "no output"
|
||||
check("subprocess ran", False, tail)
|
||||
return
|
||||
check("auto falls back to triton", payload["backend"] == "triton", payload["backend"])
|
||||
check("forcing tilelang raises", payload.get("forced") == "raised")
|
||||
check("message explains why", payload.get("message_shown", False))
|
||||
|
||||
|
||||
def test_prefill_entry_point() -> None:
|
||||
"""`prefill_attn_varlen_sparse`, both backends, with and without the optional arguments.
|
||||
|
||||
`index_length` and `attn_sink` are covered because they are the two places the entry point does
|
||||
work around the kernel -- `_mask_beyond_length` before it, `_apply_sink` after -- and a backend
|
||||
swap that got either wrong would look fine without them.
|
||||
"""
|
||||
print("\n[attn_sm90.prefill_attn_varlen_sparse]")
|
||||
torch.manual_seed(10)
|
||||
total_q, total_k, topk = 48, 1536, 512
|
||||
q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
kv = (torch.randn((total_k, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
idx = torch.stack([torch.randperm(total_k, device=DEV)[:topk] for _ in range(total_q)])
|
||||
idx = idx.to(torch.int32).contiguous()
|
||||
length = torch.randint(64, topk, (total_q,), device=DEV, dtype=torch.int32)
|
||||
sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5
|
||||
kwargs = {"attn_sink": sink, "index_length": length}
|
||||
args = (q, kv, idx, D**-0.5)
|
||||
|
||||
tl_out, tl_lse = prefill_attn_varlen_sparse(*args, **kwargs)
|
||||
tri_out, tri_lse = run_with_kernel(sparse_mla_triton, prefill_attn_varlen_sparse, args, kwargs)
|
||||
torch.cuda.synchronize()
|
||||
close("prefill out", tl_out, tri_out)
|
||||
close("prefill lse", tl_lse, tri_lse)
|
||||
|
||||
tl_out2, tl_lse2 = prefill_attn_varlen_sparse(*args)
|
||||
tri_out2, tri_lse2 = run_with_kernel(sparse_mla_triton, prefill_attn_varlen_sparse, args, {})
|
||||
torch.cuda.synchronize()
|
||||
close("prefill out (no sink/length)", tl_out2, tri_out2)
|
||||
close("prefill lse (no sink/length)", tl_lse2, tri_lse2)
|
||||
|
||||
|
||||
def test_decode_entry_point() -> None:
|
||||
"""`decode_mla_mqa_sparse`, both backends, single cache and SWA + full two-cache."""
|
||||
print("\n[attn_sm90.decode_mla_mqa_sparse]")
|
||||
torch.manual_seed(11)
|
||||
total_q, topk = 24, 512
|
||||
q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
|
||||
cache = make_cache(8, 64, 11)
|
||||
idx = torch.randint(0, 512, (total_q, topk), device=DEV, dtype=torch.int32)
|
||||
sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5
|
||||
args = (q, cache, D, idx, None, D**-0.5)
|
||||
|
||||
tl_out, tl_lse = decode_mla_mqa_sparse(*args, attn_sink=sink)
|
||||
tri_out, tri_lse = run_with_kernel(sparse_mla_triton, decode_mla_mqa_sparse, args, {"attn_sink": sink})
|
||||
torch.cuda.synchronize()
|
||||
close("decode out", tl_out, tri_out)
|
||||
close("decode lse", tl_lse, tri_lse)
|
||||
|
||||
cache2 = make_cache(8, 64, 111)
|
||||
idx2 = torch.randint(0, 512, (total_q, topk // 2), device=DEV, dtype=torch.int32)
|
||||
kwargs2 = {"attn_sink": sink, "key_cache_extra": cache2, "indices_in_kvcache_extra": idx2}
|
||||
tl_out2, tl_lse2 = decode_mla_mqa_sparse(*args, **kwargs2)
|
||||
tri_out2, tri_lse2 = run_with_kernel(sparse_mla_triton, decode_mla_mqa_sparse, args, kwargs2)
|
||||
torch.cuda.synchronize()
|
||||
close("decode two-cache out", tl_out2, tri_out2)
|
||||
close("decode two-cache lse", tl_lse2, tri_lse2)
|
||||
|
||||
|
||||
TESTS = [
|
||||
test_selection_logic,
|
||||
test_selection_without_tilelang,
|
||||
test_prefill_entry_point,
|
||||
test_decode_entry_point,
|
||||
]
|
||||
|
||||
if __name__ == "__main__":
|
||||
for test in TESTS:
|
||||
test()
|
||||
print()
|
||||
if FAILED:
|
||||
print(f"FAILED {len(FAILED)}: {FAILED}")
|
||||
sys.exit(1)
|
||||
print(f"ALL_PASSED ({len(TESTS)} tests)")
|
||||
1206
scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py
vendored
1206
scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py
vendored
File diff suppressed because it is too large
Load Diff
@ -1 +0,0 @@
|
||||
"""Minimal stand-in so the FlagGems fused kernels can be imported standalone."""
|
||||
@ -1,27 +0,0 @@
|
||||
"""Stand-in for FlagGems' version probe.
|
||||
|
||||
The real module gates TLE (Triton Language Extensions) on the triton version. Here it reports
|
||||
whatever the installed triton actually supports, because the point is running the kernel, not
|
||||
reproducing FlagGems' packaging.
|
||||
"""
|
||||
|
||||
import triton
|
||||
|
||||
|
||||
def _version_tuple() -> tuple[int, ...]:
|
||||
parts = []
|
||||
for piece in triton.__version__.split(".")[:3]:
|
||||
digits = "".join(c for c in piece if c.isdigit())
|
||||
parts.append(int(digits) if digits else 0)
|
||||
return tuple(parts)
|
||||
|
||||
|
||||
def has_triton_tle(*minimum: int) -> bool:
|
||||
try:
|
||||
import triton.experimental.tle.language # noqa: F401
|
||||
except Exception:
|
||||
return False
|
||||
return _version_tuple() >= tuple(minimum)
|
||||
|
||||
|
||||
HAS_TLE = has_triton_tle(3, 6, 0)
|
||||
Loading…
x
Reference in New Issue
Block a user