fix(tf-kernel): harden builds and enable SM90 attention

Refine architecture-specific wheel builds, runtime ABI validation, GPU capability dispatch, isolated wheel checks, and manual packaging documentation.

Fix the Hopper SageAttention v2 build by disabling relocatable device code for SM90, add synchronous CUDA launch diagnostics, restore the validated SM90 backend, and prefer tf-kernel from TeleFuser ops with standalone SageAttention fallback. Preserve causal and LSE behavior through the public attention interface.

Verification: tf-kernel wheel smoke tests; 384-case focused SM90 SageAttention matrix; architecture and symbol validation; pytest tests/unit/ops/test_attention_backends.py tests/unit/ops/test_tf_kernel_sage_attention.py -q; pytest tests/unit/ops with non-GPU markers; ruff check; ruff format --check; git diff --check.

Author: lzx1413

CLAUDE.md

@@ -81,6 +81,10 @@ telefuser/
   dependency to published project metadata.
 - Do not add GitHub Actions workflows that compile or publish `tf-kernel`. Kernel wheels require an explicitly
   provisioned CUDA/NVCC build host and manual validation.
+- Load only the wheel extension matching the single visible GPU architecture family. Keep build/runtime compatibility
+  facts in `tf_kernel._build_info` and validated SageAttention dispatch in `tf_kernel.capabilities`.
+- `telefuser.ops.attention` prefers the optional `tf_kernel` SageAttention backend and falls back to the standalone
+  `sageattention` package. Model code must continue calling the public ops layer rather than either package directly.
 
 ### LingBot Streaming State
 

docs/en/attention.md

@@ -207,7 +207,7 @@ pipe_config.dit_config.attention_config = config
 | `FLASH_ATTN_2` | Flash Attention 2 | `flash_attn` package |
 | `FLASH_ATTN_3` | Flash Attention 3 | `flash_attn_interface` package |
 | `FLASH_ATTN_4` | Flash Attention 4 | `flash_attn` package (with `cute` submodule) |
-| `SAGE_ATTN_*` | SageAttention variants | `sageattention` package |
+| `SAGE_ATTN_*` | SageAttention variants | `tf-kernel` or `sageattention` package |
 | `SPARGE_ATTN` | Sparge Attention | `spas_sage_attn` package |
 
 **Note on Flash Attention 4**: Flash Attention 4 is optimized for **Hopper (SM90, H100)** and **Blackwell (SM100+, B100/B200)** GPUs. It provides significant performance improvements on these architectures. For older GPUs (Ampere, Ada Lovelace), use Flash Attention 2 or 3 instead.
@@ -394,8 +394,8 @@ Source builds require Python 3.10+, PyTorch 2.11.0, CUDA Toolkit 12.8+, and CMak
 on SM100+. See the [tf-kernel installation and usage guide](./tf_kernel.md) for the complete compatibility matrix,
 H100 build command, installation verification, usage examples, and troubleshooting.
 
-The currently validated H100 wheel has a known `misaligned address` failure in the architecture-selected SM90
-SageAttention path. Use another attention backend until the focused SM90 GPU test passes on the deployed wheel.
+The SM90 path is validated on H100 through both tf-kernel and the TeleFuser public ops layer. The focused GPU test
+synchronizes the kernel and covers causal output and log-sum-exp results.
 
 ### Checking Available Backends
 

docs/en/tf_kernel.md

@@ -23,17 +23,19 @@ are implemented. Install `tf-kernel` when a pipeline uses one of its optimized C
 | CMake | 3.26 or newer for source builds |
 | GPU targets | SM80, SM90, and SM100 |
 
-Kernel availability depends on the selected build target. FP4 kernels require Blackwell (SM100 or newer); seeing
-`no fp4 operator available` on Ampere or Hopper is expected. Core operations are currently validated with Python
-3.11, PyTorch 2.11.0+cu128, CUDA 12.8, and H100 (SM90a). Other targets and operation families should be validated on
-their target GPU before production use.
+Kernel availability depends on the selected build target. FP4 kernels require Blackwell (SM100 or newer), and
+`tf_kernel.FP4_AVAILABLE` is false without an import-time warning on Ampere and Hopper. Core operations are currently
+validated with Python 3.11, PyTorch 2.11.0+cu128, CUDA 12.8, and H100 (SM90a). Other targets and operation families
+should be validated on their target GPU before production use.
 
-!!! warning "Current H100 SageAttention limitation"
+The wheel records and verifies its PyTorch public version, PyTorch CUDA version, C++11 ABI, and target GPU family at
+import. A process exposing GPUs from different architecture families is rejected. SageAttention v2 auto-dispatch is
+enabled only for SM80, SM86, SM89, SM90, SM120, and SM121.
 
