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