gfx1250 opus gemm splitk fuse

aiter/ops/opus/__init__.py

@@ -42,6 +42,8 @@ def _stub(*_args, **_kwargs):
         gemm_a16w16_opus,
         opus_gemm_a16w16_tune,
         opus_gemm_workspace_init,
+        opus_gemm_workspace_release,
+        opus_gemm_workspace_release_all,
     )
 
     def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
@@ -61,6 +63,12 @@ def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
         "opus_gemm_a8w8_blockscale_bpreshuffle_tune"
     )
     opus_gemm_workspace_init = _make_unsupported_arch_stub("opus_gemm_workspace_init")
+    opus_gemm_workspace_release = _make_unsupported_arch_stub(
+        "opus_gemm_workspace_release"
+    )
+    opus_gemm_workspace_release_all = _make_unsupported_arch_stub(
+        "opus_gemm_workspace_release_all"
+    )
 
 
 __all__ = [
@@ -69,4 +77,6 @@ def opus_gemm_a8w8_blockscale_bpreshuffle_tune(*args, **kwargs):
     "opus_gemm_a8w8_blockscale_bpreshuffle_tune",
     "opus_gemm_a16w16_tune",
     "opus_gemm_workspace_init",
+    "opus_gemm_workspace_release",
+    "opus_gemm_workspace_release_all",
 ]

aiter/ops/opus/gemm_op_a16w16.py