-    In the currently validated H100 build, the architecture-selected `tf_kernel.sageattn()` path chooses the
-    SM90-specific FP8 kernel and can fail with `CUDA error: misaligned address`. RMSNorm, fused activations, FP8
-    quantization, and the generic FP8 SageAttention path pass smoke tests. Do not enable the SM90-specific SageAttention
-    backend in production until its focused GPU test passes on the wheel being deployed.
+!!! note "H100 SageAttention dispatch"
+
+    On H100, `tf_kernel.sageattn()` selects the validated SM90 FP8 implementation. TeleFuser uses the same kernel
+    when `SAGE_ATTN_2_8_8_SM90` is configured and `tf-kernel` is available.
 
 ## Build and install from source
 
@@ -55,6 +57,9 @@ The local build is independent of the TeleFuser installation. Make builds a corr
 into `PYTHON`. Direct `pip install .` and `pip install -e .` source builds fail with instructions to use Make; pip
 package-index installation is not available.
 
+Local builds use a `linux_*` platform tag. The container build may use `manylinux_2_28` only after checking every
+shared object's GLIBC symbol versions against that policy.
+
 For a reproducible target-specific build:
 
 ```bash
@@ -111,6 +116,10 @@ PY
 An H100-specific wheel should load its common extension from an `sm90` package directory. Also run
 `python -m pip check` to expose dependency conflicts in the environment.
 
+For development validation, run `make test-cpu`, `make test-smoke`, and `make test-wheel`. The smoke and GPU targets
+install the wheel into an isolated temporary directory before collecting tests. `make test` is the bounded GPU suite;
+reserve the 6,000+ case `make test-full` matrix for a dedicated validation host.
+
 ## Usage
 
 TeleFuser users should call the public ops layer; it selects `tf-kernel` for supported eager CUDA paths and keeps the
@@ -137,18 +146,18 @@ x = torch.randn(8, 1024, device="cuda", dtype=torch.float16)
 weight = torch.ones(1024, device="cuda", dtype=torch.float16)
 y = tf_kernel.rmsnorm(x, weight, eps=1e-6)
 
-# H100-tested generic FP8 SageAttention path.
+# H100-tested SM90 FP8 SageAttention path.
 # HND layout: [batch, heads, sequence, head_dim]
 q = torch.randn(1, 8, 128, 64, device="cuda", dtype=torch.float16)
 k = torch.randn_like(q)
 v = torch.randn_like(q)
-attn_output = tf_kernel.sageattn_qk_int8_pv_fp8_cuda(
+attn_output = tf_kernel.sageattn_qk_int8_pv_fp8_cuda_sm90(
     q,
     k,
     v,
     tensor_layout="HND",
     is_causal=False,
-    pv_accum_dtype="fp32",
+    pv_accum_dtype="fp32+fp32",
 )
 ```
 
