[Feature] Add experimental Ascend support for Ouro

README.md

@@ -134,7 +134,7 @@ No PyPI publication is implied.
 
 | Checkpoint | Implementation | Documented validation scope | Guide |
 | --- | --- | --- | --- |
-| Ouro-1.4B | Native Ouro | Existing runtime and benchmark coverage under their recorded configurations | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.md) |
+| Ouro-1.4B | Native Ouro | Ascend inference support | [Runtime](docs/runtime.md), [Ascend recipe](docs/recipes/ouro_ascend.md) |
 | Ouro-2.6B | Shared Ouro | Shares the Ouro implementation; not separately qualified, and 1.4B quality/performance results are not inherited | [Scope notes](docs/recipes/ouro_thinking.md#6-evidence-and-limitations) |
 | Ouro-1.4B-Thinking | Reuses Ouro | Bounded BF16, fixed-four-loop synchronous basic integration | [Shared recipe](docs/recipes/ouro_thinking.md) |
 | Ouro-2.6B-Thinking | Reuses Ouro | Bounded BF16, fixed-four-loop synchronous basic integration | [Shared recipe](docs/recipes/ouro_thinking.md) |

docs/recipes/ouro_ascend.md

@@ -0,0 +1,71 @@
+# Ouro on Ascend
+
+Run Ouro on one Ascend NPU with BF16, synchronous greedy decoding and four
+recurrent loops. Use the explicit `ascend_eager` attention backend.
+
+## 1. Prepare the environment and model
+
+Use a compatible Torch, Torch-NPU and CANN environment with the project's
+Python dependencies installed. Keep the vendor Torch pair; CUDA/Triton
+packages are not needed for this backend.
+
+~~~bash
+source /path/to/cann/set_env.sh
+source /path/to/npu-environment/bin/activate
+python -m pip install --no-deps --no-build-isolation -e .
+export MODEL=/path/to/prepared/Ouro-1.4B
+~~~
+
+Use the checkpoint's own config and tokenizer. Model loading uses the shared
+native Ouro loader, including its supported sharded checkpoint layout.
+
+## 2. Start the server
+
+Select an allocated physical NPU using `ASCEND_RT_VISIBLE_DEVICES`.
+
+~~~bash
+ASCEND_RT_VISIBLE_DEVICES=<allocated-device> \
+python -m vllm_rlt.entrypoints.serve \
+  --model "$MODEL" --tokenizer "$MODEL" --served-model-name ouro \
+  --device npu --dtype bfloat16 --attention-backend ascend_eager \
+  --num-blocks 1024 --max-num-seqs 1 \
+  --max-num-batched-tokens 4096 --prefill-chunk-size 4096 \
+  --host 127.0.0.1 --port 8123
+~~~
+
+Explicit KV capacity is required. Adjust it for the model, available memory,
+prompt/output lengths and concurrency. To compare with full-prefill inference,
+both the chunk size and batched-token budget must fit the whole prompt.
+
+## 3. Send a request
+
+~~~bash
+curl --fail http://127.0.0.1:8123/health
+curl --fail http://127.0.0.1:8123/v1/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"ouro","prompt":"The capital of France is","max_tokens":8,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}'
+~~~
+
+Wait for HTTP 200 before sending requests. Add `"stream": true` for SSE.
+The `ascend_eager` backend follows HF-style BF16 attention rounding;
+`--attention-backend torch` selects the existing FP32-intermediate reference.
+
+## 4. Run checks
+
+With an assigned NPU visible to the test process:
+
+~~~bash
+pre-commit run --all-files
+VLLM_RLT_OURO_CHECKPOINT="$MODEL" python -m pytest \
+  tests/test_ascend.py tests/test_ascend_checkpoint.py \
+  tests/test_ascend_eager.py tests/test_ascend_kv.py \
+  tests/test_ascend_serving.py --run-npu -q
+~~~
+
+These checks cover causal attention, cross-page KV, request reuse and HTTP/SSE.
+They do not establish model quality or performance; use the
+[accuracy guide](../accuracy.md) and [serving benchmark guide](../serving.md)
+for those evaluations.
+
+Async scheduling, graphs, speculation, PD, preemption, prefix caching and
+adaptive exit are currently unsupported on this path.

pyproject.toml

@@ -43,7 +43,10 @@ license-files = ["LICENSE", "NOTICE"]
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 pythonpath = ["."]
-markers = ["gpu: requires a scheduler-reserved CUDA device"]
+markers = [
+    "gpu: requires a scheduler-reserved CUDA device",
+    "npu: requires an assigned Ascend device and torch-npu/CANN",
+]
 
 [tool.ruff]
 line-length = 100

tests/conftest.py

@@ -8,15 +8,21 @@ def pytest_addoption(parser):
     parser.addoption(
         "--run-gpu", action="store_true", help="Run CUDA tests inside a GPU reservation"
     )
+    parser.addoption("--run-npu", action="store_true", help="Run Ascend tests on an assigned NPU")
 
 
 def pytest_collection_modifyitems(config, items):
     enabled = config.getoption("--run-gpu")
     if enabled and not os.environ.get("CUDA_VISIBLE_DEVICES"):
         raise pytest.UsageError("--run-gpu requires scheduler-provided CUDA_VISIBLE_DEVICES")
+    npu_enabled = config.getoption("--run-npu")
+    if npu_enabled and not os.environ.get("ASCEND_RT_VISIBLE_DEVICES"):
+        raise pytest.UsageError("--run-npu requires assigned ASCEND_RT_VISIBLE_DEVICES")
     for item in items:
         if "gpu" in item.keywords and not enabled:
             item.add_marker(pytest.mark.skip(reason="use --run-gpu inside a GPU reservation"))
+        if "npu" in item.keywords and not npu_enabled:
+            item.add_marker(pytest.mark.skip(reason="use --run-npu on an assigned NPU"))
 
 
 @pytest.fixture(autouse=True, scope="session")

tests/test_ascend.py

@@ -0,0 +1,170 @@
+"""Opt-in synchronous BF16 Ouro checks using explicit Torch attention on Ascend."""
+
+import pytest
+import torch
+
+from tests.helpers import tiny_ouro_config
+from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig
+from vllm_rlt.core.kv_cache_manager import KVCacheManager
+from vllm_rlt.device import resolve_device
+from vllm_rlt.engine.stage_engine import StageEngine
+from vllm_rlt.kernels.paged_attention import torch_paged_attention
+from vllm_rlt.models import OuroForCausalLM
+
+pytestmark = pytest.mark.npu
+
+
+@pytest.fixture
+def npu_device():
+    resolve_device("npu")
+    torch.npu.set_device(0)
+    yield torch.device("npu:0")
+    torch.npu.synchronize()
+    torch.npu.empty_cache()
+
+
+def assert_bf16_close(actual, expected):
+    a, b = actual.detach().float().cpu(), expected.detach().float().cpu()
+    assert torch.isfinite(a).all() and torch.isfinite(b).all()
+    assert (a - b).norm() / b.norm().clamp_min(1e-8) <= 0.02
+    torch.testing.assert_close(a, b, atol=0.02, rtol=0.02)
+
+
+@pytest.mark.parametrize("kv_heads", [1, 2, 4])
+def test_torch_attention_on_npu_matches_cpu_prefix_oracle(npu_device, kv_heads):
+    torch.manual_seed(83)
+    keys = torch.randn(8, 2, 16, kv_heads, 64, dtype=torch.bfloat16)
+    values = torch.randn_like(keys)
+    q = torch.randn(64, 4, 64, dtype=torch.bfloat16)
+    lengths = torch.tensor([0, 1, 15, 16, 17, 31, 32, 3] * 8, dtype=torch.int32)
+    tables = torch.tensor([[5, 2], [1, 7]] * 32, dtype=torch.int32)
+    expected = torch_paged_attention(q, keys[:, 1], values[:, 1], tables, lengths)
+    actual = torch_paged_attention(
+        q.to(npu_device),
+        keys.to(npu_device)[:, 1],
+        values.to(npu_device)[:, 1],
+        tables.to(npu_device),
+        lengths.to(npu_device),
+    )
+    assert_bf16_close(actual, expected)
+    assert torch.count_nonzero(actual[0]).item() == 0
+
+
+def make_npu_model(device):
+    torch.manual_seed(41)
+    config = tiny_ouro_config(hidden_size=256, intermediate_size=512, head_dim=64)
+    return OuroForCausalLM(config).to(device=device, dtype=torch.bfloat16).eval()
+
+
+@torch.inference_mode()
+def recurrent_states(model, chunks, *, tokens=None, incremental=False):
+    c = model.config
+    cache = KVCacheManager(
+        c.num_hidden_layers,
+        c.num_key_value_heads,
+        c.head_dim,
+        32,
+        16,
+        4,
+        device=next(model.parameters()).device,
+        dtype=torch.bfloat16,
+        backend="torch",
+        incremental_allocation=incremental,
+    )
+    tokens = tokens if tokens is not None else ([3, 7, 2, 11, 5, 17, 8] * 3)[:21]
+    assert cache.allocate("a", len(tokens), initial_tokens=16 if incremental else None)
+    states, kv, start = {}, {}, 0
+    try:
+        for size in chunks:
+            assert cache.ensure_capacity("a", start + size)
+            positions = list(range(start, start + size))
+            hidden = model.prelude(torch.tensor(tokens[start : start + size], device=cache.device))
+            for depth in range(4):
+                hidden, _ = model.recurrent(hidden, ["a"] * size, [depth] * size, positions, cache)
+                states[(start, depth)] = (hidden.clone(), model.coda(hidden).clone())
+            start += size
+        assert start == len(tokens)
+        for depth in range(4):
+            for layer in range(c.num_hidden_layers):
+                assert all(cache.token_written("a", p, depth) for p in range(len(tokens)))
+                kv[(layer, depth)] = tuple(
+                    t.clone() for t in cache.read(layer, "a", depth, len(tokens))
+                )
+        assert len({tuple(cache.get_block_table("a", d)) for d in range(4)}) == 4
+        return states, kv
+    finally:
+        cache.free("a")
+        assert cache.num_used_blocks == 0
+
+
+def test_npu_chunked_prefill_and_incremental_kv(npu_device):
+    model = make_npu_model(npu_device)
+    expected, expected_kv = recurrent_states(model, [19, 1, 1])
+    incremental, incremental_kv = recurrent_states(model, [19, 1, 1], incremental=True)
+    for key in expected:
+        for a, b in zip(incremental[key], expected[key], strict=True):
+            torch.testing.assert_close(a, b, atol=0, rtol=0)
+    for key in expected_kv:
+        for a, b in zip(incremental_kv[key], expected_kv[key], strict=True):
+            torch.testing.assert_close(a, b, atol=0, rtol=0)
+    chunked, chunked_kv = recurrent_states(model, [5, 14, 1, 1], incremental=True)
+    for key in ((19, d) for d in range(4)):
+        for a, b in zip(chunked[key], expected[key], strict=True):
+            assert_bf16_close(a, b)
+    for key in expected_kv:
+        for a, b in zip(chunked_kv[key], expected_kv[key], strict=True):
+            assert_bf16_close(a, b)
+
+
+def test_npu_cancel_reuse_and_multiple_requests(npu_device):
+    engine = StageEngine(
+        make_npu_model(npu_device),
+        attention_backend="torch",
+        cache_config=CacheConfig(num_blocks=32, block_size=16, incremental_allocation=True),
+        scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=32),
+    )
+    try:
+        params = SamplingParams(max_tokens=3, temperature=0, ignore_eos=True)
+        engine.add_request("cancel", [3, 7, 2] * 6, params)
+        engine.add_request("keep", [5, 9, 11], params)
+        engine.step()
+        engine.abort_request("cancel")
+        engine.add_request("cancel", [13, 14, 15], params)
+        finished = {}
+        for _ in range(200):
+            for output in engine.step():
+                if output.finished:
+                    finished[output.request_id] = output
+            if not engine.has_unfinished_requests():
+                break
+        assert set(finished) == {"keep", "cancel"}
+        assert all(
+            o.exit_depths == [4] * 3 and o.finish_reason == "length" for o in finished.values()
+        )
+        assert engine.cache_manager.num_used_blocks == 0
+    finally:
+        engine.close()
+
+
+@pytest.mark.parametrize("ignore_eos", [False, True])
+def test_npu_eos_and_length_handling(npu_device, ignore_eos):
+    model = make_npu_model(npu_device)
+    model.lm_head.weight.zero_()  # Argmax 0 is the tiny model's actual EOS token.
+    engine = StageEngine(model, attention_backend="torch", cache_config=CacheConfig(num_blocks=16))
+    try:
+        engine.add_request(
+            "eos", [3, 7, 2], SamplingParams(max_tokens=3, temperature=0, ignore_eos=ignore_eos)
+        )
+        output = None
+        for _ in range(100):
+            for item in engine.step():
+                if item.finished:
+                    output = item
+            if not engine.has_unfinished_requests():
+                break
+        assert output is not None
+        assert output.token_ids == [0] * (3 if ignore_eos else 1)
+        assert output.finish_reason == ("length" if ignore_eos else "stop")
+        assert engine.cache_manager.num_used_blocks == 0
+    finally:
+        engine.close()