@@ -39,12 +39,15 @@
 module when that lands.
 """
 
+import functools
 import logging
+import os
+import sys
 
 import torch
 from torch import Tensor
 
-from ...jit.core import compile_ops
+from ...jit.core import AITER_ROOT_DIR, compile_ops
 from . import common as _opus_common
 
 logger = logging.getLogger("aiter")
@@ -57,6 +60,7 @@ def _gen_opus_gemm_a16w16_tune_fake_tensors(
     WQ: torch.Tensor,
     Y: torch.Tensor,
     bias: torch.Tensor | None = None,
+    workspace: torch.Tensor | None = None,
     kernelId: int = 0,
     splitK: int = 0,
 ) -> torch.Tensor:
@@ -80,6 +84,7 @@ def _opus_gemm_a16w16_tune_raw(
     WQ: torch.Tensor,
     Y: torch.Tensor,
     bias: torch.Tensor | None = None,
+    workspace: torch.Tensor | None = None,
     kernelId: int = 0,
     splitK: int = 0,
 ) -> torch.Tensor: ...
@@ -162,6 +167,112 @@ def _check_a16w16_tune_layout(XQ: torch.Tensor, WQ: torch.Tensor, Y: torch.Tenso
         )
 
 
+# gfx1250 split-K workspace element dtype (kid property) -> torch dtype. The
+# main kernel WRITES partials in this dtype and the reduce kernel READS them
+# back, so the buffer MUST be allocated with the matching element type: a
+# bf16-sized buffer handed to an fp32-workspace kid is half the bytes the kernel
+# writes -> global-memory overrun -> machine hang. Sizing is therefore done by
+# ELEMENT COUNT in this dtype, never by a raw byte size.
+_OPUS_WS_TORCH_DTYPE = {
+    "bf16_t": torch.bfloat16,
+    "fp32_t": torch.float32,
+}
+
+
+@functools.lru_cache(maxsize=1)
+def _gfx1250_kids() -> dict:
+    """Lazily load the opus kid table (csrc/opus_gemm/opus_gemm_common.py).
+
+    Lets the split-K workspace be sized from each kid's ACTUAL kernel
+    definition (tile B_M/B_N, split_k, workspace dtype) instead of a byte
+    guess. The module is pure-Python (no torch/JIT deps); returns ``{}`` if it
+    can't be located so the caller can fall back to a safe over-estimate.
+    """
+    csrc = os.path.join(AITER_ROOT_DIR, "csrc", "opus_gemm")
+    if csrc not in sys.path:
+        sys.path.insert(0, csrc)
+    try:
+        from opus_gemm_common import kernels_list  # type: ignore[import-not-found]
+
+        return kernels_list
+    except Exception:  # noqa: BLE001
+        return {}
+
+
+@functools.cache
+def _get_opus_workspace(
+    device: torch.device, ws_shape: tuple, dtype: torch.dtype
+) -> torch.Tensor:
+    """Cached split-K workspace with a data_ptr() stable across HIP graph
+    capture and replay.
+
+    Allocated with its natural ``[batch, split_k, padded_M, padded_N]`` element
+    shape (never a raw byte count) so the tensor is self-describing and matches
+    the kernel's ``ptr_ws`` layout; the launcher only reads ``data_ptr()``, so
+    the extra dims cost nothing.
+
+    A single torch.empty path serves eager AND capture: torch's caching
+    allocator is HIP graph-capture aware, so a torch.empty issued while a
+    capture is active is drawn from the graph's mempool and gets a stable
+    address that stays valid on replay -- exactly how a captured graph allocates
+    all of its other intermediates. lru_cache (keyed by device/ws_shape/dtype)
+    keeps the tensor alive for the process lifetime, so a shape first seen in an
+    eager pass is simply reused (cache hit) when its cudagraph is later captured,
+    and a shape first seen inside capture keeps that buffer pinned for every
+    subsequent replay. No eager pre-warm is required.
+    """
+    return torch.empty(ws_shape, dtype=dtype, device=device)
+
+
+def _alloc_splitk_workspace(
+    kernelId: int,
+    batch: int,
+    M: int,
+    N: int,
+    splitK: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Allocate the gfx1250 split-K partial workspace by ELEMENT COUNT in
+    ``[batch, split_k, padded_M, padded_N]`` -- never by a raw byte size.
+
+    Extents and element dtype come from the selected kid's own kernel
+    definition, so the buffer is exactly what the kernel writes and reads back:
+
+      * fuse kids (``a16w16_clusterlaunch_tdm_splitk_fuse``): ``split_k`` and the
+        workspace dtype are COMPILE-TIME per kid and the runtime ``splitK`` arg
+        is IGNORED, so the buffer is sized from ``fuse_split_k`` /
+        ``fuse_ws_dtype`` (a bf16- or fp32-workspace kid).
+      * ws-variant kids (``*_tdm_splitk_ws``): fp32 workspace with a runtime
+        ``split_k`` the launcher clamps DOWN from ``splitK`` (so ``splitK`` is a
+        safe upper bound).
+
+    Falls back to a safe over-estimate (fp32 element, split_k=16, 128x512 tile
+    padding -- the widest gfx1250 split-K tile is B_M<=128, B_N<=256) when the
+    kid table can't be loaded.
+    """
+    inst = _gfx1250_kids().get(int(kernelId))
+    if inst is not None:
+        b_m, b_n = int(inst.B_M), int(inst.B_N)
+        if inst.kernel_tag == "a16w16_clusterlaunch_tdm_splitk_fuse":
+            split_k = max(int(inst.fuse_split_k), 1)
+            ws_dtype = _OPUS_WS_TORCH_DTYPE.get(inst.fuse_ws_dtype, torch.float32)
+        else:
+            # ws-variant: fp32 workspace; launcher clamps split_k down from splitK.
+            split_k = max(1, int(splitK))
+            ws_dtype = torch.float32
+    else:
+        # Kid table unavailable: widest-element (fp32) upper bound. split_k must
+        # cover a fuse kid's max baked split_k (15) and any runtime splitK.
+        b_m, b_n = 128, 512
+        split_k = max(int(splitK), 16)
+        ws_dtype = torch.float32
+
+    padded_M = ((int(M) + b_m - 1) // b_m) * b_m
+    padded_N = ((int(N) + b_n - 1) // b_n) * b_n
+    ws_shape = (int(batch), split_k, padded_M, padded_N)
+    return _get_opus_workspace(device, ws_shape, ws_dtype)
+
+
 def opus_gemm_a16w16_tune(
     XQ: torch.Tensor,
     WQ: torch.Tensor,
@@ -210,6 +321,12 @@ def opus_gemm_a16w16_tune(
         splitK = new_splitK
         bias = None
     _check_a16w16_tune_layout(XQ, WQ, Y)
+    # gfx1250 split-K kids [20000, 30000) need a workspace tensor allocated
+    # externally (torch.empty) and passed to the C++ launcher.
+    workspace = None
+    if 20000 <= kernelId < 30000:
+        batch, M, N = Y.shape
+        workspace = _alloc_splitk_workspace(kernelId, batch, M, N, splitK, XQ.device)
     # Mono-tile kid guard: the launcher requires N / K to be tile-aligned
     # (the kernel has no N-tail mask and no K-tail mask; M-tail IS handled
     # via the bounded gmem desc). A CSV winner picked through
@@ -234,7 +351,7 @@ def opus_gemm_a16w16_tune(
     # refactor to aiter_tensor_t). Keep the wrapper's `return Y`
     # contract so callers that did `Y = opus_gemm_a16w16_tune(...)`
     # still see the populated Y.
-    _opus_gemm_a16w16_tune_raw(XQ, WQ, Y, bias, kernelId, splitK)
+    _opus_gemm_a16w16_tune_raw(XQ, WQ, Y, bias, workspace, kernelId, splitK)
     return Y
 
 
@@ -393,7 +510,7 @@ def _validate_and_reshape(A: Tensor, B: Tensor, bias, dtype, out):
         WQ = B
     else:
         raise ValueError(
-            f"B must be 2D [N, K] or 3D [batch, N, K] (got shape " f"{tuple(B.shape)})"
+            f"B must be 2D [N, K] or 3D [batch, N, K] (got shape {tuple(B.shape)})"
         )
 
     if out is not None:
@@ -551,9 +668,31 @@ def gemm_a16w16_opus(
 def opus_gemm_workspace_init() -> None: ...
 
 
+# Free the per-stream splitk workspace registered by opus_gemm_workspace_init
+# (and grown by the splitk launchers). Call inside `with torch.cuda.stream(s):`
+# in eager mode (not during HIP graph capture) to reclaim the GPU buffer +
+# handles for that stream; no-op if the stream was never registered. Use this
+# for explicit teardown of streams the framework will not reuse.
+@compile_ops(
+    "module_deepgemm_opus", fc_name="opus_gemm_workspace_release", develop=True
+)
+def opus_gemm_workspace_release() -> None: ...
+
+
+# Free the splitk workspace for all registered streams and clear the registry.
+# Eager mode only. Use for a full teardown before the framework reclaims its
+# stream pool / at process shutdown.
+@compile_ops(
+    "module_deepgemm_opus", fc_name="opus_gemm_workspace_release_all", develop=True
+)
+def opus_gemm_workspace_release_all() -> None: ...
+
+
 __all__ = [
     "gemm_a16w16_opus",
     "is_splitk_kid",
     "opus_gemm_a16w16_tune",
     "opus_gemm_workspace_init",
+    "opus_gemm_workspace_release",
+    "opus_gemm_workspace_release_all",
 ]

aiter/ops/triton/attention/fp8_mqa_logits.py

@@ -216,6 +216,8 @@ def fp8_mqa_logits(
             num_chains = 8 if USE_FOLDED_REDUCTION else 0
             num_warps = 4
             block_kv = 128
+            # This kernel has no BLOCK_M: it walks one query row per program.
+            block_m = 1
             other = {"LOOP_VARIANT": loop_variant}
 
         # Buffer ops use a 32-bit byte offset (2 GiB resource descriptor cap).

aiter/tuned_gemm.py

@@ -41,100 +41,22 @@ def is_flydsl_available():
 from aiter.ops.gemm_op_common import get_padded_m
 
 try:
-    from aiter.ops.opus.gemm_op_a16w16 import is_splitk_kid as _opus_is_splitk_kid
     from aiter.ops.opus.gemm_op_a16w16 import opus_gemm_a16w16_tune as _opus_tune
-    from aiter.ops.opus.gemm_op_a16w16 import (
-        opus_gemm_workspace_init as _opus_workspace_init,
-    )
 except Exception:  # noqa: BLE001  blanket catch is intentional here
     _opus_tune = None
-    _opus_workspace_init = None
-    _opus_is_splitk_kid = None
-
-# Every opus split-K arch (gfx950 / gfx942 / gfx1250) owns a per-stream fp32
-# workspace (process-global `opus_splitk_ws_get` registry, backed by raw
-# hipMalloc) that must be registered AND grown to the shape's size *eagerly*
-# before HIP graph capture -- hipMalloc/hipFree are stream-capture-illegal, so a
-# grow inside capture aborts the capture, leaving an empty graph whose replay
-# silently writes zeros (garbage logits). torch.cuda.graph captures on a
-# process-global stream (`torch.cuda.graphs.graph.default_capture_stream`) when
-# no explicit stream is passed (the vLLM/ATOM CUDAGraphWrapper case); we warm
-# that stream here during the eager pass so a later capture of the same shape
-# finds a ready workspace. (The opus launcher reads a stable device-resident
-# handle, so the captured graph stays valid across replays / post-capture grows
-# -- which is exactly why opus keeps a persistent workspace instead of a
-# per-call hipMallocAsync that would not survive capture; the only cost is this
-# one-time warm.)
-_OPUS_WS_ARCHS = {"gfx950", "gfx942", "gfx1250"}
-_opus_ws_warmed_sigs = set()
-
-
-@functools.lru_cache(maxsize=1)
-def _opus_needs_ws_prewarm() -> bool:
-    if _opus_tune is None or _opus_workspace_init is None:
-        return False
-    try:
-        return get_gfx() in _OPUS_WS_ARCHS
-    except Exception:  # noqa: BLE001
-        return False
 
-
-def _opus_graph_capture_stream():
-    """The stream torch.cuda.graph captures on when called without `stream=`.
-
-    Mirrors torch's own lazy-init so we register the opus workspace on the exact
-    stream a later `with torch.cuda.graph(g):` will use.
-    """
-    g = torch.cuda.graphs.graph
-    if getattr(g, "default_capture_stream", None) is None:
-        g.default_capture_stream = torch.cuda.Stream()
-    return g.default_capture_stream
-
-
-def _opus_prewarm_capture_workspace(inp, weights, solidx, splitK, bias, otype):
-    """Eagerly size the opus split-K workspace on the graph capture stream.
-
-    No-op when already capturing (too late to allocate), on non-registry archs,
-    for a non-split-K kid (never touches the workspace), or when this
-    (shape, kid, splitK, bias) was already warmed.
-    """
-    if not _opus_needs_ws_prewarm():
-        return
-    # Only split-K kids allocate/read the fp32 workspace; every other kid family
-    # (flatmm / persistent / mono_tile / nosplit) launches straight to its kernel
-    # and never touches the registry, so warming it for them is pure waste.
-    if _opus_is_splitk_kid is not None and not _opus_is_splitk_kid(solidx):
-        return
-    if torch.cuda.is_current_stream_capturing():
-        return
-    m, k = inp.shape
-    n = weights.shape[0]
-    sig = (int(solidx), m, n, k, int(splitK), bias is not None, str(otype))
-    if sig in _opus_ws_warmed_sigs:
-        return
-    try:
-        s = _opus_graph_capture_stream()
-        with torch.cuda.stream(s):
-            _opus_workspace_init()
-            Yw = torch.empty(m, n, dtype=otype or inp.dtype, device=inp.device)
-            _opus_tune(
-                inp.unsqueeze(0),
-                weights.unsqueeze(0),
-                Yw.unsqueeze(0),
-                bias=bias,
-                kernelId=int(solidx),
-                splitK=int(splitK),
-            )
-        s.synchronize()
-        _opus_ws_warmed_sigs.add(sig)
-    # Don't break eager callers; capture would re-surface it.
-    except Exception as e:  # noqa: BLE001
-        logger.warning(
-            f"opus split-K workspace prewarm on the graph capture stream failed "
-            f"({type(e).__name__}: {e}); HIP graph capture of this opus shape may "
-            f"produce zeros. Call aiter.opus_gemm_workspace_init() on the capture "
-            f"stream manually if you capture with a custom stream."
-        )
+# NOTE: gfx1250 split-K kids allocate their partial-sum workspace as a plain
+# torch.empty tensor (see aiter.ops.opus.gemm_op_a16w16._get_opus_workspace)
+# passed explicitly to the launcher. torch's caching allocator is HIP graph-
+# capture aware, so that single torch.empty path serves both eager and capture
+# (a buffer first touched inside capture comes from the graph mempool with a
+# replay-stable address) and no eager pre-warm of the shape is required. (The
+# old per-stream hipMalloc registry -- opus_gemm_workspace_init /
+# opus_splitk_ws_get -- used by the gfx942/gfx950 a16w16 split-K path still needs
+# an eager warm before capture; if that path is ever exercised under cudagraphs,
+# warm it via aiter.opus_gemm_workspace_init() on the capture stream. It fails
+# loudly ("splitk workspace not initialized") rather than silently corrupting,
+# so its absence here is safe to detect.)
 
 
 this_dir = os.path.dirname(os.path.abspath(__file__))
@@ -626,11 +548,8 @@ def opus_gemm(
     splitK = int(config.get("splitK", 0)) if config is not None else 0
     m, _k = inp.shape
     n = weights.shape[0]
-    # Eagerly size the per-stream split-K workspace on torch's graph capture
-    # stream so a later HIP graph capture of this shape doesn't abort (which
-    # would leave the captured graph empty -> replay writes zeros). No-op when
-    # already capturing, on gfx950, or for an already-warmed shape.
-    _opus_prewarm_capture_workspace(inp, weights, solidx, splitK, bias, otype)
+    # The split-K workspace (if any) is allocated capture-safely inside
+    # opus_gemm_a16w16_tune -> _get_opus_workspace; no eager pre-warm needed.
     Y = torch.empty(m, n, dtype=otype or inp.dtype, device=inp.device)
     _opus_tune(
         inp.unsqueeze(0),