From c6a65fe6addefd691cc417abac6781453a5356ee Mon Sep 17 00:00:00 2001 From: loopyt Date: Sun, 25 Oct 2026 08:52:03 +0800 Subject: [PATCH] feat(h20): FlagGems' Triton sparse MLA matches the TileLang kernel at V4.1's shape 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. --- TODO.md | 42 +- scripts/sparse_mla_port/README.md | 64 + scripts/sparse_mla_port/ab_vs_tilelang.py | 153 +++ scripts/sparse_mla_port/probe.py | 55 + .../vendor/fg_flashmla_sparse.py | 1206 +++++++++++++++++ .../vendor/flag_gems/__init__.py | 1 + .../vendor/flag_gems/utils/__init__.py | 0 .../flag_gems/utils/triton_version_utils.py | 27 + 8 files changed, 1543 insertions(+), 5 deletions(-) create mode 100644 scripts/sparse_mla_port/README.md create mode 100644 scripts/sparse_mla_port/ab_vs_tilelang.py create mode 100644 scripts/sparse_mla_port/probe.py create mode 100644 scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py create mode 100644 scripts/sparse_mla_port/vendor/flag_gems/__init__.py create mode 100644 scripts/sparse_mla_port/vendor/flag_gems/utils/__init__.py create mode 100644 scripts/sparse_mla_port/vendor/flag_gems/utils/triton_version_utils.py diff --git a/TODO.md b/TODO.md index bd6319b..4f5d8ae 100644 --- a/TODO.md +++ b/TODO.md @@ -630,13 +630,45 @@ quant/MoE/MHC/engram，25 处模块级 import / 23 个 API）。 4. prompt B 第 1 个 token 就分叉，暗示 prefill 阶段已存在实质差异，需要对比首个 token 的 logits 才能定位，而不是继续猜。 -### 6.6 还差什么才能真正脱离 tilelang +### 6.6 稀疏 MLA 的 Triton 方案已找到，且已数值验证（2026-10-25） + +原判断「SM90 稀疏 MLA 没有 Triton 版本、工作量远大于 compat 层」**需要修正**： +FlagGems 里已经有，而且是 drop-in 的。 + +**`flagos-ai/FlagGems` → `src/flag_gems/fused/flashmla_sparse.py` → `flash_mla_sparse_fwd`** +纯 Triton，实测在 H20 上编译执行通过，**7 组形状与现有 TileLang kernel 全部一致（≤2 bf16 ULP）**。 + +参数逐项对应，连两个可选参数都有： + +| DeepInfer `prefill_attn_varlen_sparse` | FlagGems `flash_mla_sparse_fwd` | +|---|---| +| `query [total_q, H, D]` | `q [s_q, h_q, d]` | +| `kv_a_with_mqa [total_k, 1, D]` | `kv [s_kv, h_kv, d_qk]` | +| `index_topk [total_q, topk]` | `indices [s_q, h_kv, topk]` | +| `softmax_scale` | `sm_scale` | +| `attn_sink [H]` | `attn_sink [h_q]` | +| `index_length [total_q]` | `topk_length [s_q]` | + +**语义也对得上**：`attn_sm90.py:110` 明确让 prefill 走 `is_causal=False`（索引由 indexer 保证 +causal，用行号当因果边界会导致整层 NaN，这个 bug 已经发生过一次）；FlagGems 这个 kernel 同样 +没有 causal 参数，理由一致。V4.1 的 `h_q=64 / h_kv=1 / d_qk=512` 恰好满足它所有 assert。 + +**不能用的那个**：同仓库 `DSA/sparse_mla.py`（`triton_sparse_mla_fwd_interface`）硬编码 +`is_causal=True`，且要求 V3.2 的 576 拆分布局（latent 512 + rope tail 64）。V4.1 把 RoPE 并进 +512 维内，`TD = 512-512 = 0`，`tl.arange(0,0)` 直接编译失败。别选错。 + +验证脚本与 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 倍。 -`h20_ops/sparse_mla_sm90.py` 是 `@tilelang.jit` 实现，且 `h20-sm90-triton` 分支上也是如此 -（该分支只是没屏蔽 tilelang，不是真的无依赖）。**SM90 稀疏 MLA 没有 Triton 版本**，这是 -「国产卡完全不用 tilelang」的最后一个硬缺口，工作量远大于这次的 compat 层。 `tilelang_ops/cuda/` 那 29 个模块已有 `tilelang_ops/triton/` 对应版本，靠 -`DEEPINFER_OPS_BACKEND=triton` 切换，不是缺口。 +`DEEPINFER_OPS_BACKEND=triton` 切换，本来就不是缺口。 ### 6.4 端到端曾经为什么跑不起来（环境 + 排查记录） diff --git a/scripts/sparse_mla_port/README.md b/scripts/sparse_mla_port/README.md new file mode 100644 index 0000000..81ba68e --- /dev/null +++ b/scripts/sparse_mla_port/README.md @@ -0,0 +1,64 @@ +# 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. + +## 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/`. + +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 +576-wide split layout (latent 512 + rope tail 64). V4.1 folds RoPE inside a single 512-wide +latent, so `TD = 512 - 512 = 0` and the kernel fails to compile on `tl.arange(0, 0)`. It is not a +candidate. + +## Why the pairing works + +V4.1's dims satisfy every assert in the FlagGems kernel: `h_q = 64`, `h_kv = 1`, `d_qk = 512`. + +The argument lists correspond item for item, including the two optional ones: + +| DeepInfer `prefill_attn_varlen_sparse` | FlagGems `flash_mla_sparse_fwd` | +|---|---| +| `query [total_q, H, D]` | `q [s_q, h_q, d]` | +| `kv_a_with_mqa [total_k, 1, D]` | `kv [s_kv, h_kv, d_qk]` | +| `index_topk [total_q, topk]` | `indices [s_q, h_kv, topk]` | +| `softmax_scale` | `sm_scale` | +| `attn_sink [H]` | `attn_sink [h_q]` | +| `index_length [total_q]` | `topk_length [s_q]` | + +Causality also lines up: `attn_sm90.prefill_attn_varlen_sparse` passes `is_causal=False` and +documents why -- 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 already happened +once, see the comment at `attn_sm90.py:110`). The FlagGems kernel has no causal parameter for the +same reason. + +## Result + +`ab_vs_tilelang.py`, 7 shapes, TileLang as reference: **all match** within 2 bf16 ULP, including +`topk_length` and `attn_sink`. + +| case | max abs diff | tol | +|---|---|---| +| prefill, distinct top-k (SQ=64, SKV=2048) | 1.22e-04 | 1.45e-04 | +| prefill, duplicate top-k | 1.22e-04 | 1.97e-04 | +| decode (SQ=1, SKV=4096) | 6.10e-05 | 1.20e-04 | +| odd SQ (7) | 6.10e-05 | 1.43e-04 | +| short KV < topk | 2.44e-04 | 2.90e-04 | +| with topk_length | 2.44e-04 | 3.20e-04 | +| with attn_sink | 1.10e-04 | 1.41e-04 | + +## 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. diff --git a/scripts/sparse_mla_port/ab_vs_tilelang.py b/scripts/sparse_mla_port/ab_vs_tilelang.py new file mode 100644 index 0000000..c704193 --- /dev/null +++ b/scripts/sparse_mla_port/ab_vs_tilelang.py @@ -0,0 +1,153 @@ +"""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)") diff --git a/scripts/sparse_mla_port/probe.py b/scripts/sparse_mla_port/probe.py new file mode 100644 index 0000000..b090a87 --- /dev/null +++ b/scripts/sparse_mla_port/probe.py @@ -0,0 +1,55 @@ +"""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) diff --git a/scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py b/scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py new file mode 100644 index 0000000..7b5e4e9 --- /dev/null +++ b/scripts/sparse_mla_port/vendor/fg_flashmla_sparse.py @@ -0,0 +1,1206 @@ +# 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. + +import os +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from flag_gems.utils.triton_version_utils import has_triton_tle + +if has_triton_tle(3, 6, 0): + try: + import triton.experimental.tle.language as tle + + HAS_TLE_FLASHMLA_SPARSE = True + except ImportError: + tle = None + HAS_TLE_FLASHMLA_SPARSE = False +else: + tle = None + HAS_TLE_FLASHMLA_SPARSE = False + + +TLE_FLASHMLA_PREFILL_BK = 64 +TLE_FLASHMLA_PREFILL_BH = 64 +TLE_FLASHMLA_PREFILL_PAIR_BLOCKS = 2 +TLE_FLASHMLA_PREFILL_WORKER_NUM_WARPS = 4 + + +@triton.autotune( + configs=[ + triton.Config({"BK": 64, "BH": 64}, num_warps=8, num_stages=2), + triton.Config({"BK": 64, "BH": 64}, num_warps=8, num_stages=4), + ], + key=["SQ", "HQ", "DQK", "SKV", "TOPK", "HAVE_ATTN_SINK", "HAVE_TOPK_LENGTH"], +) +@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_kvg, + stride_kvn, + stride_tg, + stride_tm, + stride_oh, + stride_om, + stride_mm, + stride_lm, + SQ, # s_q + HQ: tl.constexpr, # h_q=64 or 128 + DQK: tl.constexpr, # d_qk=512 or 576 + SKV, # s_kv + TOPK: tl.constexpr, # topk + HAVE_ATTN_SINK: tl.constexpr, + HAVE_TOPK_LENGTH: tl.constexpr, + BK: tl.constexpr, + BH: tl.constexpr, +): + num_head_blocks: tl.constexpr = (HQ + BH - 1) // BH + pid = tl.program_id(0) + i_sq = pid // num_head_blocks + i_sq = i_sq.to(tl.int64) # prevent mul overflow + i_gbh = pid % num_head_blocks + gbh_base = i_gbh * BH + DP: tl.constexpr = 512 + BDP: tl.constexpr = 256 + + q_base = q + i_sq * stride_qm + gbh_base * stride_qh + kv_base = kv + tkv_base = kv + DP + t_base = indices + i_sq * stride_tm + attn_sink_ptr = attn_sink + gbh_base if HAVE_ATTN_SINK else 0 + topk_length_ptr = topk_length + i_sq if HAVE_TOPK_LENGTH else 0 + o_base = output + i_sq * stride_om + gbh_base * stride_oh + max_log_base = max_logits + i_sq * stride_mm + gbh_base + l_base = lse + i_sq * stride_lm + gbh_base + + offs_h = tl.arange(0, BH) + offs_d = tl.arange(0, BDP) + if DQK == 576: + offs_td = tl.arange(0, 64) + offs_t = tl.arange(0, BK) + + # `[BH, 256] x 2` delivers better performance than `[BH, 512]` when BH=64 + q_ptr = q_base + offs_h[:, None] * stride_qh + offs_d[None, :] + q_blk0 = tl.load(q_ptr, eviction_policy="evict_first") + q_blk1 = tl.load(q_ptr + BDP, eviction_policy="evict_first") + if DQK == 576: + tq_ptr = q_base + DP + offs_h[:, None] * stride_qh + offs_td[None, :] + tq_blk = tl.load(tq_ptr, eviction_policy="evict_first") + + max_log = tl.full([BH], float("-inf"), dtype=tl.float32) + sum_exp = tl.full([BH], 0.0, dtype=tl.float32) + acc0 = tl.zeros([BH, BDP], dtype=tl.float32) + acc1 = tl.zeros([BH, BDP], dtype=tl.float32) + + topk_len = tl.load(topk_length_ptr) if HAVE_TOPK_LENGTH else TOPK + NK = tl.cdiv(topk_len, BK) + for ck in range(NK): + # step1: load indices + t_ptr = BK * ck + offs_t # [BK] + t_msk = t_ptr < topk_len + t_ptr += t_base + kv_ids = tl.load(t_ptr, t_msk, other=-1) + mask_ids = (kv_ids < SKV) & (kv_ids >= 0) + # filter invalid index that may cause overflow in mul + kv_ids = tl.where(mask_ids, kv_ids, 0) + + # step2: gather kv with indices + kv_ptr = kv_base + offs_d[:, None] + kv_ids[None, :] * stride_kvn + kv_blk0 = tl.load(kv_ptr, cache_modifier=".cg") # [BDP, BK] + kv_blk1 = tl.load(kv_ptr + BDP, cache_modifier=".cg") # [BDP, BK] + # step3: (q @ kv) * sm_scale + qk = tl.dot( + q_blk0, kv_blk0, out_dtype=tl.float32 + ) # [BH, BDP]@[BDP, BK] -> [BH, BK] + qk = tl.dot(q_blk1, kv_blk1, qk, out_dtype=tl.float32) + if DQK == 576: + tkv_ptr = tkv_base + offs_td[:, None] + kv_ids[None, :] * stride_kvn + tkv_blk = tl.load(tkv_ptr, cache_modifier=".cg") # [TDP, BK] + qk = tl.dot(tq_blk, tkv_blk, qk, out_dtype=tl.float32) + qk *= sm_scale + + # step4: preprocess for logsumexp + qk = tl.where(mask_ids[None, :], qk, float("-inf")) # [BH, BK] + # step5: lse=logsumexp(qk), loop part + new_max = tl.maximum(max_log, tl.max(qk, axis=1)) # [BH] + exp_qk = tl.math.exp(qk - new_max[:, None]) # [BH, BK] + sum_qk = tl.sum(exp_qk, axis=1) # [BH] + alpha = tl.math.exp(max_log - new_max) # [BH] + sum_exp = sum_exp * alpha + sum_qk # [BH] + # step6: exp(qk-lse) @ gathered_kv.trans(), loop part + acc0 = tl.dot( + exp_qk.to(tl.bfloat16), + kv_blk0.trans(), + acc0 * alpha[:, None], + out_dtype=tl.float32, + ) # [BH, BK]@[BK, BDP]->[BH, BDP] + acc1 = tl.dot( + exp_qk.to(tl.bfloat16), + kv_blk1.trans(), + acc1 * alpha[:, None], + out_dtype=tl.float32, + ) # [BH, BK]@[BK, BDP]->[BH, BDP] + max_log = new_max + + # step7: store max_logits + valid_mask = max_log != float("-inf") + max_log = tl.where(valid_mask, max_log, float("-inf")) + tl.store(max_log_base + offs_h, max_log) # [BH], float32 + + # step8: lse=logsumexp(qk) final part, store lse + orig_lse = max_log + tl.math.log(sum_exp) + lse_out = tl.where(valid_mask, orig_lse, float("inf")) + tl.store(l_base + offs_h, lse_out) # [BH], float32 + + # step9: exp(qk-lse) @ gathered_kv.trans(), final part + if HAVE_ATTN_SINK: + # step10: attn_sink + sink = tl.load(attn_sink_ptr + offs_h) # [BH] + sum_exp_new_lse = tl.math.exp(orig_lse) + tl.math.exp(sink) + factor = tl.math.exp(max_log) / sum_exp_new_lse + else: + factor = 1.0 / sum_exp + + out_vals0 = tl.where(valid_mask[:, None], acc0 * factor[:, None], 0.0) + out_vals1 = tl.where(valid_mask[:, None], acc1 * factor[:, None], 0.0) + # step11: store output + o_ptr = o_base + offs_h[:, None] * stride_oh + offs_d[None, :] # [BH, BDP] + tl.store(o_ptr, out_vals0.to(tl.bfloat16)) + tl.store(o_ptr + BDP, out_vals1.to(tl.bfloat16)) + + +if HAS_TLE_FLASHMLA_SPARSE: + + @triton.jit + def _tle_flashmla_prefill_producer( + k0_l_writer, + k0_r_writer, + k1_l_writer, + k1_r_writer, + valid_writer, + kv_base, + tkv_base, + t_base, + topk_len_ptr, + D: tl.constexpr, + TD: tl.constexpr, + DPH: tl.constexpr, + TDP: tl.constexpr, + VG: tl.constexpr, + SKV, + TOPK: tl.constexpr, + HAVE_TOPK_LENGTH: tl.constexpr, + HAVE_TAIL: tl.constexpr, + BK: tl.constexpr, + ): + topk_len = tl.load(topk_len_ptr) if HAVE_TOPK_LENGTH else TOPK + max_col = SKV - 1 + stride_kvn: tl.constexpr = VG * (TD + D) + NK = tl.cdiv(topk_len, BK) + NPAIRS = tl.cdiv(NK, 2) + offs_t = tl.arange(0, BK) + offs_tile = tl.arange(0, 64) + kv_tile_rows = tl.broadcast_to(offs_t[:, None], (BK, 64)) + for pair in tl.range(NPAIRS): + ck0 = pair * 2 + ck1 = ck0 + 1 + t_offs0 = BK * ck0 + offs_t + t_msk0 = t_offs0 < topk_len + kv_ids0 = tl.load(t_base + t_offs0, t_msk0, other=-1) + valid0 = t_msk0 & (kv_ids0 <= max_col) & (kv_ids0 >= 0) + kv_offsets0 = tl.where(valid0, kv_ids0, 0).to(tl.int64) * stride_kvn + + t_offs1 = BK * ck1 + offs_t + t_msk1 = t_offs1 < topk_len + kv_ids1 = tl.load(t_base + t_offs1, t_msk1, other=-1) + valid1 = t_msk1 & (kv_ids1 <= max_col) & (kv_ids1 >= 0) + kv_offsets1 = tl.where(valid1, kv_ids1, 0).to(tl.int64) * stride_kvn + + k0_l_slot = k0_l_writer.acquire(pair) + for tile in tl.static_range(0, DPH, 64): + k_cols = tile + offs_tile + k_cols_b = tl.broadcast_to(k_cols[None, :], (BK, 64)) + k0_l_ptr = kv_base + kv_offsets0[:, None] + k_cols[None, :] + k0_l_msk = valid0[:, None] & (k_cols < D)[None, :] + k0_l_blk = tl.load( + k0_l_ptr, + mask=k0_l_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k0_l_slot.sK, (kv_tile_rows, k_cols_b)), + k0_l_blk, + mask=k0_l_msk, + ) + k0_l_writer.commit(pair) + + k1_r_slot = k1_r_writer.acquire(pair) + for tile in tl.static_range(0, DPH, 64): + k_cols = DPH + tile + offs_tile + k_cols_b = tl.broadcast_to(k_cols[None, :], (BK, 64)) + k1_r_ptr = kv_base + kv_offsets1[:, None] + k_cols[None, :] + k1_r_msk = valid1[:, None] & (k_cols < D)[None, :] + k1_r_blk = tl.load( + k1_r_ptr, + mask=k1_r_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k1_r_slot.sK, (kv_tile_rows, k_cols_b)), + k1_r_blk, + mask=k1_r_msk, + ) + if HAVE_TAIL: + offs_td = tl.arange(0, TDP) + k1_r_tail_ptr = tkv_base + kv_offsets1[:, None] + offs_td[None, :] + k1_r_tail_msk = valid1[:, None] & (offs_td < TD)[None, :] + k1_r_tail_blk = tl.load( + k1_r_tail_ptr, + mask=k1_r_tail_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k1_r_slot.sK_tail), + k1_r_tail_blk, + mask=k1_r_tail_msk, + ) + k1_r_writer.commit(pair) + + k0_r_slot = k0_r_writer.acquire(pair) + for tile in tl.static_range(0, DPH, 64): + k_cols = DPH + tile + offs_tile + k_cols_b = tl.broadcast_to(k_cols[None, :], (BK, 64)) + k0_r_ptr = kv_base + kv_offsets0[:, None] + k_cols[None, :] + k0_r_msk = valid0[:, None] & (k_cols < D)[None, :] + k0_r_blk = tl.load( + k0_r_ptr, + mask=k0_r_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k0_r_slot.sK, (kv_tile_rows, k_cols_b)), + k0_r_blk, + mask=k0_r_msk, + ) + if HAVE_TAIL: + offs_td = tl.arange(0, TDP) + k0_r_tail_ptr = tkv_base + kv_offsets0[:, None] + offs_td[None, :] + k0_r_tail_msk = valid0[:, None] & (offs_td < TD)[None, :] + k0_r_tail_blk = tl.load( + k0_r_tail_ptr, + mask=k0_r_tail_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k0_r_slot.sK_tail), + k0_r_tail_blk, + mask=k0_r_tail_msk, + ) + k0_r_writer.commit(pair) + + k1_l_slot = k1_l_writer.acquire(pair) + for tile in tl.static_range(0, DPH, 64): + k_cols = tile + offs_tile + k_cols_b = tl.broadcast_to(k_cols[None, :], (BK, 64)) + k1_l_ptr = kv_base + kv_offsets1[:, None] + k_cols[None, :] + k1_l_msk = valid1[:, None] & (k_cols < D)[None, :] + k1_l_blk = tl.load( + k1_l_ptr, + mask=k1_l_msk, + other=0.0, + eviction_policy="evict_last", + ) + tl.store( + tle.gpu.local_ptr(k1_l_slot.sK, (kv_tile_rows, k_cols_b)), + k1_l_blk, + mask=k1_l_msk, + ) + k1_l_writer.commit(pair) + + valid_slot = valid_writer.acquire(pair) + valid_row0 = tl.full([BK], 0, dtype=tl.int32) + valid_row1 = tl.full([BK], 1, dtype=tl.int32) + valid_ptr0 = tle.gpu.local_ptr(valid_slot.is_kv_valid, (valid_row0, offs_t)) + valid_ptr1 = tle.gpu.local_ptr(valid_slot.is_kv_valid, (valid_row1, offs_t)) + tl.store(valid_ptr0, valid0.to(tl.int8)) + tl.store(valid_ptr1, valid1.to(tl.int8)) + valid_writer.commit(pair) + + @triton.jit + def _tle_flashmla_prefill_consumer0( + q_writer, + q_reader, + q_desc, + tq_desc, + k0_l_reader, + k0_r_qk_reader, + k1_l_remote_reader, + valid_reader, + sM_wg0_writer, + sM_wg1_reader, + sS0_writer, + sS1_reader, + sL_wg0_writer, + sL_wg1_reader, + output_desc, + output_row, + h_base, + topk_len_ptr, + attn_sink_base, + log_scale: tl.constexpr, + D: tl.constexpr, + TD: tl.constexpr, + OUT_DTYPE: tl.constexpr, + HAVE_ATTN_SINK: tl.constexpr, + TOPK: tl.constexpr, + HAVE_TOPK_LENGTH: tl.constexpr, + HAVE_TAIL: tl.constexpr, + BK: tl.constexpr, + BH: tl.constexpr, + DPH: tl.constexpr, + TDP: tl.constexpr, + G: tl.constexpr, + ): + topk_len = tl.load(topk_len_ptr) if HAVE_TOPK_LENGTH else TOPK + offs_h = tl.arange(0, BH) + offs_dh = tl.arange(0, DPH) + mask_h = h_base + offs_h < G + mask_od_l = offs_dh < D + kv_rows = tl.broadcast_to(tl.arange(0, BK)[:, None], (BK, DPH)) + kv_cols_l = tl.broadcast_to(offs_dh[None, :], (BK, DPH)) + kv_cols_r = tl.broadcast_to((DPH + offs_dh)[None, :], (BK, DPH)) + + q_write_slot = q_writer.acquire(0) + tle.gpu.copy(q_desc, q_write_slot.sQ_l, [BH, DPH], [output_row, 0]) + tle.gpu.copy(q_desc, q_write_slot.sQ_r, [BH, DPH], [output_row, DPH]) + if HAVE_TAIL: + tle.gpu.copy(tq_desc, q_write_slot.sQ_tail, [BH, TDP], [output_row, D]) + q_writer.commit(0) + + q_slot = q_reader.wait(0).slot + q_l_smem_ptr = tle.gpu.local_ptr(q_slot.sQ_l) + q_r_smem_ptr = tle.gpu.local_ptr(q_slot.sQ_r) + max_prev = tl.full([BH], -1.0e30, dtype=tl.float32) + sum_exp = tl.full([BH], 0.0, dtype=tl.float32) + acc_l = tl.zeros([BH, DPH], dtype=tl.float32) + + NK = tl.cdiv(topk_len, BK) + NPAIRS = tl.cdiv(NK, 2) + + for pair in tl.range(NPAIRS): + k0_l_wait = k0_l_reader.wait(pair) + k0_l_slot = k0_l_wait.slot + + q_l_blk = tl.load(q_l_smem_ptr) + q_r_blk = tl.load(q_r_smem_ptr) + k0_l_blk = tl.load(tle.gpu.local_ptr(k0_l_slot.sK, (kv_rows, kv_cols_l))) + + qk0 = tl.full([BH, BK], 0.0, dtype=tl.float32) + qk0 = tl.dot(q_l_blk, tl.trans(k0_l_blk), qk0, out_dtype=tl.float32) + + k0_r_wait = k0_r_qk_reader.wait(pair) + k0_r_slot = k0_r_wait.slot + k0_r_blk = tl.load(tle.gpu.local_ptr(k0_r_slot.sK, (kv_rows, kv_cols_r))) + qk0 = tl.dot(q_r_blk, tl.trans(k0_r_blk), qk0, out_dtype=tl.float32) + if HAVE_TAIL: + q_tail_blk = tl.load(tle.gpu.local_ptr(q_slot.sQ_tail)) + k0_t_blk = tl.load(tle.gpu.local_ptr(k0_r_slot.sK_tail)) + qk0 = tl.dot(q_tail_blk, tl.trans(k0_t_blk), qk0, out_dtype=tl.float32) + + valid_wait = valid_reader.wait(pair) + row0 = tl.full([BK], 0, dtype=tl.int32) + valid0 = ( + tl.load( + tle.gpu.local_ptr( + valid_wait.slot.is_kv_valid, (row0, tl.arange(0, BK)) + ) + ) + != 0 + ) + qk0 = tl.where(valid0[None, :], qk0, float("-inf")) + valid_reader.release(pair) + + local_max = tl.maximum(max_prev, tl.max(qk0, axis=1)) + alpha = tl.math.exp2((max_prev - local_max) * log_scale) + prob0 = tl.math.exp2(qk0 * log_scale - local_max[:, None] * log_scale) + sum_exp = sum_exp * alpha + tl.sum(prob0, axis=1) + acc_l = acc_l * alpha[:, None] + prob0_b = prob0.to(OUT_DTYPE) + + sM_wg0_slot = sM_wg0_writer.acquire(pair) + tl.store(tle.gpu.local_ptr(sM_wg0_slot.sM), local_max) + sM_wg0_writer.commit(pair) + + k0_l_blk = tl.load(tle.gpu.local_ptr(k0_l_slot.sK, (kv_rows, kv_cols_l))) + acc_l = tl.dot(prob0_b, k0_l_blk, acc_l, out_dtype=tl.float32) + k0_l_reader.release(pair) + k0_r_qk_reader.release(pair) + + sM_wg1_wait = sM_wg1_reader.wait(pair) + max_next = tl.load(tle.gpu.local_ptr(sM_wg1_wait.slot.sM)) + sM_wg1_reader.release(pair) + + final_scale = tl.math.exp2((local_max - max_next) * log_scale) + sum_exp = sum_exp * final_scale + acc_l = acc_l * final_scale[:, None] + + prob0_scaled = prob0 * final_scale[:, None] + sS0_slot = sS0_writer.acquire(pair) + tl.store(tle.gpu.local_ptr(sS0_slot.sS0), prob0_scaled.to(OUT_DTYPE)) + sS0_writer.commit(pair) + + sS1_wait = sS1_reader.wait(pair) + prob1 = tl.load(tle.gpu.local_ptr(sS1_wait.slot.sS1)) + k1_l_wait = k1_l_remote_reader.wait(pair) + k1_l_blk = tl.load( + tle.gpu.local_ptr(k1_l_wait.slot.sK, (kv_rows, kv_cols_l)) + ) + acc_l = tl.dot(prob1, k1_l_blk, acc_l, out_dtype=tl.float32) + sS1_reader.release(pair) + k1_l_remote_reader.release(pair) + + max_prev = max_next + + sL_wg0_slot = sL_wg0_writer.acquire(0) + tl.store(tle.gpu.local_ptr(sL_wg0_slot.sL), sum_exp) + sL_wg0_writer.commit(0) + sL_wg1_wait = sL_wg1_reader.wait(1) + peer_sum = tl.load(tle.gpu.local_ptr(sL_wg1_wait.slot.sL)) + total_sum = sum_exp + peer_sum + sL_wg1_reader.release(1) + + is_no_valid_tokens = total_sum == 0.0 + inv_total_sum = tl.fdiv(1.0, total_sum) + out_l_vals = acc_l * inv_total_sum[:, None] + if HAVE_ATTN_SINK: + fin_log = ( + max_prev * log_scale + tl.math.log2(total_sum) + ) * 0.6931471805599453 + sink = tl.load(attn_sink_base + h_base + offs_h, mask_h, other=0.0) + sink_scale = tl.fdiv(1.0, 1.0 + tl.math.exp(sink - fin_log)) + out_l_vals = out_l_vals * sink_scale[:, None] + out_l_vals = tl.where(is_no_valid_tokens[:, None], 0.0, out_l_vals) + o_l_msk = mask_h[:, None] & mask_od_l[None, :] + tl.store(q_l_smem_ptr, out_l_vals.to(OUT_DTYPE), o_l_msk) + tle.gpu.copy(q_slot.sQ_l, output_desc, [BH, DPH], [output_row, 0]) + + @triton.jit + def _tle_flashmla_prefill_consumer1( + q_reader, + k1_r_reader, + k1_l_qk_reader, + k0_r_remote_reader, + valid_reader, + sM_wg1_writer, + sM_wg0_reader, + sS1_writer, + sS0_reader, + sL_wg1_writer, + sL_wg0_reader, + final_max_logits_smem, + final_lse_smem, + output_desc, + output_row, + max_logits_base, + l_base, + h_base, + topk_len_ptr, + attn_sink_base, + log_scale: tl.constexpr, + D: tl.constexpr, + TD: tl.constexpr, + OUT_DTYPE: tl.constexpr, + HAVE_ATTN_SINK: tl.constexpr, + TOPK: tl.constexpr, + HAVE_TOPK_LENGTH: tl.constexpr, + HAVE_TAIL: tl.constexpr, + BK: tl.constexpr, + BH: tl.constexpr, + DPH: tl.constexpr, + TDP: tl.constexpr, + G: tl.constexpr, + ): + topk_len = tl.load(topk_len_ptr) if HAVE_TOPK_LENGTH else TOPK + offs_h = tl.arange(0, BH) + offs_dh = tl.arange(0, DPH) + mask_h = h_base + offs_h < G + mask_od_r = DPH + offs_dh < D + kv_rows = tl.broadcast_to(tl.arange(0, BK)[:, None], (BK, DPH)) + kv_cols_l = tl.broadcast_to(offs_dh[None, :], (BK, DPH)) + kv_cols_r = tl.broadcast_to((DPH + offs_dh)[None, :], (BK, DPH)) + q_slot = q_reader.wait(0).slot + q_l_smem_ptr = tle.gpu.local_ptr(q_slot.sQ_l) + q_r_smem_ptr = tle.gpu.local_ptr(q_slot.sQ_r) + max_prev = tl.full([BH], -1.0e30, dtype=tl.float32) + sum_exp = tl.full([BH], 0.0, dtype=tl.float32) + acc_r = tl.zeros([BH, DPH], dtype=tl.float32) + + NK = tl.cdiv(topk_len, BK) + NPAIRS = tl.cdiv(NK, 2) + for pair in tl.range(NPAIRS): + k1_r_wait = k1_r_reader.wait(pair) + k1_r_slot = k1_r_wait.slot + + q_l_blk = tl.load(q_l_smem_ptr) + q_r_blk = tl.load(q_r_smem_ptr) + k1_r_blk = tl.load(tle.gpu.local_ptr(k1_r_slot.sK, (kv_rows, kv_cols_r))) + + qk1 = tl.full([BH, BK], 0.0, dtype=tl.float32) + qk1 = tl.dot(q_r_blk, tl.trans(k1_r_blk), qk1, out_dtype=tl.float32) + if HAVE_TAIL: + q_tail_blk = tl.load(tle.gpu.local_ptr(q_slot.sQ_tail)) + k1_t_blk = tl.load(tle.gpu.local_ptr(k1_r_slot.sK_tail)) + qk1 = tl.dot(q_tail_blk, tl.trans(k1_t_blk), qk1, out_dtype=tl.float32) + k1_l_wait = k1_l_qk_reader.wait(pair) + k1_l_slot = k1_l_wait.slot + k1_l_blk = tl.load(tle.gpu.local_ptr(k1_l_slot.sK, (kv_rows, kv_cols_l))) + qk1 = tl.dot(q_l_blk, tl.trans(k1_l_blk), qk1, out_dtype=tl.float32) + + valid_wait = valid_reader.wait(pair) + row1 = tl.full([BK], 1, dtype=tl.int32) + valid1 = ( + tl.load( + tle.gpu.local_ptr( + valid_wait.slot.is_kv_valid, (row1, tl.arange(0, BK)) + ) + ) + != 0 + ) + qk1 = tl.where(valid1[None, :], qk1, float("-inf")) + valid_reader.release(pair) + + sM_wg0_wait = sM_wg0_reader.wait(pair) + candidate0 = tl.load(tle.gpu.local_ptr(sM_wg0_wait.slot.sM)) + sM_wg0_reader.release(pair) + + candidate1 = tl.maximum(max_prev, tl.max(qk1, axis=1)) + max_next = tl.maximum(candidate1, candidate0) + sM_wg1_slot = sM_wg1_writer.acquire(pair) + tl.store(tle.gpu.local_ptr(sM_wg1_slot.sM), max_next) + sM_wg1_writer.commit(pair) + + alpha = tl.math.exp2((max_prev - max_next) * log_scale) + prob1 = tl.math.exp2(qk1 * log_scale - max_next[:, None] * log_scale) + sum_exp = sum_exp * alpha + tl.sum(prob1, axis=1) + acc_r = acc_r * alpha[:, None] + prob1_b = prob1.to(OUT_DTYPE) + + k1_l_qk_reader.release(pair) + + acc_r = tl.dot(prob1_b, k1_r_blk, acc_r, out_dtype=tl.float32) + + sS1_slot = sS1_writer.acquire(pair) + tl.store(tle.gpu.local_ptr(sS1_slot.sS1), prob1_b) + sS1_writer.commit(pair) + + sS0_wait = sS0_reader.wait(pair) + prob0 = tl.load(tle.gpu.local_ptr(sS0_wait.slot.sS0)) + k0_r_wait = k0_r_remote_reader.wait(pair) + k0_r_blk = tl.load( + tle.gpu.local_ptr(k0_r_wait.slot.sK, (kv_rows, kv_cols_r)) + ) + acc_r = tl.dot(prob0, k0_r_blk, acc_r, out_dtype=tl.float32) + k1_r_reader.release(pair) + sS0_reader.release(pair) + k0_r_remote_reader.release(pair) + max_prev = max_next + + sL_wg1_slot = sL_wg1_writer.acquire(1) + tl.store(tle.gpu.local_ptr(sL_wg1_slot.sL), sum_exp) + sL_wg1_writer.commit(1) + sL_wg0_wait = sL_wg0_reader.wait(0) + peer_sum = tl.load(tle.gpu.local_ptr(sL_wg0_wait.slot.sL)) + total_sum = sum_exp + peer_sum + sL_wg0_reader.release(0) + + is_no_valid_tokens = total_sum == 0.0 + inv_total_sum = tl.fdiv(1.0, total_sum) + out_r_vals = acc_r * inv_total_sum[:, None] + final_max_logits_log2 = max_prev * log_scale + final_max_logits = final_max_logits_log2 * 0.6931471805599453 + fin_log = (final_max_logits_log2 + tl.math.log2(total_sum)) * 0.6931471805599453 + if HAVE_ATTN_SINK: + sink = tl.load(attn_sink_base + h_base + offs_h, mask_h, other=0.0) + sink_scale = tl.fdiv(1.0, 1.0 + tl.math.exp(sink - fin_log)) + out_r_vals = out_r_vals * sink_scale[:, None] + out_r_vals = tl.where(is_no_valid_tokens[:, None], 0.0, out_r_vals) + o_r_msk = mask_h[:, None] & mask_od_r[None, :] + tl.store(q_r_smem_ptr, out_r_vals.to(OUT_DTYPE), o_r_msk) + tle.gpu.copy(q_slot.sQ_r, output_desc, [BH, DPH], [output_row, DPH]) + + final_max_logits = tl.where(is_no_valid_tokens, float("-inf"), final_max_logits) + fin_log = tl.where(is_no_valid_tokens, float("inf"), fin_log) + tl.store(tle.gpu.local_ptr(final_max_logits_smem), final_max_logits, mask_h) + tl.store(tle.gpu.local_ptr(final_lse_smem), fin_log, mask_h) + final_max_logits = tl.load( + tle.gpu.local_ptr(final_max_logits_smem), mask_h, other=float("-inf") + ) + fin_log = tl.load(tle.gpu.local_ptr(final_lse_smem), mask_h, other=float("inf")) + tl.store(max_logits_base + offs_h, final_max_logits, mask_h) + tl.store(l_base + offs_h, fin_log, mask_h) + + @triton.jit + def _tle_flashmla_prefill_fwd( + q_desc, + tq_desc, + output_desc, + kv, + indices, + attn_sink, + topk_length, + sm_scale: tl.constexpr, + output, + max_logits, + lse, + SQ, + H: tl.constexpr, + DQK: tl.constexpr, + SKV, + TOPK: tl.constexpr, + HAVE_ATTN_SINK: tl.constexpr, + HAVE_TOPK_LENGTH: tl.constexpr, + D: tl.constexpr, + TD: tl.constexpr, + DP: tl.constexpr, + TDP: tl.constexpr, + G: tl.constexpr, + VG: tl.constexpr, + RH: tl.constexpr, + HAVE_TAIL: tl.constexpr, + BK: tl.constexpr, + BH: tl.constexpr, + PAIR_BLOCKS: tl.constexpr, + ): + DPH: tl.constexpr = DP // 2 + stride_kvg: tl.constexpr = TD + D + stride_tg = TOPK + stride_tm = VG * stride_tg + stride_lm = H + stride_mm = H + + pid = tl.program_id(0) + programs_per_q: tl.constexpr = VG * RH + i_sq = pid // programs_per_q + i_grh = pid % programs_per_q + i_g = i_grh // RH + i_rh = i_grh % RH + h_base = i_rh * BH + q_head_base = i_g * G + h_base + i_sq64 = i_sq.to(tl.int64) + i_g64 = i_g.to(tl.int64) + q_head_base64 = q_head_base.to(tl.int64) + kv_base = kv + i_g64 * stride_kvg + tkv_base = kv_base + D + t_base = indices + i_sq64 * stride_tm + i_g64 * stride_tg + topk_len_ptr = topk_length + i_sq64 if HAVE_TOPK_LENGTH else indices + attn_sink_base = attn_sink if HAVE_ATTN_SINK else max_logits + max_logits_base = max_logits + i_sq64 * stride_mm + q_head_base64 + l_base = lse + i_sq64 * stride_lm + q_head_base64 + q_row = i_sq * H + q_head_base + _ = output + _ = SQ + _ = DQK + + sQ_l_smem = tle.gpu.alloc( + [1, BH, DPH], dtype=kv.dtype.element_ty, layout=None, scope=tle.gpu.smem + ) + sQ_r_smem = tle.gpu.alloc( + [1, BH, DPH], dtype=kv.dtype.element_ty, layout=None, scope=tle.gpu.smem + ) + if HAVE_TAIL: + sQ_tail_smem = tle.gpu.alloc( + [1, BH, TDP], + dtype=kv.dtype.element_ty, + layout=None, + scope=tle.gpu.smem, + ) + q_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sQ", + readers=("wg0", "wg1"), + one_shot=True, + sQ_l=sQ_l_smem, + sQ_r=sQ_r_smem, + sQ_tail=sQ_tail_smem, + ) + else: + q_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sQ", + readers=("wg0", "wg1"), + one_shot=True, + sQ_l=sQ_l_smem, + sQ_r=sQ_r_smem, + ) + + sK0_smem = tle.gpu.alloc( + [1, BK, DP], dtype=kv.dtype.element_ty, layout=None, scope=tle.gpu.smem + ) + sK1_smem = tle.gpu.alloc( + [1, BK, DP], dtype=kv.dtype.element_ty, layout=None, scope=tle.gpu.smem + ) + if HAVE_TAIL: + sK0_tail_smem = tle.gpu.alloc( + [1, BK, TDP], + dtype=kv.dtype.element_ty, + layout=None, + scope=tle.gpu.smem, + ) + sK1_tail_smem = tle.gpu.alloc( + [1, BK, TDP], + dtype=kv.dtype.element_ty, + layout=None, + scope=tle.gpu.smem, + ) + sS0_smem = sK0_tail_smem + else: + sS0_smem = tle.gpu.alloc( + [1, BH, BK], + dtype=kv.dtype.element_ty, + layout=None, + scope=tle.gpu.smem, + ) + is_kv_valid_smem = tle.gpu.alloc( + [1, PAIR_BLOCKS, BK], + dtype=tl.int8, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + k0_l_pipe = tle.pipe( + capacity=1, scope="cta", name="flashmla_sK0_l", sK=sK0_smem + ) + if HAVE_TAIL: + k0_r_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sK0_r", + readers=("qk", "remote"), + sK=sK0_smem, + sK_tail=sK0_tail_smem, + ) + else: + k0_r_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sK0_r", + readers=("qk", "remote"), + sK=sK0_smem, + ) + k1_l_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sK1_l", + readers=("qk", "remote"), + sK=sK1_smem, + ) + if HAVE_TAIL: + k1_r_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sK1_r", + sK=sK1_smem, + sK_tail=sK1_tail_smem, + ) + else: + k1_r_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_sK1_r", + sK=sK1_smem, + ) + is_kv_valid_pipe = tle.pipe( + capacity=1, + scope="cta", + name="flashmla_is_kv_valid_ready", + readers=("wg0", "wg1"), + is_kv_valid=is_kv_valid_smem, + ) + + sM_smem = tle.gpu.alloc( + [1, BH], + dtype=tl.float32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + sS1_smem = tle.gpu.alloc( + [1, BH, BK], dtype=kv.dtype.element_ty, layout=None, scope=tle.gpu.smem + ) + sL_smem = tle.gpu.alloc( + [2, BH], + dtype=tl.float32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + final_max_logits_smem = tle.gpu.alloc( + [BH], + dtype=tl.float32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + final_lse_smem = tle.gpu.alloc( + [BH], + dtype=tl.float32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + sM_wg0_pipe = tle.pipe( + capacity=1, scope="cta", name="flashmla_wg0_bunch_0_ready", sM=sM_smem + ) + sM_wg1_pipe = tle.pipe( + capacity=1, scope="cta", name="flashmla_wg1_bunch_0_ready", sM=sM_smem + ) + sS0_pipe = tle.pipe(capacity=1, scope="cta", name="flashmla_sS0", sS0=sS0_smem) + sS1_pipe = tle.pipe(capacity=1, scope="cta", name="flashmla_sS1", sS1=sS1_smem) + sL_wg0_pipe = tle.pipe( + capacity=2, scope="cta", name="flashmla_sL_wg0", sL=sL_smem + ) + sL_wg1_pipe = tle.pipe( + capacity=2, scope="cta", name="flashmla_sL_wg1", sL=sL_smem + ) + + log_scale: tl.constexpr = sm_scale * 1.4426950408889634 + + tle.gpu.warp_specialize( + [ + ( + _tle_flashmla_prefill_consumer0, + ( + q_pipe.writer(), + q_pipe.reader("wg0"), + q_desc, + tq_desc, + k0_l_pipe.reader(), + k0_r_pipe.reader("qk"), + k1_l_pipe.reader("remote", fields=("sK",)), + is_kv_valid_pipe.reader("wg0"), + sM_wg0_pipe.writer(), + sM_wg1_pipe.reader(), + sS0_pipe.writer(), + sS1_pipe.reader(), + sL_wg0_pipe.writer(), + sL_wg1_pipe.reader(), + output_desc, + q_row, + h_base, + topk_len_ptr, + attn_sink_base, + log_scale, + D, + TD, + kv.dtype.element_ty, + HAVE_ATTN_SINK, + TOPK, + HAVE_TOPK_LENGTH, + HAVE_TAIL, + BK, + BH, + DPH, + TDP, + G, + ), + ), + ( + _tle_flashmla_prefill_consumer1, + ( + q_pipe.reader("wg1"), + k1_r_pipe.reader(), + k1_l_pipe.reader("qk"), + k0_r_pipe.reader("remote", fields=("sK",)), + is_kv_valid_pipe.reader("wg1"), + sM_wg1_pipe.writer(), + sM_wg0_pipe.reader(), + sS1_pipe.writer(), + sS0_pipe.reader(), + sL_wg1_pipe.writer(), + sL_wg0_pipe.reader(), + final_max_logits_smem, + final_lse_smem, + output_desc, + q_row, + max_logits_base, + l_base, + h_base, + topk_len_ptr, + attn_sink_base, + log_scale, + D, + TD, + kv.dtype.element_ty, + HAVE_ATTN_SINK, + TOPK, + HAVE_TOPK_LENGTH, + HAVE_TAIL, + BK, + BH, + DPH, + TDP, + G, + ), + ), + ( + _tle_flashmla_prefill_producer, + ( + k0_l_pipe.writer(), + k0_r_pipe.writer(), + k1_l_pipe.writer(), + k1_r_pipe.writer(), + is_kv_valid_pipe.writer(), + kv_base, + tkv_base, + t_base, + topk_len_ptr, + D, + TD, + DPH, + TDP, + VG, + SKV, + TOPK, + HAVE_TOPK_LENGTH, + HAVE_TAIL, + BK, + ), + ), + ], + [4, 4], + [216, 72], + ) + + +def _flash_mla_sparse_tle_enabled() -> bool: + value = os.environ.get("FLAGGEMS_FLASHMLA_SPARSE_TLE", "1").lower() + return value not in {"0", "false", "off", "no"} + + +def _can_use_tle_flash_mla_sparse_fwd( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + d_v: int, + topk_length: Optional[torch.Tensor] = None, +) -> bool: + if not (HAS_TLE_FLASHMLA_SPARSE and _flash_mla_sparse_tle_enabled()): + return False + if q.device.type != "cuda": + return False + SQ, HQ, DQK = q.shape + _ = SQ + HKV = kv.shape[1] + TOPK = indices.shape[-1] + return ( + d_v == 512 + and HKV == 1 + and DQK in (512, 576) + and HQ % TLE_FLASHMLA_PREFILL_BH == 0 + and TOPK > 0 + and TOPK % 128 == 0 + ) + + +def _set_triton_descriptor_allocator(device: torch.device) -> None: + def alloc_fn(size: int, align: int, stream): + _ = align + _ = stream + return torch.empty(size, dtype=torch.int8, device=device) + + triton.set_allocator(alloc_fn) + + +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 + """ + assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous() + assert ( + q.dtype == torch.bfloat16 + and kv.dtype == torch.bfloat16 + and indices.dtype == torch.int32 + ) + SQ, HQ, DQK = q.shape + SKV, HKV, _ = kv.shape + + assert d_v == 512, "Unsupported d_v" + DV = d_v + + assert kv.shape[-1] == DQK + _, _, TOPK = indices.shape + assert indices.shape == (SQ, HKV, TOPK) + if attn_sink is not None: + assert attn_sink.is_contiguous() + assert attn_sink.dtype == torch.float32 + assert attn_sink.shape == (HQ,), "attn_sink error shape" + if topk_length is not None: + assert topk_length.is_contiguous() + assert topk_length.dtype == torch.int32 + assert topk_length.shape == (SQ,), "topk_length error shape" + + # check from FlashMLA + assert HKV == 1, "h_kv is expected to be 1" + assert HQ == 64 or HQ == 128, "Unsupported h_q" + assert DQK == 576 or DQK == 512, "Unsupported d_qk" + + _ = SKV + D = DV + TD = DQK - D + DP = triton.next_power_of_2(D) + HAVE_TAIL = TD > 0 + TDP = triton.next_power_of_2(TD) if HAVE_TAIL else 1 + G = HQ // HKV + BH = TLE_FLASHMLA_PREFILL_BH + RH = G // BH + BK = TLE_FLASHMLA_PREFILL_BK + output = torch.empty((SQ, HQ, DV), device=q.device, dtype=q.dtype) + max_logits = torch.empty((SQ, HQ), device=q.device, dtype=torch.float32) + lse = torch.empty((SQ, HQ), device=q.device, dtype=torch.float32) + + def triton_grid(META): + return (triton.cdiv(HQ, META["BH"]) * SQ,) + + if _can_use_tle_flash_mla_sparse_fwd(q, kv, indices, d_v, topk_length): + from triton.tools.tensor_descriptor import TensorDescriptor + + _set_triton_descriptor_allocator(q.device) + q_desc = TensorDescriptor( + q, shape=[SQ * HQ, DQK], strides=[DQK, 1], block_shape=[BH, DP // 2] + ) + if HAVE_TAIL: + tq_desc = TensorDescriptor( + q, shape=[SQ * HQ, DQK], strides=[DQK, 1], block_shape=[BH, TDP] + ) + else: + tq_desc = q_desc + output_desc = TensorDescriptor( + output, shape=[SQ * HQ, D], strides=[D, 1], block_shape=[BH, DP // 2] + ) + _tle_flashmla_prefill_fwd[triton_grid]( + q_desc, + tq_desc, + output_desc, + kv, + indices, + attn_sink, + topk_length, + sm_scale, + output, + max_logits, + lse, + SQ, + HQ, + DQK, + SKV, + TOPK, + attn_sink is not None, + topk_length is not None, + D, + TD, + DP, + TDP, + G, + HKV, + RH, + HAVE_TAIL, + BK, + BH, + TLE_FLASHMLA_PREFILL_PAIR_BLOCKS, + num_warps=TLE_FLASHMLA_PREFILL_WORKER_NUM_WARPS, + num_stages=1, + ) + return output, max_logits, lse + + triton_flash_mla_sparse_fwd[triton_grid]( + q, + kv, + indices, + attn_sink, + topk_length, + sm_scale, + output, + max_logits, + lse, + q.stride(1), + q.stride(0), + kv.stride(1), + kv.stride(0), + indices.stride(1), + indices.stride(0), + output.stride(1), + output.stride(0), + max_logits.stride(0), + lse.stride(0), + SQ, + HQ, + DQK, + SKV, + TOPK, + attn_sink is not None, + topk_length is not None, + ) + return output, max_logits, lse diff --git a/scripts/sparse_mla_port/vendor/flag_gems/__init__.py b/scripts/sparse_mla_port/vendor/flag_gems/__init__.py new file mode 100644 index 0000000..64b3aad --- /dev/null +++ b/scripts/sparse_mla_port/vendor/flag_gems/__init__.py @@ -0,0 +1 @@ +"""Minimal stand-in so the FlagGems fused kernels can be imported standalone.""" diff --git a/scripts/sparse_mla_port/vendor/flag_gems/utils/__init__.py b/scripts/sparse_mla_port/vendor/flag_gems/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/sparse_mla_port/vendor/flag_gems/utils/triton_version_utils.py b/scripts/sparse_mla_port/vendor/flag_gems/utils/triton_version_utils.py new file mode 100644 index 0000000..e44b19d --- /dev/null +++ b/scripts/sparse_mla_port/vendor/flag_gems/utils/triton_version_utils.py @@ -0,0 +1,27 @@ +"""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)