adapt(metax): MetaX C550 backend adaptation for vLLM 0.24.0 (#294)

## Background

This PR upgrades the MetaX C550 backend from vLLM 0.20.2 to 0.24.0,
fixing two regressions that blocked offline inference on MetaX C550
hardware.

---

## Fix 1: Missing `USE_EXP2` parameter in MetaX delta-rule kernel

**File:**
`vllm_fl/dispatch/backends/vendor/metax/patches/chunk_delta_h.py`

vLLM 0.24.0 updated the upstream fla-org
`chunk_gated_delta_rule_fwd_kernel_h_blockdim64` kernel signature to add
`USE_EXP2: tl.constexpr`. The MetaX-patched copy was not updated
accordingly, causing a `TypeError` at kernel dispatch time.

**Fix:** Add `USE_EXP2: tl.constexpr = False` to the MetaX kernel
signature. The parameter is accepted but unused — MetaX uses its own exp
implementation.

---

## Fix 2: Incomplete `torch.accelerator` memory API patch for torch
2.8+metax

**File:**
`vllm_fl/dispatch/backends/vendor/metax/patches/accelerator_compat.py`

vLLM 0.24.0 uses several `torch.accelerator` memory APIs added in
PyTorch 2.9+. The MetaX torch distribution ships as `2.8+metax` and is
missing these APIs, causing `AttributeError` during worker startup:

```
AttributeError: module 'torch.accelerator' has no attribute 'max_memory_allocated'
```

The existing `accelerator_compat.py` only patched `empty_cache` and
missed 5 other APIs.

**Fix:** Patch all 6 missing APIs to their `torch.cuda.*` equivalents:
`empty_cache`, `memory_stats`, `memory_reserved`, `memory_allocated`,
`reset_peak_memory_stats`, `max_memory_allocated`

Reference: [MetaX vLLM-metax
v0.21.0](https://github.com/MetaX-MACA/vLLM-metax/blob/releases/v0.21.0/vllm_metax/patch/torch_fix/fix_standalone_compile.py)

---

## Testing

| Model | Type | TP | dtype | max_model_len | Result |
|-------|------|----|-------|---------------|--------|
| Qwen3.6-27B | Mamba hybrid | 4 | bfloat16 | 512 | ✅ PASS |
| Qwen3.6-35B-A3B | Mamba + MoE (256 experts) | 4 | bfloat16 | 512 | ✅
PASS |

Hardware: MetaX C550 × 8, MACA 3.7.0, torch 2.8+metax

Author: physics31415926

vllm_fl/dispatch/backends/vendor/metax/patches/accelerator_compat.py

@@ -1,13 +1,73 @@
 # SPDX-License-Identifier: Apache-2.0
 # 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
 
-# -----------------------------------------------------
-# Note: torch 2.8+metax does not have torch.accelerator.empty_cache
-#       (added in PyTorch 2.10). Patch it to use torch.cuda.empty_cache.
-# _____________________________________________________
+# --------------------------------------------------------------------------------
+# Hotfix: torch 2.8+metax is missing several torch.accelerator memory APIs that
+# were added in PyTorch 2.9+. Patch them to the equivalent torch.cuda calls.
+#
+# Guard: only applied when torch version is < 2.9 AND the API is actually absent.
+# - Version check avoids silently overriding APIs that exist (and may have changed
+#   signatures) in torch >= 2.9.
+# - hasattr check is the authoritative guard: if the API exists we never touch it,
+#   regardless of version.
+#
+# Reference: https://github.com/MetaX-MACA/vLLM-metax/blob/releases/v0.21.0/vllm_metax/patch/torch_fix/fix_standalone_compile.py
+# TODO: remove when MetaX ships torch >= 2.9 with full accelerator API support.
+# --------------------------------------------------------------------------------
+
+import logging
+from typing import Tuple
 
 import torch
 
+logger = logging.getLogger(__name__)
+
+
+def _torch_version_tuple() -> Tuple[int, ...]:
+    """Return (major, minor) of torch version, ignoring vendor suffixes like '+metax'."""
+    ver = torch.__version__.split("+")[0]  # strip "+metax", "+cu124", etc.
+    parts = ver.split(".")[:2]
+    try:
+        return tuple(int(p) for p in parts)
+    except ValueError:
+        return (0, 0)
+
+
+_torch_ver = _torch_version_tuple()
+
+# Only apply this patch on torch < 2.9. On torch >= 2.9, these APIs should exist
+# natively; patching them would risk silently overriding a potentially changed
+# signature or implementation.
+if _torch_ver < (2, 9):
+    _MISSING_APIS = {
+        "empty_cache": torch.cuda.empty_cache,
+        "memory_stats": torch.cuda.memory_stats,
+        "memory_reserved": torch.cuda.memory_reserved,
+        "memory_allocated": torch.cuda.memory_allocated,
+        "reset_peak_memory_stats": torch.cuda.reset_peak_memory_stats,
+        "max_memory_allocated": torch.cuda.max_memory_allocated,
+    }
 
-if not hasattr(torch.accelerator, "empty_cache"):
-    torch.accelerator.empty_cache = torch.cuda.empty_cache
+    for _name, _impl in _MISSING_APIS.items():
+        if not hasattr(torch.accelerator, _name):
+            setattr(torch.accelerator, _name, _impl)
+            logger.debug(
+                "accelerator_compat: patched torch.accelerator.%s -> torch.cuda.%s "
+                "(torch %s lacks this API)",
+                _name, _name, torch.__version__,
+            )
+else:
+    # Sanity check: warn if any expected API is still missing on torch >= 2.9,
+    # which would indicate an unexpected regression.
+    _EXPECTED_APIS = [
+        "empty_cache", "memory_stats", "memory_reserved",
+        "memory_allocated", "reset_peak_memory_stats", "max_memory_allocated",
+    ]
+    for _name in _EXPECTED_APIS:
+        if not hasattr(torch.accelerator, _name):
+            logger.warning(
+                "accelerator_compat: torch.accelerator.%s is missing on torch %s "
+                "(expected it to exist on >= 2.9). "
+                "Please update this patch file.",
+                _name, torch.__version__,
+            )

vllm_fl/dispatch/backends/vendor/metax/patches/chunk_delta_h.py

@@ -68,6 +68,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     STORE_FINAL_STATE: tl.constexpr,
     SAVE_NEW_VALUE: tl.constexpr,
     IS_VARLEN: tl.constexpr,
+    USE_EXP2: tl.constexpr = False,  # accepted but unused on MetaX
 ):
     i_v, i_nh = tl.program_id(0), tl.program_id(1)
     i_n, i_h = i_nh // H, i_nh % H