Q8KV4 prefill: build for SM107, speed up K1, unify the combine, place partials by locality domain

README.md

@@ -119,15 +119,15 @@ python -c 'import cutlass; print(cutlass.__version__, cutlass.CUDA_VERSION)'
 ```
 
 Q8KV4 C++ operators use the standard `CUDA_HOME`, `CUTLASS_ROOT`, and
-`TORCH_EXTENSIONS_DIR` configuration. Runtime JIT always selects `sm_100a` or
-`sm_103a` from the CUDA device of the input tensor. Tensor-less AOT/offline
-builds use `MM_SPARSE_TARGET_ARCH=100a|103a` to select a target (default:
-`103a`). `TORCH_CUDA_ARCH_LIST` is internal compiler plumbing rather than the
+`TORCH_EXTENSIONS_DIR` configuration. Runtime JIT always selects `sm_100a`,
+`sm_103a` or `sm_107a` (operators with SM107 support) from the CUDA device of
+the input tensor. Tensor-less AOT/offline builds use
+`MM_SPARSE_TARGET_ARCH=100a|103a|107a` to select a target (default: `103a`). `TORCH_CUDA_ARCH_LIST` is internal compiler plumbing rather than the
 runtime architecture-selection interface. Q8KV4 Decode Attention, Decode
 Indexer, and the shared dequant package support CUDA Toolkit 12.9 or newer.
 They use QMUL4 when the selected toolchain supports it and otherwise select an
 exact fallback automatically. Q8KV4 Prefill Attention requires CUDA Toolkit
-13.4 or newer.
+13.4 or newer (13.5 or newer on SM107).
 
 BF16 and Q8KV8 sparse decode are provided through an optional external FlashInfer
 TRTLLM-GEN backend. This repository does not include FlashInfer source code or
@@ -136,7 +136,7 @@ Q8KV4 decode uses native CUTLASS C++. All three decode attention formats support
 output on B200/B300, selecting the implementation by the actual `Hq/Hkv`.
 Q8KV4 decode also retains SM107 GQA=16 support, requiring CUDA Toolkit 13.5 or newer.
 
-Rubin (SM107) retains BF16/Q8KV8 prefill attention, Q8KV8 prefill indexer, and Q8KV4 decode attention paths,
+Rubin (SM107) retains BF16/Q8KV8/Q8KV4 prefill attention, Q8KV8 prefill indexer, and Q8KV4 decode attention paths,
 selected from the input tensor device. See the operator READMEs for dependencies and restrictions.
 This does not imply Rubin support for the other operators.
 

README.zh-CN.md

@@ -112,20 +112,20 @@ python -c 'import cutlass; print(cutlass.__version__, cutlass.CUDA_VERSION)'
 ```
 
 Q8KV4 C++ 算子使用标准的 `CUDA_HOME`、`CUTLASS_ROOT` 和 `TORCH_EXTENSIONS_DIR`
-配置。运行时 JIT 始终根据输入 tensor 所在的 CUDA device 选择 `sm_100a` 或
-`sm_103a`；无 tensor 的 AOT/离线构建使用 `MM_SPARSE_TARGET_ARCH=100a|103a` 指定
-目标（默认 `103a`）。`TORCH_CUDA_ARCH_LIST` 只由内部编译流程设置，不作为运行时
+配置。运行时 JIT 始终根据输入 tensor 所在的 CUDA device 选择 `sm_100a`、
+`sm_103a` 或 `sm_107a`（支持 SM107 的算子）；无 tensor 的 AOT/离线构建使用
+`MM_SPARSE_TARGET_ARCH=100a|103a|107a` 指定目标（默认 `103a`）。`TORCH_CUDA_ARCH_LIST` 只由内部编译流程设置，不作为运行时
 架构选择接口。Q8KV4 Decode Attention、Decode Indexer 和共享 dequant 支持 CUDA
 Toolkit 12.9 或更高版本；工具链支持 QMUL4 时使用该路径，否则自动使用精确