@@ -186,12 +195,11 @@ Check `nvcc --version`, set `CUDA_HOME` to the CUDA 12.8+ toolkit, and put `$CUD
 Rebuild with `make build-auto` on the target machine or use the explicit `build-sm80`, `build-sm90`, or
 `build-sm100` target. Architecture-specific wheels cannot provide kernels that were omitted at build time.
 
-### SageAttention fails with `misaligned address` on H100
+### Validate an SM90 SageAttention deployment
+
+The SM90-specific kernel is enabled on H100. After building the wheel, run the synchronized tf-kernel smoke test and
+the TeleFuser public-ops GPU integration test before deploying that artifact on a new host.
 
-The architecture selector currently routes H100 to the SM90-specific FP8 implementation. Treat a failure in that
-path as an unsupported backend for the current wheel; select another TeleFuser attention implementation instead of
-continuing in the CUDA process after the asynchronous error. The generic FP8 function shown above is useful for
-isolated validation, but production enablement still requires parity and workload benchmarks.
 
 ### The build exhausts CPU or memory
 

docs/zh/attention.md

@@ -207,7 +207,7 @@ pipe_config.dit_config.attention_config = config
 | `FLASH_ATTN_2` | Flash Attention 2 | `flash_attn` 包 |
 | `FLASH_ATTN_3` | Flash Attention 3 | `flash_attn_interface` 包 |
 | `FLASH_ATTN_4` | Flash Attention 4 | `flash_attn` 包（含 `cute` 子模块） |
-| `SAGE_ATTN_*` | SageAttention 变体 | `sageattention` 包 |
+| `SAGE_ATTN_*` | SageAttention 变体 | `tf-kernel` 或 `sageattention` 包 |
 | `SPARGE_ATTN` | Sparge Attention | `spas_sage_attn` 包 |
 
 **Flash Attention 4 说明**: Flash Attention 4 针对 **Hopper (SM90, H100)** 和 **Blackwell (SM100+, B100/B200)** GPU 架构进行了优化，在这些架构上提供显著的性能提升。对于旧版 GPU（Ampere、Ada Lovelace），请使用 Flash Attention 2 或 3。
@@ -393,8 +393,8 @@ make build-sm100 PYTHON=/path/to/venv/bin/python  # Blackwell
 可用。完整兼容矩阵、H100 编译命令、安装验证、使用示例和常见问题见
 [tf-kernel 安装与使用指南](./tf_kernel.md)。
 
-当前验证的 H100 wheel 在架构选择的 SM90 SageAttention 路径存在已知 `misaligned address` 错误；部署的
-wheel 通过专项 SM90 GPU 测试前，应选择其他注意力后端。
+SM90 路径已在 H100 上通过 tf-kernel 和 TeleFuser 公共 ops 层验证。专项 GPU 测试会同步内核，并覆盖
+causal 输出和 log-sum-exp 结果。
 
 ### 检查可用后端
 

docs/zh/tf_kernel.md

@@ -22,15 +22,18 @@ TeleFuser 不安装 `tf-kernel` 也可以运行：对于已经实现回退的算
 | CMake | 从源码编译要求 3.26 或更高版本 |
 | GPU 目标 | SM80、SM90 和 SM100 |
 
-具体可用内核取决于编译目标。FP4 内核要求 Blackwell（SM100 或更高）；在 Ampere 或 Hopper 上看到
-`no fp4 operator available` 属于预期行为。目前核心算子已在 Python 3.11、PyTorch 2.11.0+cu128、
-CUDA 12.8 和 H100（SM90a）组合上验证。其他目标和算子族用于生产前仍应在目标 GPU 上验证。
+具体可用内核取决于编译目标。FP4 内核要求 Blackwell（SM100 或更高）；在 Ampere 或 Hopper 上
+`tf_kernel.FP4_AVAILABLE` 为 false，导入时不会打印提示。目前核心算子已在 Python 3.11、PyTorch
+2.11.0+cu128、CUDA 12.8 和 H100（SM90a）组合上验证。其他目标和算子族用于生产前仍应在目标 GPU 上验证。
 
-!!! warning "当前 H100 SageAttention 限制"
+wheel 导入时会校验构建时记录的 PyTorch 公共版本、PyTorch CUDA 版本、C++11 ABI 和目标 GPU 架构族。
+一个进程暴露不同架构族的 GPU 时会明确失败。SageAttention v2 自动分派目前仅支持 SM80、SM86、
+SM89、SM90、SM120 和 SM121。
 
-    当前验证的 H100 build 中，架构选择接口 `tf_kernel.sageattn()` 会进入 SM90 专用 FP8 内核，并可能报
-    `CUDA error: misaligned address`。RMSNorm、融合激活、FP8 量化和通用 FP8 SageAttention 路径已通过
-    smoke test。在部署 wheel 的专项 GPU 测试通过前，不应在生产环境启用 SM90 专用 SageAttention 后端。
+!!! note "H100 SageAttention 分派"
+
+    H100 上的 `tf_kernel.sageattn()` 会选择已验证的 SM90 FP8 实现。配置 `SAGE_ATTN_2_8_8_SM90` 且
+    `tf-kernel` 可用时，TeleFuser 会使用同一个内核。
 
 ## 从源码编译和安装
 
@@ -50,6 +53,9 @@ make build-auto PYTHON=/path/to/venv/bin/python
 本地构建独立于 TeleFuser 安装。Make 会构建带有正确 tag 的 wheel，并将其安装到 `PYTHON` 指定的解释器。
 直接执行 `pip install .` 或 `pip install -e .` 会失败并提示改用 Make；当前不提供 pip 包索引安装。
 
+本地构建使用 `linux_*` platform tag。容器构建只有在检查所有共享库的 GLIBC 符号版本满足策略后，
+才可以使用 `manylinux_2_28` tag。
+
 需要可复现的指定架构编译时：
 
 ```bash
@@ -105,6 +111,10 @@ PY
 H100 专用 wheel 的 common extension 应从 `sm90` 包目录加载。还应运行 `python -m pip check` 检查环境中的
 依赖冲突。
 
+开发验证依次运行 `make test-cpu`、`make test-smoke` 和 `make test-wheel`。smoke 与 GPU target 会先将
+wheel 安装到隔离的临时目录再收集用例。`make test` 是有界 GPU 测试；超过 6,000 条用例的
+`make test-full` 应只在专用验证机器执行。
+
 ## 使用示例
 
 TeleFuser 用户应调用公共 ops 层；它会在支持的 eager CUDA 路径选择 `tf-kernel`，同时保留框架回退：
@@ -130,18 +140,18 @@ x = torch.randn(8, 1024, device="cuda", dtype=torch.float16)
 weight = torch.ones(1024, device="cuda", dtype=torch.float16)
 y = tf_kernel.rmsnorm(x, weight, eps=1e-6)
 
