延迟导入 FlashMLA 扩展

[ghangz]

该 PR 将 FlashMLA 扩展导入延后到真正执行测试时，使命令行参数解析、环境诊断和文档示例在扩展未编译时仍可运行。

这个修改面向沐曦 GPU 适配场景中比较容易影响开发、构建或验证稳定性的环节，把原来需要人工排查的问题前移到工具链、运行前检查或基准脚本中处理。实现上保持对现有默认行为的兼容，只在检测到明确配置、输入或环境异常时给出更直接的诊断，避免引入额外运行依赖，也方便维护者独立审阅该分支。

已在沐曦算力环境中完成对应分支验证，验证记录包含真实运行日志、命令输出和失败路径检查，本地归档目录为：E:/Documents/muxi/测试报告/FlashMLA_real_maca_validation_20260608。提交分支：mengz/lazy-extension-import，目标仓库：MetaX-MACA/FlashMLA。

[gemini-code-assist]

Code Review

This pull request introduces lazy loading for the flash_mla package, allowing it to be imported even when the compiled CUDA extension is not available, and adds a corresponding test. The reviewer recommended caching the dynamically loaded functions to avoid performance overhead in high-frequency execution paths, and explicitly defining the function signatures instead of using *args, **kwargs to preserve IDE type hints and autocompletion.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

问题分析

1. 性能开销（Efficiency）：在当前的设计中，每次调用 get_mla_metadata 或 flash_mla_with_kvcache 都会执行 _load_interface()。虽然 Python 会缓存已导入的模块，但在高频调用的 CUDA 算子核心路径（hot path）中，频繁执行 try-except 和 import 语句仍会带来不必要的 CPU 额外开销。
2. 易用性与类型提示（Maintainability）：使用 *args, **kwargs 包装函数会导致 IDE（如 VS Code、PyCharm）和静态类型检查工具（如 mypy）无法获取原始函数的参数签名和 Docstring，降低了开发体验。

解决方案

建议使用模块级缓存来避免重复导入的开销。同时，为了保留参数签名和文档，可以直接在包装函数中显式声明参数。

_funcs_cache = None


def _load_interface():
    global _funcs_cache
    if _funcs_cache is None:
        try:
            from flash_mla.flash_mla_interface import flash_mla_with_kvcache, get_mla_metadata
            _funcs_cache = (get_mla_metadata, flash_mla_with_kvcache)
        except ImportError as exc:
            raise ImportError(
                "flash_mla_cuda is not available. Build and install FlashMLA from source "
                "in a configured MACA environment before calling FlashMLA kernels."
            ) from exc
    return _funcs_cache


def get_mla_metadata(cache_seqlens, num_heads_per_head_k: int, num_heads_k: int):
    metadata_func, _ = _load_interface()
    return metadata_func(cache_seqlens, num_heads_per_head_k, num_heads_k)


def flash_mla_with_kvcache(
    q,
    k_cache,
    block_table,
    cache_seqlens,
    head_dim_v: int,
    tile_scheduler_metadata,
    num_splits,
    softmax_scale=None,
    causal: bool = False,
):
    _, flash_func = _load_interface()
    return flash_func(
        q,
        k_cache,
        block_table,
        cache_seqlens,
        head_dim_v,
        tile_scheduler_metadata,
        num_splits,
        softmax_scale,
        causal,
    )