-fallback。Q8KV4 Prefill Attention 要求 CUDA Toolkit 13.4 或更高版本。
+fallback。Q8KV4 Prefill Attention 要求 CUDA Toolkit 13.4 或更高版本（SM107 上要求 13.5 或更高版本）。
 
 BF16 和 Q8KV8 sparse decode 通过可选的外部 FlashInfer TRTLLM-GEN backend 提供。本仓库不包含
 FlashInfer 源码或 cubin；使用 `python -m pip install -e '.[flashinfer]'` 安装
 `flashinfer-python==0.6.17`。Q8KV4 decode 使用原生 CUTLASS C++；三种 decode attention
 均支持 B200/B300 上的 GQA=8/16 和 BF16 输出，按实际 `Hq/Hkv` 选择实现。
 Q8KV4 decode 另保留 SM107 的 GQA=16 支持，需要 CUDA Toolkit 13.5 或更新版本。
 
-Rubin（SM107）保留 BF16/Q8KV8 prefill attention、Q8KV8 prefill indexer 和 Q8KV4 decode attention 路径，
+Rubin（SM107）保留 BF16/Q8KV8/Q8KV4 prefill attention、Q8KV8 prefill indexer 和 Q8KV4 decode attention 路径，
 通过输入 tensor 所在设备自动选择；具体依赖和限制见各算子 README。其余算子不据此声明 Rubin 支持。
 
 ## 3. 主要接口

benchmarks/inference/msa_v1/attention/prefill/q8kv4/benchmark.py

@@ -371,8 +371,8 @@ def main() -> None:
     args = parser.parse_args()
     if args.slots < 0 or args.warmup < 1 or args.replays < 2:
         parser.error("slots must be non-negative, warmup positive, and replays >= 2")
-    if torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}:
-        raise RuntimeError("SM100-family GPU (B200/GB300) required")
+    if torch.cuda.get_device_capability() not in {(10, 0), (10, 3), (10, 7)}:
+        raise RuntimeError("SM100, SM103 or SM107 GPU required")
 
     device = torch.device("cuda")
     _warm_compilation(device)

inference/msa_v1/README.md

@@ -32,7 +32,7 @@ Available formats and constraints:
   CUDA 13.5 or newer.
 - Q8KV8 prefill attention supports GQA group sizes 1, 2, 4, 8, and 16.
 
-BF16/Q8KV8 prefill attention also retains SM107 Rubin paths selected by the actual device architecture.
+BF16/Q8KV8/Q8KV4 prefill attention also retains SM107 Rubin paths selected by the actual device architecture.
 
 See the README in each operator directory for exact tensor shapes, dtypes, and usage examples.
 

inference/msa_v1/README.zh-CN.md

@@ -31,7 +31,7 @@ MSA v1 提供 paged sparse attention 和 indexer 的 decode/prefill 接口。所
   Q8KV4 使用原生 CUTLASS C++，并保留 CUDA 13.5 或更新版本上的 SM107 GQA=16 支持。
 - Q8KV8 prefill attention 支持 GQA group size 1、2、4、8 或 16。
 
-BF16/Q8KV8 prefill attention 还保留 SM107 Rubin 路径，由实际设备架构选择。
+BF16/Q8KV8/Q8KV4 prefill attention 还保留 SM107 Rubin 路径，由实际设备架构选择。
 
 具体张量 shape、dtype 和调用示例见各算子目录中的 README。
 

inference/msa_v1/_build_utils.py

@@ -2,15 +2,35 @@
 
 from __future__ import annotations
 
+import fcntl
+import hashlib
+import logging
 import os
 import re
 import shutil
 import subprocess
+from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
 
+logger = logging.getLogger(__name__)
+
 _REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
 MM_SPARSE_TARGET_ARCH = "MM_SPARSE_TARGET_ARCH"
