Support W8A8 hy4-preview inference on PPU 810E

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 build/
 dist/
 vllm_fl/_version.py
+vllm_fl/dispatch/backends/vendor/thead/lib/*.so
 
 # Coverage
 .coverage

README.md

@@ -120,6 +120,25 @@ Operator adapters use the plugin dispatch manager, so backend selection,
 fallback, per-op policy, operator-list recording, and I/O diagnostics continue
 to follow the common FlagOS controls.
 
+### Optional T-Head native kernels
+
+The T-Head backend can use an optional, version-pinned native extension bundle.
+Those binaries are not stored in this repository and are not included in the
+Python wheel. To enable them, provision the complete bundle in a deployment
+directory and point the plugin to its absolute path before starting vLLM:
+
+```sh
+export VLLM_FL_THEAD_NATIVE_LIB_DIR=/absolute/path/to/thead-native-libs
+```
+
+The directory must contain `_C_stable_libtorch.abi3.so`, `_C.abi3.so`, and
+`_moe_C.abi3.so`. If the bundle is absent or incomplete, the plugin logs one
+warning per process and uses the configured FlagGems/reference fallbacks.
+Library load failures caused by an incompatible ABI or a missing transitive
+dependency are not suppressed. See
+[`PROVENANCE.md`](./vllm_fl/dispatch/backends/vendor/thead/lib/PROVENANCE.md)
+for the exact bundle that has been validated.
+
 4. (Optional) Install [FlagCX](https://github.com/flagos-ai/FlagCX/blob/main/docs/getting_started.md#build-and-installation)
 
     4.1 Clone the repository:

tests/unit_tests/dispatch/test_bf16_indexer_backend_ownership.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Contract tests for the composite BF16 Indexer decode backend."""
+
+from __future__ import annotations
+
+import importlib
+
+import pytest
+import torch
+
+
+@pytest.mark.parametrize(
+    ("candidate_indices", "expected"),
+    [
+        ([0], [0]),
+        ([0, 2, 1], [1, 2, 0]),
+        ([3, 0, 2, 1], [1, 2, 0, 3]),
+    ],
+)
+def test_flaggems_decode_owns_candidate_ordering(
+    monkeypatch, candidate_indices, expected
+):
+    module = importlib.import_module(
+        "vllm_fl.dispatch.backends.flaggems.impl.bf16_indexer"
+    )
+    logits = torch.tensor([[0.5, 4.0, 2.0, -1.0]], dtype=torch.float32)
+
+    monkeypatch.setattr(
+        module, "_bf16_paged_mqa_logits_flaggems", lambda *args, **kwargs: logits
+    )
+
+    def fake_topk(_logits, _seq_lens, indices, *, next_n):
+        assert next_n == 1
+        indices.copy_(torch.tensor([candidate_indices], dtype=torch.int32))
+
+    monkeypatch.setattr(module, "top_k_per_row_decode", fake_topk)
+
+    def reject_global_topk(*args, **kwargs):
+        raise AssertionError("candidate ordering must not re-enter global torch.topk")
+
+    monkeypatch.setattr(torch, "topk", reject_global_topk)
+    indices = torch.empty((1, len(candidate_indices)), dtype=torch.int32)
+    module.bf16_indexer_decode_flaggems(
+        None,
+        None,
+        None,
+        torch.tensor([4], dtype=torch.int32),
+        None,
+        None,
+        indices,
+        next_n=1,
+        max_context_len=4,
+    )
+    assert indices.tolist() == [expected]
+
+
+def test_reference_decode_honors_each_valid_length(monkeypatch):
+    module = importlib.import_module(
+        "vllm_fl.dispatch.backends.reference.impl.bf16_indexer"
+    )
+    logits = torch.tensor(
+        [[1.0, 5.0, 3.0, 99.0], [7.0, 80.0, 60.0, 40.0]],
+        dtype=torch.float32,
+    )
+    monkeypatch.setattr(
+        module, "_bf16_paged_mqa_logits_torch", lambda *args, **kwargs: logits
+    )
+    q = torch.empty((1, 2, 1, 1), dtype=torch.bfloat16)
+    indices = torch.empty((2, 3), dtype=torch.int32)
+    module.bf16_indexer_decode_torch(
+        q,
+        None,
+        None,
+        torch.tensor([[3, 1]], dtype=torch.int32),
+        None,
+        None,
+        indices,
+        next_n=2,
+        max_context_len=4,
+    )
+    assert indices.tolist() == [[1, 2, 0], [0, -1, -1]]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires accelerator")
+def test_flaggems_decode_graph_replay_uses_changed_logits(monkeypatch):
+    module = importlib.import_module(
+        "vllm_fl.dispatch.backends.flaggems.impl.bf16_indexer"
+    )
+    device = torch.device("cuda")
+    logits = torch.tensor(
+        [[1.0, 5.0, 3.0, 2.0, 4.0, -1.0], [6.0, 2.0, 5.0, 1.0, -1.0, -2.0]],
+        device=device,
+        dtype=torch.float32,
+    )
+    seq_lens = torch.tensor([5], device=device, dtype=torch.int32)
+    indices = torch.empty((2, 3), device=device, dtype=torch.int32)
+    monkeypatch.setattr(
+        module, "_bf16_paged_mqa_logits_flaggems", lambda *args, **kwargs: logits
+    )
+
+    def run():
+        module.bf16_indexer_decode_flaggems(
+            None,
+            None,
+            None,
+            seq_lens,
+            None,
+            None,
+            indices,
+            next_n=2,
+            max_context_len=6,
+        )
+
+    for _ in range(3):
+        run()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run()
+    graph.replay()
+    torch.cuda.synchronize()
+    first = indices.cpu().clone()
+
+    logits.copy_(
+        torch.tensor(
+            [[9.0, 1.0, 8.0, 2.0, 7.0, -1.0], [1.0, 9.0, 2.0, 8.0, -1.0, -2.0]],
+            device=device,
+        )
+    )
+    graph.replay()
+    torch.cuda.synchronize()
+    second = indices.cpu()
+    assert not torch.equal(first, second)
+    for row, valid_len in enumerate((4, 5)):
+        chosen = second[row].to(torch.long)
+        assert bool(((chosen >= 0) & (chosen < valid_len)).all())
+        scores = logits[row].cpu().index_select(0, chosen)
+        assert bool((scores[:-1] >= scores[1:]).all())

tests/unit_tests/dispatch/test_flaggems_mla_sparse.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+import sys
+from types import ModuleType
+
+from vllm_fl.dispatch.backends.flaggems.impl import mla_sparse
+
+
+def test_sparse_mla_prefers_flaggems_vllm(monkeypatch):
+    package = ModuleType("flaggems_vllm")
+
+    def specialized_impl(**kwargs):
+        return kwargs
+
+    package.flash_mla_sparse_fwd = specialized_impl
+    monkeypatch.setitem(sys.modules, "flaggems_vllm", package)
+
+    assert mla_sparse._resolve_flash_mla_sparse_fwd() is specialized_impl
+
+
+def test_sparse_mla_falls_back_when_specialized_symbol_is_missing(monkeypatch):
+    specialized_package = ModuleType("flaggems_vllm")
+    generic_package = ModuleType("flag_gems")
+
+    def generic_impl(**kwargs):
+        return kwargs
+
+    generic_package.flash_mla_sparse_fwd = generic_impl
+    monkeypatch.setitem(sys.modules, "flaggems_vllm", specialized_package)
+    monkeypatch.setitem(sys.modules, "flag_gems", generic_package)
+
+    assert mla_sparse._resolve_flash_mla_sparse_fwd() is generic_impl

tests/unit_tests/dispatch/test_thead_native_extensions.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Contract tests for optional T-Head native-extension initialization."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_fl.dispatch.backends.vendor.thead import bootstrap
+from vllm_fl.dispatch.backends.vendor.thead.impl import native_extensions
+
+
+def test_native_library_directory_requires_absolute_override(monkeypatch):
+    monkeypatch.setenv("VLLM_FL_THEAD_NATIVE_LIB_DIR", "relative/libs")
+
+    with pytest.raises(ValueError, match="must be an absolute path"):
+        native_extensions.load_all_native_extensions()
+
+
+def test_missing_bundle_reports_configured_directory(tmp_path, monkeypatch):
+    monkeypatch.setenv("VLLM_FL_THEAD_NATIVE_LIB_DIR", str(tmp_path))
+
+    with pytest.raises(native_extensions.NativeExtensionBundleMissingError) as error:
+        native_extensions.load_all_native_extensions()
+
+    assert {path.parent for path in error.value.missing_paths} == {tmp_path}
+    assert {path.name for path in error.value.missing_paths} == set(
+        native_extensions._FILES.values()
+    )
+
+
+def test_missing_bundle_warns_once_and_uses_fallback(tmp_path, monkeypatch):
+    monkeypatch.setenv("PPU_SDK", "/opt/ppu")
+    monkeypatch.setenv("VLLM_FL_THEAD_NATIVE_LIB_DIR", str(tmp_path))
+    monkeypatch.setattr(bootstrap, "_WARNED_MISSING_BUNDLES", set())
+    messages = []
+
+    def record_warning(message, *args):
+        messages.append(message % args)
+
+    monkeypatch.setattr(bootstrap.logger, "warning", record_warning)
+
+    assert bootstrap.initialize_native_extensions() is False
+    assert bootstrap.initialize_native_extensions() is False
+
+    assert len(messages) == 1
+    assert "T-Head native extensions are unavailable" in messages[0]
+    assert str(tmp_path) in messages[0]
+
+
+def test_non_missing_load_failure_is_not_suppressed(monkeypatch):
+    def fail_load():
+        raise OSError("incompatible ABI")
+
+    monkeypatch.setenv("PPU_SDK", "/opt/ppu")
+    monkeypatch.setattr(
+        native_extensions,
+        "load_all_native_extensions",
+        fail_load,
+    )
+
+    with pytest.raises(OSError, match="incompatible ABI"):
+        bootstrap.initialize_native_extensions()