|
| 1 | +# pyright: reportInvalidTypeForm=false, reportMissingImports=false |
| 2 | +"""Real fused DeepSeek-V3.2 Sparse-MLA forward+backward (TileLang-CUDA). |
| 3 | +
|
| 4 | +This module wraps the *real* training-complete fused Sparse-MLA kernels we own |
| 5 | +at ``tilelang/examples/deepseek_v32/sparse_mla_fwd.py`` + |
| 6 | +``sparse_mla_bwd.py`` and exposes them over MLX bf16 ``q``/``kv``/``indices`` |
| 7 | +buffers so Path C can dispatch the fused O(seq*topk) attention (forward AND |
| 8 | +backward, with vectorized ``T.atomic_addx4`` dKV) instead of the per-element |
| 9 | +string-emitted reference. |
| 10 | +
|
| 11 | +It is **env-gated** (``CPPMEGA_SPARSE_MLA_V32_FUSED=1``). When the gate is OFF |
| 12 | +(default) callers keep using the existing reference path; nothing here runs. |
| 13 | +
|
| 14 | +RULE #1 (no silent fallback): when the gate is ON and a caller forces Path C, |
| 15 | +every failure here RAISES with where+what. We never silently fall back to the |
| 16 | +reference and claim the fused kernel ran. |
| 17 | +
|
| 18 | +ABI / layout |
| 19 | +------------ |
| 20 | +* ``q`` : bf16 ``(B, S, H, DQK)`` where ``DQK = d_v + tail_dim`` (576 = 512+64). |
| 21 | +* ``kv`` : bf16 ``(B, SKV, G, DQK)`` (kv_group ``G``). |
| 22 | +* ``indices`` : int32 ``(B, S, G, topk)``. |
| 23 | +* fwd -> ``out`` bf16 ``(B, S, H, d_v)``, ``lse`` fp32 ``(B, S, H)``. |
| 24 | +* bwd -> ``dq`` bf16 ``(B, S, H, DQK)``, ``dkv`` bf16 ``(B, SKV, G, DQK)``. |
| 25 | +
|
| 26 | +The kernel math is byte-for-byte the upstream v32 example (log2-domain online |
| 27 | +softmax forward; preprocess-delta + atomic-scatter dKV backward). |
| 28 | +
|
| 29 | +gb10 / sm_121 note |
| 30 | +------------------ |
| 31 | +The v32 fwd/bwd tiles request >48 KB dynamic shared memory. On gb10/sm_121 the |
| 32 | +TileLang/TVM runtime's ``cuFuncSetAttribute(MAX_DYNAMIC_SHARED_SIZE_BYTES)`` |
| 33 | +opt-in is rejected by the driver above ~48 KB even though the device reports a |
| 34 | +99 KB opt-in carve-out, so the fused kernel does not yet run there (it lowers + |
| 35 | +compiles for sm_121a; only the runtime smem opt-in is blocked). On H100 (sm_90) |
| 36 | +/ B200 (sm_100) the carve-out fits and the fused path runs. This is surfaced |
| 37 | +loudly when forced, never silently degraded. |
| 38 | +""" |
| 39 | + |
| 40 | +from __future__ import annotations |
| 41 | + |
| 42 | +import os |
| 43 | +from functools import lru_cache |
| 44 | +from typing import Any |
| 45 | + |
| 46 | +import mlx.core as mx |
| 47 | + |
| 48 | + |
| 49 | +SPARSE_MLA_V32_FUSED_ENV = "CPPMEGA_SPARSE_MLA_V32_FUSED" |
| 50 | + |
| 51 | +# v32 layout constant: q/kv last dim is d_v + tail_dim; the upstream kernels hard |
| 52 | +# assume d_v == 512 (and validate it), so tail_dim == DQK - 512. |
| 53 | +_V32_D_V = 512 |
| 54 | + |
| 55 | + |
| 56 | +def v32_fused_enabled() -> bool: |
| 57 | + """Return True iff the real fused v32 Sparse-MLA path is env-gated ON.""" |
| 58 | + |
| 59 | + return os.environ.get(SPARSE_MLA_V32_FUSED_ENV, "").strip().lower() in { |
| 60 | + "1", |
| 61 | + "true", |
| 62 | + "yes", |
| 63 | + "on", |
| 64 | + } |
| 65 | + |
| 66 | + |
| 67 | +def _torch(): |
| 68 | + try: |
| 69 | + import torch # noqa: PLC0415 |
| 70 | + |
| 71 | + return torch |
| 72 | + except Exception as exc: # pragma: no cover - hosts without torch |
| 73 | + raise RuntimeError( |
| 74 | + "_sparse_mla_v32_fused: PyTorch (CUDA) is required for the fused v32 " |
| 75 | + f"Sparse-MLA kernels but could not be imported ({type(exc).__name__}: {exc})." |
| 76 | + ) from exc |
| 77 | + |
| 78 | + |
| 79 | +def _mlx_to_torch_cuda(arr: mx.array): |
| 80 | + """Materialize an MLX array as a contiguous CUDA torch tensor (no host copy |
| 81 | + when a CUDA->CUDA DLPack bridge is available; otherwise via DLPack).""" |
| 82 | + |
| 83 | + torch = _torch() |
| 84 | + from cppmega_mlx.nn._tilelang._cuda_eager import _mlx_to_torch_cuda as _bridge |
| 85 | + |
| 86 | + t = _bridge(arr) |
| 87 | + if not t.is_cuda: |
| 88 | + raise RuntimeError( |
| 89 | + "_sparse_mla_v32_fused: MLX->torch bridge did not yield a CUDA tensor; " |
| 90 | + "the fused v32 Sparse-MLA path requires CUDA buffers." |
| 91 | + ) |
| 92 | + return t.contiguous() |
| 93 | + |
| 94 | + |
| 95 | +def _torch_cuda_to_mlx(tensor, dtype: mx.Dtype) -> mx.array: |
| 96 | + from cppmega_mlx.nn._tilelang._cuda_eager import _torch_cuda_to_mlx as _bridge |
| 97 | + |
| 98 | + return _bridge(tensor, dtype) |
| 99 | + |
| 100 | + |
| 101 | +@lru_cache(maxsize=1) |
| 102 | +def _v32_modules() -> tuple[Any, Any]: |
| 103 | + """Import the vendored v32 fwd/bwd example modules from the TileLang tree. |
| 104 | +
|
| 105 | + The examples live under ``tilelang/examples/deepseek_v32`` and import each |
| 106 | + other + ``utils`` by bare name, so the directory must be on ``sys.path``. |
| 107 | + RULE #1: a missing TileLang examples tree RAISES (no silent skip). |
| 108 | + """ |
| 109 | + |
| 110 | + import importlib |
| 111 | + import sys |
| 112 | + |
| 113 | + try: |
| 114 | + import tilelang # noqa: PLC0415,F401 |
| 115 | + except Exception as exc: |
| 116 | + raise RuntimeError( |
| 117 | + "_sparse_mla_v32_fused: tilelang is not importable; the fused v32 " |
| 118 | + f"Sparse-MLA path cannot run ({type(exc).__name__}: {exc})." |
| 119 | + ) from exc |
| 120 | + |
| 121 | + example_dir = os.environ.get("CPPMEGA_TILELANG_V32_EXAMPLES_DIR") |
| 122 | + candidates = [] |
| 123 | + if example_dir: |
| 124 | + candidates.append(example_dir) |
| 125 | + # Best-effort discovery relative to an installed/dev tilelang checkout. |
| 126 | + tl_file = getattr(importlib.import_module("tilelang"), "__file__", None) |
| 127 | + if tl_file: |
| 128 | + root = os.path.dirname(os.path.dirname(os.path.abspath(tl_file))) |
| 129 | + candidates.append(os.path.join(root, "examples", "deepseek_v32")) |
| 130 | + for env_root in ("TILELANG_ROOT", "TILELANG_SOURCE_DIR"): |
| 131 | + r = os.environ.get(env_root) |
| 132 | + if r: |
| 133 | + candidates.append(os.path.join(r, "examples", "deepseek_v32")) |
| 134 | + |
| 135 | + chosen = next( |
| 136 | + (c for c in candidates if c and os.path.isfile(os.path.join(c, "sparse_mla_fwd.py"))), |
| 137 | + None, |
| 138 | + ) |
| 139 | + if chosen is None: |
| 140 | + raise RuntimeError( |
| 141 | + "_sparse_mla_v32_fused: could not locate the TileLang deepseek_v32 " |
| 142 | + "examples (sparse_mla_fwd.py/sparse_mla_bwd.py). Set " |
| 143 | + "CPPMEGA_TILELANG_V32_EXAMPLES_DIR to the examples/deepseek_v32 " |
| 144 | + f"directory. Searched: {candidates}" |
| 145 | + ) |
| 146 | + if chosen not in sys.path: |
| 147 | + sys.path.insert(0, chosen) |
| 148 | + fwd = importlib.import_module("sparse_mla_fwd") |
| 149 | + bwd = importlib.import_module("sparse_mla_bwd") |
| 150 | + return fwd, bwd |
| 151 | + |
| 152 | + |
| 153 | +def _resolve_dims(q: mx.array, kv: mx.array, indices: mx.array): |
| 154 | + if q.ndim != 4 or kv.ndim != 4 or indices.ndim != 4: |
| 155 | + raise ValueError( |
| 156 | + "_sparse_mla_v32_fused: expected 4D q/kv/indices " |
| 157 | + f"(got q={tuple(q.shape)} kv={tuple(kv.shape)} indices={tuple(indices.shape)})." |
| 158 | + ) |
| 159 | + B, S, H, DQK = (int(x) for x in q.shape) |
| 160 | + Bk, SKV, G, DQKk = (int(x) for x in kv.shape) |
| 161 | + if Bk != B or DQKk != DQK: |
| 162 | + raise ValueError( |
| 163 | + "_sparse_mla_v32_fused: q/kv batch or qk-dim mismatch " |
| 164 | + f"(q={tuple(q.shape)} kv={tuple(kv.shape)})." |
| 165 | + ) |
| 166 | + Bi, Si, Gi, topk = (int(x) for x in indices.shape) |
| 167 | + if (Bi, Si, Gi) != (B, S, G): |
| 168 | + raise ValueError( |
| 169 | + "_sparse_mla_v32_fused: indices shape must be (B, S, kv_group, topk); " |
| 170 | + f"got {tuple(indices.shape)} for q={tuple(q.shape)} kv={tuple(kv.shape)}." |
| 171 | + ) |
| 172 | + if DQK <= _V32_D_V: |
| 173 | + raise ValueError( |
| 174 | + "_sparse_mla_v32_fused: the v32 kernels require qk_dim > d_v=512 " |
| 175 | + f"(got DQK={DQK}); this is the d_v(512)+tail layout." |
| 176 | + ) |
| 177 | + return B, S, SKV, H, G, DQK, topk |
| 178 | + |
| 179 | + |
| 180 | +def v32_fused_apply( |
| 181 | + q: mx.array, |
| 182 | + kv: mx.array, |
| 183 | + indices: mx.array, |
| 184 | + *, |
| 185 | + sm_scale: float, |
| 186 | + return_lse: bool = False, |
| 187 | +) -> mx.array | tuple[mx.array, mx.array]: |
| 188 | + """Run the real fused v32 Sparse-MLA forward over bf16 q/kv/indices. |
| 189 | +
|
| 190 | + Returns ``out`` bf16 ``(B, S, H, d_v)`` (and ``lse`` fp32 ``(B, S, H)`` when |
| 191 | + ``return_lse``). RAISES on any kernel failure (RULE #1). |
| 192 | + """ |
| 193 | + |
| 194 | + fwd, _bwd = _v32_modules() |
| 195 | + _resolve_dims(q, kv, indices) |
| 196 | + torch = _torch() |
| 197 | + |
| 198 | + q_t = _mlx_to_torch_cuda(q.astype(mx.bfloat16)) |
| 199 | + kv_t = _mlx_to_torch_cuda(kv.astype(mx.bfloat16)) |
| 200 | + idx_t = _mlx_to_torch_cuda(indices.astype(mx.int32)) |
| 201 | + try: |
| 202 | + out_t, lse_t = fwd.sparse_mla_fwd_interface( |
| 203 | + q_t.contiguous(), |
| 204 | + kv_t.contiguous(), |
| 205 | + idx_t.contiguous(), |
| 206 | + sm_scale=float(sm_scale), |
| 207 | + ) |
| 208 | + torch.cuda.synchronize() |
| 209 | + except Exception as exc: |
| 210 | + raise RuntimeError( |
| 211 | + "_sparse_mla_v32_fused.v32_fused_apply: the fused DeepSeek-V3.2 " |
| 212 | + f"Sparse-MLA forward kernel failed ({type(exc).__name__}: {exc}). " |
| 213 | + "On gb10/sm_121 the >48 KB dynamic-smem opt-in is rejected by the " |
| 214 | + "CUDA driver; on sm_90/sm_100 it should run." |
| 215 | + ) from exc |
| 216 | + |
| 217 | + out = _torch_cuda_to_mlx(out_t, mx.bfloat16) |
| 218 | + if return_lse: |
| 219 | + lse = _torch_cuda_to_mlx(lse_t, mx.float32) |
| 220 | + return out, lse |
| 221 | + return out |
| 222 | + |
| 223 | + |
| 224 | +def v32_fused_bwd( |
| 225 | + q: mx.array, |
| 226 | + kv: mx.array, |
| 227 | + out: mx.array, |
| 228 | + d_out: mx.array, |
| 229 | + indices: mx.array, |
| 230 | + lse: mx.array, |
| 231 | + *, |
| 232 | + sm_scale: float, |
| 233 | +) -> tuple[mx.array, mx.array]: |
| 234 | + """Run the real fused v32 Sparse-MLA backward -> ``(dq, dkv)`` bf16. |
| 235 | +
|
| 236 | + Exercises the vectorized ``T.atomic_addx4`` dKV scatter. RAISES on any |
| 237 | + kernel failure (RULE #1). |
| 238 | + """ |
| 239 | + |
| 240 | + _fwd, bwd = _v32_modules() |
| 241 | + _resolve_dims(q, kv, indices) |
| 242 | + torch = _torch() |
| 243 | + |
| 244 | + q_t = _mlx_to_torch_cuda(q.astype(mx.bfloat16)).contiguous() |
| 245 | + kv_t = _mlx_to_torch_cuda(kv.astype(mx.bfloat16)).contiguous() |
| 246 | + out_t = _mlx_to_torch_cuda(out.astype(mx.bfloat16)).contiguous() |
| 247 | + do_t = _mlx_to_torch_cuda(d_out.astype(mx.bfloat16)).contiguous() |
| 248 | + idx_t = _mlx_to_torch_cuda(indices.astype(mx.int32)).contiguous() |
| 249 | + lse_t = _mlx_to_torch_cuda(lse.astype(mx.float32)).contiguous() |
| 250 | + try: |
| 251 | + dq_t, dkv_t = bwd.sparse_mla_bwd( |
| 252 | + q_t, |
| 253 | + kv_t, |
| 254 | + out_t, |
| 255 | + do_t, |
| 256 | + idx_t, |
| 257 | + lse_t, |
| 258 | + sm_scale=float(sm_scale), |
| 259 | + ) |
| 260 | + torch.cuda.synchronize() |
| 261 | + except Exception as exc: |
| 262 | + raise RuntimeError( |
| 263 | + "_sparse_mla_v32_fused.v32_fused_bwd: the fused DeepSeek-V3.2 " |
| 264 | + f"Sparse-MLA backward kernel failed ({type(exc).__name__}: {exc}). " |
| 265 | + "This path exercises T.atomic_addx4 (vectorized dKV scatter); on " |
| 266 | + "gb10/sm_121 the >48 KB dynamic-smem opt-in is rejected by the CUDA " |
| 267 | + "driver before atomics are reached, on sm_90/sm_100 it should run." |
| 268 | + ) from exc |
| 269 | + |
| 270 | + dq = _torch_cuda_to_mlx(dq_t, mx.bfloat16) |
| 271 | + dkv = _torch_cuda_to_mlx(dkv_t.astype(torch.bfloat16), mx.bfloat16) |
| 272 | + return dq, dkv |
0 commit comments