+QMUL4_PROBE_SOURCE = r"""
+#include <cstdint>
+
+__global__ void qmul4_probe(uint32_t* output) {
+  uint32_t result;
+  uint16_t packed_e2m1 = 0;
+  uint32_t scale_e4m3 = 0;
+  asm volatile(
+      "mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4 %0, %1, %2;"
+      : "=r"(result)
+      : "h"(packed_e2m1), "r"(scale_e4m3));
+  output[0] = result;
+}
+""".lstrip()
 
 
 def normalize_target_arch(
@@ -181,13 +201,90 @@ def cutlass_root() -> Path:
     )
 
 
+@contextmanager
+def build_lock(cache_dir: Path):
+    """Keep a shared library immutable while another rank builds or loads it."""
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    with (cache_dir / "build.lock").open("a") as lock_file:
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+def _write_text_if_changed(path: Path, content: str) -> None:
+    if path.is_file() and path.read_text(encoding="utf-8") == content:
+        return
+    path.write_text(content, encoding="utf-8")
+
+
+def _probe_qmul4(nvcc: Path, arch: str, probe_dir: Path) -> bool:
+    result_path = probe_dir / "qmul4.result"
+    if result_path.is_file():
+        return result_path.read_text(encoding="utf-8").strip() == "supported"
+
+    probe_dir.mkdir(parents=True, exist_ok=True)
+    source_path = probe_dir / "qmul4_probe.cu"
+    object_path = probe_dir / f"qmul4_probe.{os.getpid()}.o"
+    _write_text_if_changed(source_path, QMUL4_PROBE_SOURCE)
+    result = subprocess.run(
+        [
+            str(nvcc),
+            "-std=c++20",
+            f"-gencode=arch=compute_{arch},code=sm_{arch}",
+            "-c",
+            str(source_path),
+            "-o",
+            str(object_path),
+        ],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    supported = result.returncode == 0
+    if supported:
+        object_path.unlink(missing_ok=True)
+    else:
+        logger.info(
+            "Selected NVCC does not accept QMUL4 for SM%s; using FP16 dequant fallback",
+            arch,
+        )
+    _write_text_if_changed(
+        result_path,
+        "supported\n" if supported else "unsupported\n",
+    )
+    return supported
+
+
+def supports_qmul4(arch: str, *, probe_root: Path) -> bool:
+    """Return whether the selected NVCC accepts the public QMUL4 PTX form for ``arch``.
+
+    The answer is recorded under ``probe_root`` per toolkit, architecture and probe source, so
+    every op that dequantizes NVFP4 shares one compile of the probe per toolchain.
+    """
+
+    nvcc = cuda_home() / "bin/nvcc"
+    probe_key = hashlib.sha256()
+    probe_key.update(str(nvcc.resolve()).encode())
+    probe_key.update(str(cuda_version()).encode())
+    probe_key.update(arch.encode())
+    probe_key.update(QMUL4_PROBE_SOURCE.encode())
+    probe_dir = probe_root / probe_key.hexdigest()[:16]
+    with build_lock(probe_dir):
+        return _probe_qmul4(nvcc, arch, probe_dir)
+
+
 __all__ = [
     "MM_SPARSE_TARGET_ARCH",
+    "QMUL4_PROBE_SOURCE",
+    "build_lock",
     "cuda_home",
     "cuda_version",
     "cutlass_root",
     "normalize_target_arch",
     "require_cuda_version",
+    "supports_qmul4",
     "target_arch",
     "torch_cuda_arch",
 ]

inference/msa_v1/attention/decode/q8kv4/jit.py

@@ -2,24 +2,24 @@
 
 from __future__ import annotations
 
-import fcntl
 import hashlib
 import logging
 import os
 import subprocess
 import sys
 import threading
 import time
-from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import cache, lru_cache
 from pathlib import Path
 
 import jinja2
 
+from inference.msa_v1._build_utils import build_lock as _build_lock
 from inference.msa_v1._build_utils import cuda_home as _cuda_home
 from inference.msa_v1._build_utils import cutlass_root as _cutlass_root
 from inference.msa_v1._build_utils import require_cuda_version
+from inference.msa_v1._build_utils import supports_qmul4 as _probe_supports_qmul4
 from inference.msa_v1._build_utils import target_arch as _resolve_target_arch
 
 logger = logging.getLogger(__name__)
@@ -35,20 +35,6 @@
 }
 _QMUL4_DEQUANT = "qmul4"
 _FP16_DEQUANT = "fp16_fallback"
