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)