mxflashattn 1.4.0__py3-none-any.whl

benchmarks/report.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import argparse
+import json
+from statistics import median
+from pathlib import Path
+
+
+def render_markdown(run: dict) -> str:
+    cases = run.get("cases", [])
+    metadata = run.get("metadata", {})
+    fallbacks = [case for case in cases if case.get("fallback_reason")]
+    latencies = [float(case["latency_ms"]) for case in cases if case.get("latency_ms") is not None]
+    errors = [float(case["max_abs_error"]) for case in cases if case.get("max_abs_error") is not None]
+    relative_errors = [
+        float(case["max_relative_error"]) for case in cases if case.get("max_relative_error") is not None
+    ]
+    comparable = [
+        case
+        for case in cases
+        if case.get("latency_ms") is not None and case.get("baseline_latency_ms") is not None
+    ]
+    reductions = [
+        (float(case["baseline_latency_ms"]) - float(case["latency_ms"]))
+        / float(case["baseline_latency_ms"])
+        * 100
+        for case in comparable
+        if float(case["baseline_latency_ms"]) > 0
+    ]
+    decode_cases = [case for case in comparable if case.get("q_len") == 1]
+    decode_reductions = [
+        (float(case["baseline_latency_ms"]) - float(case["latency_ms"]))
+        / float(case["baseline_latency_ms"])
+        * 100
+        for case in decode_cases
+        if float(case["baseline_latency_ms"]) > 0
+    ]
+    decode_target_cases = sum(value >= 20 for value in decode_reductions)
+    baseline_memory = [
+        int(case["baseline_peak_memory_bytes"])
+        for case in cases
+        if case.get("baseline_peak_memory_bytes") is not None
+    ]
+    candidate_memory = [
+        int(case["peak_memory_bytes"])
+        for case in cases
+        if case.get("peak_memory_bytes") is not None
+    ]
+    lines = [
+        "# MXFlashAttn Benchmark Report",
+        "",
+        f"- Device: {metadata.get('device_name', metadata.get('device', 'unknown'))}",
+        f"- Device total memory: {_memory(metadata.get('device_total_memory_bytes'))}",
+        f"- PyTorch: {metadata.get('torch', 'unknown')}",
+        f"- MXMACA runtime / driver: {metadata.get('runtime', 'unknown')} / {metadata.get('driver', 'unknown')}",
+        f"- FlashAttention / vLLM: {metadata.get('flash_attn', 'unknown')} / {metadata.get('vllm', 'unknown')}",
+        f"- Seed / warmup / repeats: {metadata.get('seed', 'unknown')} / {metadata.get('warmup', 'unknown')} / {metadata.get('repeats', 'unknown')}",
+        f"- Cases: {len(cases)}",
+        f"- Reference fallback cases: {len(fallbacks)}",
+    ]
+    if latencies:
+        lines.extend(["", f"- Median attention latency: {median(latencies):.4f} ms"])
+    if reductions:
+        lines.append(f"- Median latency reduction vs reference: {median(reductions):.2f}%")
+    if decode_reductions:
+        lines.append(
+            f"- Decode (q_len=1) median reduction: {median(decode_reductions):.2f}% "
+            f"({decode_target_cases}/{len(decode_reductions)} cases at or above 20%)"
+        )
+    if errors:
+        lines.append(f"- Maximum absolute error: {max(errors):.6g}")
+    if relative_errors:
+        lines.append(
+            f"- Maximum relative error (reference denominator floor 1e-6): {max(relative_errors):.6g}"
+        )
+    if baseline_memory and candidate_memory:
+        lines.append(
+            f"- Maximum measured peak memory: {_memory(max(baseline_memory))} reference, "
+            f"{_memory(max(candidate_memory))} candidate"
+        )
+    if not cases or len(fallbacks) == len(cases):
+        lines.extend(["", "**Performance target is not verified.** Every measured candidate used the reference fallback."])
+    else:
+        lines.extend(["", "Native backend measurements are present; per-case results include regressions and improvements."])
+        if decode_reductions and decode_target_cases != len(decode_reductions):
+            lines.append("The 20% decode target is not met by every measured decode case.")
+    lines.extend(
+        [
+            "",
+            "| Case | Layout | Dtype | Backend | Reference (ms) | Candidate (ms) | Reduction | Decode (ms) | Max abs error | Max rel error | Peak memory ref/candidate | Fallback reason |",
+            "|---|---|---|---|---:|---:|---:|---:|---:|---:|---|---|",
+        ]
+    )
+    for case in cases:
+        baseline_latency = case.get("baseline_latency_ms")
+        latency = case.get("latency_ms")
+        reduction = None
+        if baseline_latency is not None and latency is not None and float(baseline_latency) > 0:
+            reduction = (float(baseline_latency) - float(latency)) / float(baseline_latency) * 100
+        lines.append(
+            "| {case_id} | {layout} | {dtype} | {backend} | {baseline} | {latency} | {reduction} | {decode} | {error} | {relative_error} | {memory} | {fallback} |".format(
+                case_id=case.get("case_id", ""),
+                layout=case.get("layout", ""),
+                dtype=case.get("dtype", ""),
+                backend=case.get("backend", ""),
+                baseline=_format(baseline_latency),
+                latency=_format(case.get("latency_ms")),
+                reduction="" if reduction is None else f"{reduction:.2f}%",
+                decode=_format(case.get("decode_latency_ms")),
+                error=_format(case.get("max_abs_error")),
+                relative_error=_format(case.get("max_relative_error")),
+                memory=f"{_memory(case.get('baseline_peak_memory_bytes'))} / {_memory(case.get('peak_memory_bytes'))}",
+                fallback=case.get("fallback_reason") or "",
+            )
+        )
+    return "\n".join(lines) + "\n"
+
+
+def _format(value: object) -> str:
+    return "" if value is None else f"{float(value):.6g}"
+
+
+def _memory(value: object) -> str:
+    return "" if value is None else f"{int(value) / (1024**2):.2f} MiB"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Render a benchmark JSON report as Markdown")
+    parser.add_argument("input", type=Path)
+    parser.add_argument("--output", type=Path, default=Path("artifacts/benchmark.md"))
+    args = parser.parse_args()
+    data = json.loads(args.input.read_text(encoding="utf-8"))
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(render_markdown(data), encoding="utf-8")
+    print(f"Wrote benchmark report to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/run.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import itertools
+import json
+import os
+import re
+import statistics
+import subprocess
+import time
+from pathlib import Path
+
+import torch
+import yaml
+
+from mxflashattn import flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache, get_last_dispatch_info
+from mxflashattn import reference
+
+
+def expand_cases(config: dict) -> list[dict]:
+    matrix = config["matrix"]
+    keys = list(matrix)
+    return [
+        dict(zip(keys, values))
+        for values in itertools.product(*(matrix[key] for key in keys))
+    ]
+
+
+def _case_id(case: dict) -> str:
+    return "_".join(f"{key}-{value}" for key, value in case.items())
+
+
+def _sync(device: torch.device) -> None:
+    module_names = {"maca": "maca", "musa": "musa", "cuda": "cuda"}
+    module = getattr(torch, module_names.get(device.type, ""), None)
+    synchronize = getattr(module, "synchronize", None)
+    if synchronize is not None:
+        synchronize(device)
+
+
+def _peak_memory(device: torch.device) -> int | None:
+    module = getattr(torch, {"maca": "maca", "musa": "musa", "cuda": "cuda"}.get(device.type, ""), None)
+    getter = getattr(module, "max_memory_allocated", None)
+    if getter is None:
+        return None
+    try:
+        return int(getter(device))
+    except (RuntimeError, TypeError):
+        return None
+
+
+def _reset_peak_memory(device: torch.device) -> None:
+    module = getattr(torch, {"maca": "maca", "musa": "musa", "cuda": "cuda"}.get(device.type, ""), None)
+    reset = getattr(module, "reset_peak_memory_stats", None)
+    if reset is not None:
+        try:
+            reset(device)
+        except (RuntimeError, TypeError):
+            pass
+
+
+def _make_inputs(case: dict, device: torch.device):
+    dtype = getattr(torch, case["dtype"])
+    batch = int(case["batch_size"])
+    query_len = int(case["q_len"])
+    kv_len = int(case["kv_len"])
+    head_dim = int(case["head_dim"])
+    kv_heads = 4
+    query_heads = kv_heads * int(case["gqa_ratio"])
+    q = torch.randn(batch, query_len, query_heads, head_dim, device=device, dtype=dtype)
+    k = torch.randn(batch, kv_len, kv_heads, head_dim, device=device, dtype=dtype)
+    v = torch.randn_like(k)
+    return q, k, v
+
+
+def _prepare_candidate(case: dict, q, k, v):
+    layout = case["layout"]
+    causal = bool(case.get("causal", True))
+    if layout == "dense":
+        return lambda: flash_attn_func(q, k, v, causal=causal)
+    if layout == "varlen":
+        batch, q_len = q.shape[:2]
+        kv_len = k.shape[1]
+        cu_q = torch.arange(0, (batch + 1) * q_len, q_len, dtype=torch.int32, device=q.device)
+        cu_k = torch.arange(0, (batch + 1) * kv_len, kv_len, dtype=torch.int32, device=q.device)
+        flat_q, flat_k, flat_v = (
+            q.reshape(-1, *q.shape[2:]),
+            k.reshape(-1, *k.shape[2:]),
+            v.reshape(-1, *v.shape[2:]),
+        )
+        return lambda: flash_attn_varlen_func(
+            flat_q,
+            flat_k,
+            flat_v,
+            cu_q,
+            cu_k,
+            max_seqlen_q=q_len,
+            max_seqlen_k=kv_len,
+            causal=causal,
+        )
+    if layout == "paged":
+        batch, kv_len = k.shape[:2]
+        page_size = int(case.get("page_size", 16))
+        pages_per_batch = (kv_len + page_size - 1) // page_size
+        cache_k = torch.zeros(batch * pages_per_batch, page_size, *k.shape[2:], dtype=k.dtype, device=k.device)
+        cache_v = torch.zeros_like(cache_k)
+        table = torch.arange(batch * pages_per_batch, dtype=torch.int32, device=k.device).reshape(batch, pages_per_batch)
+        for batch_index in range(batch):
+            start = batch_index * pages_per_batch
+            flat_k = k[batch_index]
+            flat_v = v[batch_index]
+            for page_index in range(pages_per_batch):
+                source_start = page_index * page_size
+                count = min(page_size, kv_len - source_start)
+                cache_k[start + page_index, :count] = flat_k[source_start : source_start + count]
+                cache_v[start + page_index, :count] = flat_v[source_start : source_start + count]
+        lengths = torch.full((batch,), kv_len, dtype=torch.int32, device=k.device)
+        return lambda: flash_attn_with_kvcache(
+            q,
+            cache_k,
+            cache_v,
+            cache_seqlens=lengths,
+            block_table=table,
+            causal=causal,
+        )
+    raise ValueError(f"unknown layout: {layout}")
+
+
+def _reference(case: dict, q, k, v):
+    causal = bool(case.get("causal", True))
+    if case["layout"] in ("dense", "paged"):
+        return reference.attention_batched(q, k, v, causal=causal)
+    batch, q_len = q.shape[:2]
+    kv_len = k.shape[1]
+    cu_q = [index * q_len for index in range(batch + 1)]
+    cu_k = [index * kv_len for index in range(batch + 1)]
+    return reference.attention_varlen(q.reshape(-1, *q.shape[2:]), k.reshape(-1, *k.shape[2:]), v.reshape(-1, *v.shape[2:]), cu_q, cu_k, causal=causal)
+
+
+def _measure(fn, warmup: int, repeats: int, device: torch.device) -> tuple[float, object]:
+    result = None
+    for _ in range(warmup):
+        result = fn()
+    _sync(device)
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        result = fn()
+        _sync(device)
+        samples.append((time.perf_counter() - start) * 1000)
+    return statistics.median(samples), result
+
+
+def run_benchmarks(
+    config: dict,
+    *,
+    device: torch.device,
+    warmup: int,
+    repeats: int,
+    limit: int | None = None,
+    seed: int | None = None,
+) -> dict:
+    seed = int(config.get("seed", 0) if seed is None else seed)
+    if seed < 0:
+        raise ValueError("seed must be non-negative")
+    torch.manual_seed(seed)
+    cases = expand_cases(config)
+    if limit is not None:
+        cases = cases[:limit]
+    output_cases = []
+    for case in cases:
+        q, k, v = _make_inputs(case, device)
+        _reset_peak_memory(device)
+        baseline_latency, baseline_output = _measure(lambda: _reference(case, q, k, v), warmup, repeats, device)
+        baseline_peak_memory = _peak_memory(device)
+        baseline_output = baseline_output.detach().cpu()
+        candidate = _prepare_candidate(case, q, k, v)
+        _reset_peak_memory(device)
+        candidate_latency, candidate_output = _measure(candidate, warmup, repeats, device)
+        dispatch_info = get_last_dispatch_info()
+        candidate_output = candidate_output.detach().cpu()
+        difference = (candidate_output.float() - baseline_output.float()).abs()
+        max_abs_error = float(difference.max().item()) if difference.numel() else 0.0
+        denom = baseline_output.float().abs().clamp_min(1e-6)
+        max_relative_error = float((difference / denom).max().item()) if difference.numel() else 0.0
+        query_tokens = int(q.shape[0] * q.shape[1])
+        output_cases.append(
+            {
+                **case,
+                "case_id": _case_id(case),
+                "backend": dispatch_info.backend if dispatch_info else "unknown",
+                "fallback_reason": dispatch_info.fallback_reason if dispatch_info else None,
+                "max_abs_error": max_abs_error,
+                "max_relative_error": max_relative_error,
+                "baseline_latency_ms": baseline_latency,
+                "latency_ms": candidate_latency,
+                "attention_tokens_per_second": query_tokens / (candidate_latency / 1000),
+                "decode_latency_ms": candidate_latency if q.shape[1] == 1 else None,
+                "first_token_latency_ms": None,
+                "baseline_peak_memory_bytes": baseline_peak_memory,
+                "peak_memory_bytes": _peak_memory(device),
+            }
+        )
+    return {
+        "schema_version": 1,
+        "metadata": {
+            "device": str(device),
+            "device_name": _device_name(device),
+            "device_total_memory_bytes": _device_total_memory(device),
+            "torch": torch.__version__,
+            "runtime": _runtime_version(),
+            "driver": _driver_version(),
+            "flash_attn": _optional_distribution_version("flash-attn"),
+            "vllm": _optional_version("vllm"),
+            "warmup": warmup,
+            "repeats": repeats,
+            "seed": seed,
+            "fallback_enabled": os.getenv("MXFLASHATTN_ALLOW_FALLBACK") == "1",
+        },
+        "cases": output_cases,
+    }
+
+
+def _device_name(device: torch.device) -> str:
+    if device.type == "cpu":
+        return "CPU reference"
+    for module_name in (device.type, "cuda"):
+        module = getattr(torch, module_name, None)
+        getter = getattr(module, "get_device_name", None)
+        if getter is not None:
+            try:
+                return str(getter(device))
+            except (RuntimeError, TypeError):
+                pass
+    return str(device)
+
+
+def _device_total_memory(device: torch.device) -> int | None:
+    if device.type == "cpu":
+        return None
+    for module_name in (device.type, "cuda"):
+        module = getattr(torch, module_name, None)
+        getter = getattr(module, "get_device_properties", None)
+        if getter is not None:
+            try:
+                return int(getter(device).total_memory)
+            except (RuntimeError, TypeError, AttributeError, AssertionError):
+                pass
+    return None
+
+
+def _runtime_version() -> str | None:
+    for name in ("MXMACA_VERSION", "MACA_VERSION", "MUSA_VERSION"):
+        if os.getenv(name):
+            return os.environ[name]
+    return _mx_smi_version(r"MACA Version:\s*([^\s]+)")
+
+
+def _driver_version() -> str | None:
+    for name in ("MXMACA_DRIVER_VERSION", "MACA_DRIVER_VERSION", "MUSA_DRIVER_VERSION"):
+        if os.getenv(name):
+            return os.environ[name]
+    return _mx_smi_version(r"Kernel Mode Driver Version:\s*([^\s]+)")
+
+
+def _mx_smi_version(pattern: str) -> str | None:
+    try:
+        result = subprocess.run(
+            ["mx-smi"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=False,
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    match = re.search(pattern, result.stdout)
+    return match.group(1) if match else None
+
+
+def _optional_version(module_name: str) -> str | None:
+    try:
+        return importlib.metadata.version(module_name.replace("_", "-"))
+    except importlib.metadata.PackageNotFoundError:
+        try:
+            module = importlib.import_module(module_name)
+        except (ImportError, OSError):
+            return None
+    return str(getattr(module, "__version__", "unknown"))
+
+
+def _optional_distribution_version(distribution: str) -> str | None:
+    try:
+        return importlib.metadata.version(distribution)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run MXFlashAttn correctness and latency cases")
+    parser.add_argument("--config", type=Path, default=Path("benchmarks/configs/c500.yaml"))
+    parser.add_argument("--output", type=Path, default=Path("artifacts/benchmark.json"))
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--repeats", type=int, default=10)
+    parser.add_argument("--limit", type=int, default=36)
+    parser.add_argument("--seed", type=int)
+    args = parser.parse_args()
+    if args.warmup < 0 or args.repeats < 1 or args.limit < 1:
+        parser.error("warmup must be non-negative; repeats and limit must be positive")
+    with args.config.open(encoding="utf-8") as stream:
+        config = yaml.safe_load(stream)
+    result = run_benchmarks(
+        config,
+        device=torch.device(args.device),
+        warmup=args.warmup,
+        repeats=args.repeats,
+        limit=args.limit,
+        seed=args.seed,
+    )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(f"Wrote {len(result['cases'])} benchmark cases to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

integrations/__init__.py

@@ -0,0 +1 @@
+"""Optional upstream framework adapters."""

integrations/vllm/README.md

@@ -0,0 +1,7 @@
+# vLLM Integration Status
+
+The release candidate targets the checked C500 environment: `vllm 0.17.0` and `vllm_metax 0.17.0+gd10261.d20260409.maca3.5.3.20.torch2.8`. SGLang is not installed in the image, so the alternate integration path remains unevaluated. [`backend.py`](backend.py) contains a callable bridge and an isolated vLLM v1 registration probe for dense `[tokens, heads, head_dim]` Q/K/V tensors.
+
+A full integration must target one explicitly pinned and C500-compatible vLLM version. The local v0.5.0 probe reached the vLLM 0.17.0 v1 registration hook, and the candidate encoder-shaped path can call the public MXFlashAttn API. Decoder paged-cache execution remains delegated to vLLM because metadata construction, slot-based KV writes, block-table reads, cascade handling, and the version's `AttentionImpl` contract still require validation. No end-to-end MXFlashAttn acceleration claim is made.
+
+The generation script only collects a baseline by default when `--baseline-only` is passed. It refuses to run without that flag because decoder paged-cache execution has not been switched to MXFlashAttn; baseline output explicitly says that it did not exercise this project. The configured quick model is Qwen3-0.6B for the tested vLLM 0.17.0 environment; weights remain outside this repository.

integrations/vllm/__init__.py

@@ -0,0 +1 @@
+"""Experimental vLLM integration helpers."""

integrations/vllm/backend.py

@@ -0,0 +1,158 @@
+"""vLLM 0.17.0 registration and dense bridge for the local E2E candidate."""
+
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+
+from mxflashattn import flash_attn_func
+
+
+DISPATCH_EVENTS: list[dict[str, str]] = []
+
+# Recording runs once per attention call (28 layers x N steps x prompts). The
+# event payloads repeat heavily, so serialisation is cached. Set
+# MXFLASHATTN_DISPATCH_SAMPLE=N to write only every Nth event: the log still
+# shows which backend served the run, but the write cost leaves the hot path.
+_JSON_CACHE: dict[tuple[tuple[str, str], ...], str] = {}
+_SAMPLE_EVERY = int(os.getenv("MXFLASHATTN_DISPATCH_SAMPLE", "0") or 0)
+_event_count = 0
+
+
+def get_dispatch_events() -> list[dict[str, str]]:
+    return list(DISPATCH_EVENTS)
+
+
+def clear_dispatch_events(truncate_log: bool = True) -> None:
+    """Drop buffered events.
+
+    ``truncate_log=False`` keeps previously written lines, which is what
+    multi-prompt suites need: the in-memory window resets per prompt, but the
+    on-disk evidence must accumulate across the whole run.
+    """
+    DISPATCH_EVENTS.clear()
+    if not truncate_log:
+        return
+    path = os.getenv("MXFLASHATTN_DISPATCH_LOG")
+    if path:
+        try:
+            open(path, "w", encoding="utf-8").close()
+        except OSError:
+            pass
+
+
+def _record_event(event: dict[str, str]) -> None:
+    DISPATCH_EVENTS.append(event)
+    path = os.getenv("MXFLASHATTN_DISPATCH_LOG")
+    # EngineCore workers may import the plugin before the parent-side smoke
+    # runner's environment mutation is visible.  The plugin log is already
+    # worker-visible; derive a sibling dispatch path as a durable fallback.
+    if not path:
+        plugin_log = os.getenv("MXFLASHATTN_PLUGIN_IMPORT_LOG")
+        if plugin_log:
+            path = plugin_log.removesuffix(".plugin.log") + ".dispatch.jsonl"
+    if not path:
+        return
+    if _SAMPLE_EVERY > 0:
+        # Sampling keeps the evidence trail (which backend served which call)
+        # while removing the write cost from the measured hot path.
+        global _event_count
+        _event_count += 1
+        if _event_count % _SAMPLE_EVERY != 1:
+            return
+    try:
+        key = tuple(sorted(event.items()))
+        rendered = _JSON_CACHE.get(key)
+        if rendered is None:
+            rendered = json.dumps(event, sort_keys=True) + "\n"
+            _JSON_CACHE[key] = rendered
+        with open(path, "a", encoding="utf-8") as stream:
+            stream.write(rendered)
+    except OSError:
+        pass
+
+
+def attention_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    causal: bool,
+    scale: float,
+) -> torch.Tensor:
+    """Adapt dense vLLM token/head tensors to the MXFlashAttn public API."""
+    if query.ndim != 3 or key.ndim != 3 or value.shape != key.shape:
+        raise ValueError("vLLM Q/K/V must use [tokens, heads, head_dim] layout")
+    output = flash_attn_func(
+        query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0),
+        causal=causal, softmax_scale=scale,
+    )
+    return output.squeeze(0)
+
+
+def register_vllm_v1_backend() -> dict[str, str]:
+    """Register an importable backend path for vLLM 0.17.0 workers."""
+    try:
+        from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend
+    except Exception as exc:  # pragma: no cover - depends on installed vLLM
+        return {"registered": "false", "reason": f"vllm_v1_api_unavailable: {type(exc).__name__}: {exc}"}
+    class_path = "integrations.vllm.vllm_plugin.MXFlashAttnBackend"
+
+    def apply_override() -> None:
+        register_backend(AttentionBackendEnum.FLASH_ATTN, class_path=class_path)
+
+    apply_override()
+    # The MetaX selector may intentionally return its own backend class even
+    # after registry overrides.  Replace only its implementation factory so
+    # the vendor metadata builder and constructor remain intact while the
+    # decoder implementation is ours.
+    try:
+        from integrations.vllm.vllm_plugin import MXFlashAttnImpl
+        import vllm_metax.v1.attention.backends.flash_attn as metax_flash_attn
+
+        metax_flash_attn.MacaFlashAttentionBackend.get_impl_cls = staticmethod(
+            lambda: MXFlashAttnImpl
+        )
+    except Exception:
+        pass
+    # The MetaX vLLM plugin re-registers FLASH_ATTN during backend selection.
+    # Reapply our importable path after that platform hook runs in each worker.
+    try:
+        import vllm_metax.platform as metax_platform
+
+        original = metax_platform.register_attention_backends
+        if not getattr(original, "_mxflashattn_wrapped", False):
+            def wrapped_register_attention_backends() -> None:
+                original()
+                apply_override()
+
+            wrapped_register_attention_backends._mxflashattn_wrapped = True
+            metax_platform.register_attention_backends = wrapped_register_attention_backends
+
+        # MetaX's selector calls its registration hook immediately before
+        # resolving the class path.  In some worker start modes that hook is
+        # bound before the wrapper above is installed, so enforce the path at
+        # the final selector boundary as well.
+        platform_base = getattr(metax_platform, "MacaPlatformBase", None)
+        descriptor = getattr(platform_base, "__dict__", {}).get("get_attn_backend_cls")
+        selector = descriptor.__func__ if isinstance(descriptor, classmethod) else descriptor
+        if selector is not None and not getattr(selector, "_mxflashattn_wrapped", False):
+            def wrapped_selector(cls, selected_backend, attn_selector_config, num_heads=None):
+                if getattr(selected_backend, "name", None) == "FLASH_ATTN":
+                    apply_override()
+                    return class_path
+                return selector(cls, selected_backend, attn_selector_config, num_heads)
+
+            wrapped_selector._mxflashattn_wrapped = True
+            platform_base.get_attn_backend_cls = classmethod(wrapped_selector)
+    except Exception:
+        pass
+    return {
+        "registered": "true",
+        "backend_name": "MXFLASHATTN_V017",
+        "backend_path": class_path,
+        "decoder_path": "vendor_direct_with_aten_correctness_fallback",
+        "encoder_path": "mxflashattn_public_api",
+    }

integrations/vllm/compat.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import importlib.metadata
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class VllmCompatibility:
+    installed_version: str | None
+    status: str
+    reason: str
+    supported: bool
+
+
+def detect_vllm_version() -> str | None:
+    try:
+        return importlib.metadata.version("vllm")
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def classify_vllm(version: str | None) -> VllmCompatibility:
+    if version is None:
+        return VllmCompatibility(None, "not_installed", "vLLM is not installed", False)
+    normalized = version.split("+", 1)[0]
+    parts = normalized.split(".")
+    try:
+        major, minor = int(parts[0]), int(parts[1])
+    except (IndexError, ValueError):
+        return VllmCompatibility(version, "unknown", "cannot parse vLLM version", False)
+
+    if (major, minor) == (0, 17):
+        return VllmCompatibility(version, "supported", "v1.3 backend contract", True)
+    if (major, minor) == (0, 30):
+        return VllmCompatibility(
+            version,
+            "probe_required",
+            "vLLM 0.30 API must be inspected and tested in the target environment",
+            False,
+        )
+    return VllmCompatibility(
+        version,
+        "unsupported",
+        "no validated MXFlashAttn integration contract for this vLLM version",
+        False,
+    )