-_QMUL4_PROBE_SOURCE = r"""
-#include <cstdint>
-
-__global__ void qmul4_probe(uint32_t* output) {
-  uint32_t result;
-  uint16_t packed_e2m1 = 0;
-  uint32_t scale_e4m3 = 0;
-  asm volatile(
-      "mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4 %0, %1, %2;"
-      : "=r"(result)
-      : "h"(packed_e2m1), "r"(scale_e4m3));
-  output[0] = result;
-}
-""".lstrip()
 
 
 def _write_text_if_changed(path: Path, content: str) -> None:
@@ -57,18 +43,6 @@ def _write_text_if_changed(path: Path, content: str) -> None:
     path.write_text(content, encoding="utf-8")
 
 
-@contextmanager
-def _build_lock(cache_dir: Path):
-    """Keep a shared library immutable while another rank builds or loads it."""
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    with (cache_dir / "build.lock").open("a") as lock_file:
-        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
-        try:
-            yield
-        finally:
-            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
-
-
 @lru_cache(maxsize=1)
 def _cuda_version() -> tuple[int, int]:
     return require_cuda_version(
@@ -141,53 +115,7 @@ def _cache_root() -> Path:
 def _supports_qmul4(arch: str) -> bool:
     """Return whether the selected NVCC accepts the public QMUL4 PTX form."""
 
-    nvcc = _cuda_home() / "bin/nvcc"
-    probe_key = hashlib.sha256()
-    probe_key.update(str(nvcc.resolve()).encode())
-    probe_key.update(str(_cuda_version()).encode())
-    probe_key.update(arch.encode())
-    probe_key.update(_QMUL4_PROBE_SOURCE.encode())
-    probe_dir = _cache_root() / "capability_probes" / probe_key.hexdigest()[:16]
-    with _build_lock(probe_dir):
-        return _probe_qmul4(nvcc, arch, probe_dir)
-
-
-def _probe_qmul4(nvcc: Path, arch: str, probe_dir: Path) -> bool:
-    result_path = probe_dir / "qmul4.result"
-    if result_path.is_file():
-        return result_path.read_text(encoding="utf-8").strip() == "supported"
-
-    probe_dir.mkdir(parents=True, exist_ok=True)
-    source_path = probe_dir / "qmul4_probe.cu"
-    object_path = probe_dir / f"qmul4_probe.{os.getpid()}.o"
-    _write_text_if_changed(source_path, _QMUL4_PROBE_SOURCE)
-    result = subprocess.run(
-        [
-            str(nvcc),
-            "-std=c++20",
-            f"-gencode=arch=compute_{arch},code=sm_{arch}",
-            "-c",
-            str(source_path),
-            "-o",
-            str(object_path),
-        ],
-        capture_output=True,
-        text=True,
-        check=False,
-    )
-    supported = result.returncode == 0
-    if supported:
-        object_path.unlink(missing_ok=True)
-    else:
-        logger.info(
-            "Selected NVCC does not accept QMUL4 for SM%s; using FP16 dequant fallback",
-            arch,
-        )
-    _write_text_if_changed(
-        result_path,
-        "supported\n" if supported else "unsupported\n",
-    )
-    return supported
+    return _probe_supports_qmul4(arch, probe_root=_cache_root() / "capability_probes")
 
 
 @cache