tests/test_ascend_checkpoint.py

@@ -0,0 +1,100 @@
+"""Opt-in local-checkpoint KV and lifecycle regressions on Ascend."""
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.test_ascend import assert_bf16_close, recurrent_states
+from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig
+from vllm_rlt.device import resolve_device
+from vllm_rlt.engine.stage_engine import StageEngine
+from vllm_rlt.models import OuroForCausalLM, load_tokenizer
+
+pytestmark = pytest.mark.npu
+
+
+@pytest.fixture(scope="module")
+def checkpoint():
+    source = os.environ.get("VLLM_RLT_OURO_CHECKPOINT")
+    if not source:
+        pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to a prepared local Ouro checkpoint")
+    if not Path(source).is_dir():
+        pytest.fail("VLLM_RLT_OURO_CHECKPOINT must be a prepared local directory")
+    resolve_device("npu")
+    torch.npu.set_device(0)
+    model = OuroForCausalLM.from_pretrained(source, device="npu", dtype=torch.bfloat16)
+    assert all(p.dtype == torch.bfloat16 for p in model.parameters())
+    yield model, load_tokenizer(source)
+    torch.npu.synchronize()
+
+
+@torch.inference_mode()
+def test_checkpoint_incremental_pages_preserve_hidden_logits_and_kv(checkpoint):
+    model, tokenizer = checkpoint
+    tokens = tokenizer.encode(
+        "The capital of France is Paris. The capital of Japan is Tokyo. "
+        "A language model predicts the next token in a sequence.",
+        add_special_tokens=False,
+    )[:21]
+    assert len(tokens) == 21
+    expected, expected_kv = recurrent_states(model, [19, 1, 1], tokens=tokens)
+    actual, actual_kv = recurrent_states(model, [19, 1, 1], tokens=tokens, incremental=True)
+    pairs = {}
+    for key in expected:
+        for name, a, b in zip(("hidden", "logits"), actual[key], expected[key], strict=True):
+            pairs[f"{key}.{name}"] = (a, b)
+    for key in expected_kv:
+        for name, a, b in zip(("key", "value"), actual_kv[key], expected_kv[key], strict=True):
+            pairs[f"{key}.{name}"] = (a, b)
+    for a, b in pairs.values():
+        assert_bf16_close(a, b)
+        torch.testing.assert_close(a, b, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("backend", ["torch", "ascend_eager"])
+@torch.inference_mode()
+def test_checkpoint_greedy_replay_and_request_cleanup(checkpoint, backend):
+    model, tokenizer = checkpoint
+    engine = StageEngine(
+        model,
+        attention_backend=backend,
+        cache_config=CacheConfig(num_blocks=64),
+        scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=32),
+    )
+    prompts = ["The capital of France is", "One plus one equals"]
+    loops = model.config.total_ut_steps
+    params = SamplingParams(
+        max_tokens=4, temperature=0, min_loops=loops, max_loops=loops, ignore_eos=True
+    )
+    expected = None
+    try:
+        # Repeat the same packed workload and reuse IDs after all pages are freed.
+        # Expected IDs come from this checkpoint, not a particular weight file.
+        for _ in range(2):
+            finished = {}
+            for index, prompt in enumerate(prompts):
+                engine.add_request(
+                    str(index), tokenizer.encode(prompt, add_special_tokens=False), params
+                )
+            for _ in range(200):
+                for output in engine.step():
+                    if output.finished:
+                        assert len(output.token_ids) == params.max_tokens
+                        assert all(
+                            0 <= token < model.config.vocab_size for token in output.token_ids
+                        )
+                        assert output.exit_depths == [loops] * params.max_tokens
+                        assert output.finish_reason == "length"
+                        finished[output.request_id] = list(output.token_ids)
+                if not engine.has_unfinished_requests():
+                    break
+            assert set(finished) == {str(index) for index in range(len(prompts))}
+            assert engine.cache_manager.num_used_blocks == 0
+            if expected is None:
+                expected = finished
+            else:
+                assert finished == expected
+    finally:
+        engine.close()