-# 已在 H100 验证的通用 FP8 SageAttention 路径。
+# 已在 H100 验证的 SM90 FP8 SageAttention 路径。
 # HND 布局：[batch, heads, sequence, head_dim]
 q = torch.randn(1, 8, 128, 64, device="cuda", dtype=torch.float16)
 k = torch.randn_like(q)
 v = torch.randn_like(q)
-attn_output = tf_kernel.sageattn_qk_int8_pv_fp8_cuda(
+attn_output = tf_kernel.sageattn_qk_int8_pv_fp8_cuda_sm90(
     q,
     k,
     v,
     tensor_layout="HND",
     is_causal=False,
-    pv_accum_dtype="fp32",
+    pv_accum_dtype="fp32+fp32",
 )
 ```
 
@@ -179,11 +189,10 @@ PyTorch 版本，请在干净环境中安装 TeleFuser 和 `tf-kernel`。除非
 在目标机器使用 `make build-auto` 重新编译，或显式选择 `build-sm80`、`build-sm90`、`build-sm100`。
 指定架构 wheel 无法提供编译时未包含的内核。
 
-### H100 上 SageAttention 报 `misaligned address`
+### 验证 SM90 SageAttention 部署
 
-当前架构选择接口会把 H100 路由到 SM90 专用 FP8 实现。如果该路径失败，应将其视为当前 wheel 不支持的
-后端并选择其他 TeleFuser 注意力实现；异步 CUDA 错误发生后不要继续复用该进程。上文的通用 FP8 函数
-可用于独立验证，但生产启用仍需要完成精度对齐和目标工作负载 benchmark。
+SM90 专用内核已在 H100 上启用。构建 wheel 后，应在新的部署主机上运行带同步的 tf-kernel smoke test
+和 TeleFuser 公共 ops GPU 集成测试。
 
 ### 编译耗尽 CPU 或内存
 

telefuser/ops/attention/attention_impl.py

@@ -251,6 +251,7 @@ def attention(
                 k,
                 v,
                 attn_mask=attn_mask,
+                is_causal=is_causal,
                 sm_scale=scale,
                 tensor_layout=sage_tensor_layout,
                 pv_accum_dtype="fp32+fp32",
@@ -263,6 +264,7 @@ def attention(
                 k,
                 v,
                 attn_mask=attn_mask,
+                is_causal=is_causal,
                 sm_scale=scale,
                 tensor_layout=sage_tensor_layout,
                 pv_accum_dtype="fp32",
@@ -275,6 +277,7 @@ def attention(
                 k,
                 v,
                 attn_mask=attn_mask,
+                is_causal=is_causal,
                 sm_scale=scale,
                 tensor_layout=sage_tensor_layout,
                 pv_accum_dtype="fp32+fp32",
@@ -425,12 +428,12 @@ def _get_ring_attn_config(
     """Get attention config with LSE support for ring attention."""
     attn_impl = attention_config.attn_impl if attention_config else AttnImplType.TORCH_SDPA
 
-    if not supports_return_lse(attn_impl.value):
+    if not supports_return_lse(attn_impl.name):
         fallback = get_lse_fallback_impl()
         if fallback is None:
             raise RuntimeError("Ring attention requires LSE support. Install flash-attn or sageattention.")
         logger.info(f"Ring attention: falling back to {fallback}")
-        return AttentionConfig(attn_impl=AttnImplType(fallback), scale=scale, is_causal=is_causal)
+        return AttentionConfig(attn_impl=AttnImplType[fallback], scale=scale, is_causal=is_causal)
 
     return attention_config
 

telefuser/ops/attention/backends.py

@@ -85,16 +85,23 @@ def _try_import_sdpa() -> None:
 
 
 def _try_import_sage_attn() -> None:
-    """Import Sage Attention."""
+    """Import SageAttention, preferring the optional tf-kernel package."""
     global SAGE_ATTN_AVAILABLE, sageattention
 
-    try:
-        if importlib.util.find_spec("sageattention") is not None:
-            sageattention = importlib.import_module("sageattention")
-            SAGE_ATTN_AVAILABLE = True
-            logger.debug("Sage Attention available")
-    except (ModuleNotFoundError, ImportError):
-        pass
+    SAGE_ATTN_AVAILABLE = False
+    sageattention = None
+    for module_name in ("tf_kernel", "sageattention"):
+        try:
+            if importlib.util.find_spec(module_name) is None:
+                continue
+            sageattention = importlib.import_module(module_name)
+        except (ModuleNotFoundError, ImportError) as error:
+            logger.debug("SageAttention backend %s unavailable: %s", module_name, error)
+            continue
+
+        SAGE_ATTN_AVAILABLE = True
+        logger.debug("SageAttention loaded from %s", module_name)
+        return
 
 
 def _try_import_sparge_attn() -> None: