From a0f21e48a4cf71542e6cd81188472f13cd1cbdeb Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Mon, 2 Nov 2026 19:15:09 +0800 Subject: [PATCH 01/12] feat: bootstrap synchronous Ouro execution on Ascend NPU --- tests/test_device.py | 55 +++++++++++++++++++++++++++++++ vllm_rlt/core/kv_cache_manager.py | 3 +- vllm_rlt/core/memory.py | 2 ++ vllm_rlt/device.py | 13 ++++++++ vllm_rlt/engine/stage_engine.py | 24 ++++++++++++++ vllm_rlt/engine/worker.py | 7 +++- vllm_rlt/models/ouro.py | 2 ++ vllm_rlt/worker/model_runner.py | 10 ++++++ 8 files changed, 114 insertions(+), 2 deletions(-) create mode 100644 tests/test_device.py create mode 100644 vllm_rlt/device.py diff --git a/tests/test_device.py b/tests/test_device.py new file mode 100644 index 00000000..fa0cd828 --- /dev/null +++ b/tests/test_device.py @@ -0,0 +1,55 @@ +"""Device bootstrap and fail-fast NPU scope without a device dependency.""" + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from vllm_rlt import CacheConfig, ExecutionConfig, ExitConfig, SchedulerConfig +from vllm_rlt.core.memory import plan_cache +from vllm_rlt.device import resolve_device +from vllm_rlt.engine.stage_engine import StageEngine + + +def test_cpu_device_resolution_does_not_import_optional_backend(monkeypatch): + monkeypatch.setattr( + importlib, "import_module", lambda name: pytest.fail(f"unexpected import: {name}") + ) + assert resolve_device("cpu") == torch.device("cpu") + assert resolve_device("cuda:2") == torch.device("cuda:2") + + +def npu_model_stub(**overrides): + parameter = SimpleNamespace( + device=SimpleNamespace(type="npu"), dtype=torch.bfloat16, element_size=lambda: 2 + ) + return SimpleNamespace( + parameters=lambda: iter([parameter]), config=tiny_ouro_config(**overrides) + ) + + +@pytest.mark.parametrize( + "options", + [ + dict(execution_config=ExecutionConfig(async_scheduling=True)), + dict(execution_config=ExecutionConfig(cuda_graphs=True)), + dict(cache_config=CacheConfig(num_blocks=32, enable_prefix_caching=True)), + dict(scheduler_config=SchedulerConfig(enable_preemption=True)), + dict(exit_config=ExitConfig(mode="ouro_delayed")), + ], +) +def test_npu_unsupported_execution_rejected_before_allocation(options): + with pytest.raises(ValueError, match="NPU Ouro requires synchronous"): + StageEngine(npu_model_stub(), **options) + + +def test_npu_does_not_use_cpu_default_memory_budget(): + model = npu_model_stub() + with pytest.raises(ValueError, match="explicit num_blocks"): + plan_cache(model, CacheConfig(), SchedulerConfig(), ExecutionConfig(), "ascend") + blocks, info = plan_cache( + model, CacheConfig(num_blocks=17), SchedulerConfig(), ExecutionConfig(), "ascend" + ) + assert blocks == 17 and info["source"] == "blocks" diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..b0b28be3 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -23,6 +23,7 @@ create_backend, validate_backend_name, ) +from vllm_rlt.device import resolve_device from vllm_rlt.request import ExecutionBatch from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder @@ -133,7 +134,7 @@ def __init__( ): raise ValueError("recurrent_layers must be valid cache layer indices") self.storage_depths = max_loops if layout == "last_exited" else 1 - self.device = torch.device(device) + self.device = resolve_device(device) self.dtype = dtype self.attention = attention or create_backend( backend, self.device, dtype, head_dim, block_size diff --git a/vllm_rlt/core/memory.py b/vllm_rlt/core/memory.py index e9d2d4cf..21beea57 100644 --- a/vllm_rlt/core/memory.py +++ b/vllm_rlt/core/memory.py @@ -69,6 +69,8 @@ def plan_cache( if cache.kv_cache_memory_bytes is not None: blocks = budget_blocks(cache.kv_cache_memory_bytes, per_block) return blocks, {"source": "bytes", "bytes_per_block": per_block} + if device.type == "npu": + raise ValueError("NPU requires explicit num_blocks or kv_cache_memory_bytes") if device.type != "cuda": return 256, {"source": "cpu_default", "bytes_per_block": per_block} torch.cuda.synchronize(device) diff --git a/vllm_rlt/device.py b/vllm_rlt/device.py new file mode 100644 index 00000000..41244158 --- /dev/null +++ b/vllm_rlt/device.py @@ -0,0 +1,13 @@ +"""Register optional device types only when explicitly requested.""" + +import importlib + +import torch + + +def resolve_device(device) -> torch.device: + if isinstance(device, str) and device.split(":", 1)[0] == "npu": + # torch.device cannot parse private-use names until torch-npu registers + # them. Do this in each spawned worker, without importing it on CPU/CUDA. + importlib.import_module("torch_npu") + return torch.device(device) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..3001eee9 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -70,6 +70,26 @@ def __init__( self.execution_config = resolve_cuda_graph_rows( execution_config or ExecutionConfig(), scheduler_config, speculative_config ) + if parameter.device.type == "npu": + if config.to_dict()["model_type"] != "ouro" or config.total_ut_steps != 4: + raise ValueError("NPU currently supports Ouro with four loops") + if str(parameter.dtype) != "torch.bfloat16": + raise ValueError("NPU Ouro requires BF16 weights") + if ( + self.execution_config.async_scheduling + or self.execution_config.cuda_graphs + or self.execution_config.prefill_uva + or speculative_config is not None + or self.stage_config.role != "prefill_decode" + or scheduler_config.enable_preemption + or cache_config.enable_prefix_caching + or cache_config.layout != "last_exited" + or self.exit_config.mode != "ouro" + ): + raise ValueError( + "NPU Ouro requires synchronous local execution, ouro exit mode and " + "last_exited KV without speculation, graphs, preemption or prefix caching" + ) staged_prefill = requires_staged_prefill( scheduler_config, produces_tokens=self.stage_config.produces_tokens ) @@ -323,6 +343,10 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i if any(type(t) is not int or not 0 <= t < config.vocab_size for t in prompt_token_ids): raise ValueError("prompt token IDs must be integers within the model vocabulary") max_loops = params.max_loops or config.total_ut_steps + if next(self.model.parameters()).device.type == "npu" and ( + max_loops != 4 or params.exit_threshold != 1.0 or params.temperature != 0 + ): + raise ValueError("NPU Ouro requests require four loops, exit_threshold=1 and greedy") if max_loops > config.total_ut_steps or params.min_loops > max_loops: raise ValueError("requested loop bounds exceed the model's supported depth") if self.speculative_config is not None and ( diff --git a/vllm_rlt/engine/worker.py b/vllm_rlt/engine/worker.py index 9a2b2805..80e98985 100644 --- a/vllm_rlt/engine/worker.py +++ b/vllm_rlt/engine/worker.py @@ -13,6 +13,7 @@ import torch +from vllm_rlt.device import resolve_device from vllm_rlt.engine.connector import TransferConnector from vllm_rlt.engine.stage_config import StageConfig from vllm_rlt.engine.stage_engine import StageEngine @@ -33,12 +34,16 @@ def __init__( ): if type(max_control_messages) is not int or max_control_messages < 1: raise ValueError("max_control_messages must be a positive integer") - device = torch.device(f"cuda:{device}" if isinstance(device, int) else device) + device = resolve_device(f"cuda:{device}" if isinstance(device, int) else device) torch.set_num_threads(1) if device.type == "cuda": if device.index is None: device = torch.device("cuda", torch.cuda.current_device()) torch.cuda.set_device(device) + elif device.type == "npu": + if device.index is None: + device = torch.device("npu", torch.npu.current_device()) + torch.npu.set_device(device) torch.manual_seed(options["seed"]) self.channel = channel self.max_control_messages = max_control_messages diff --git a/vllm_rlt/models/ouro.py b/vllm_rlt/models/ouro.py index 31b37f9c..5041fdf3 100644 --- a/vllm_rlt/models/ouro.py +++ b/vllm_rlt/models/ouro.py @@ -19,6 +19,7 @@ import torch from torch import nn +from vllm_rlt.device import resolve_device from vllm_rlt.kernels import silu_and_mul from vllm_rlt.layers import ( RMSNorm, @@ -309,6 +310,7 @@ def from_pretrained( """ from safetensors import safe_open + device = resolve_device(device) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise ValueError("dtype must be float32, float16, or bfloat16") folder = Path(path_or_repo).expanduser() diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..93d59dac 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1152,6 +1152,9 @@ def _readback_slot(self, result): raise RuntimeError("readback buffers exhausted; retire completed submissions") def _synchronize_submission(self, stream, prepared): + if self.device.type == "npu": + torch.npu.synchronize(self.device) + return if self.device.type != "cuda": return (stream or torch.cuda.current_stream(self.device)).synchronize() @@ -1254,6 +1257,10 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu for span in batch.spans: self.events[span.request.request_id] = compute_done self.submissions.complete(compute_done) + if self.device.type == "npu": + # No CUDA completion events protect NPU metadata snapshots. + # Drain before pools/host inputs can be recycled or KV freed. + torch.npu.synchronize(self.device) if split_prefill: self.split_prefill_submissions += 1 return Submission( @@ -1390,6 +1397,9 @@ def release(self, request_id): self.rng.drop(request_id) def synchronize(self): + if self.device.type == "npu": + torch.npu.synchronize(self.device) + return for stream in ( self.core_stream, self.boundary_stream, From 4734e1963bdf7fd4e0fd66b2612511d6d4173550 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Mon, 2 Nov 2026 23:51:35 +0800 Subject: [PATCH 02/12] feat: expose synchronous BF16 Ouro on Ascend through Torch attention --- README.md | 2 +- docs/recipes/ouro_ascend.md | 113 +++++++++++++++++++++ pyproject.toml | 5 +- tests/conftest.py | 6 ++ tests/test_ascend.py | 171 ++++++++++++++++++++++++++++++++ tests/test_ascend_checkpoint.py | 113 +++++++++++++++++++++ tests/test_ascend_serving.py | 135 +++++++++++++++++++++++++ tests/test_device.py | 16 ++- vllm_rlt/engine/stage_engine.py | 6 +- vllm_rlt/entrypoints/cli.py | 2 + vllm_rlt/entrypoints/serve.py | 4 +- vllm_rlt/profiling.py | 2 + 12 files changed, 569 insertions(+), 6 deletions(-) create mode 100644 docs/recipes/ouro_ascend.md create mode 100644 tests/test_ascend.py create mode 100644 tests/test_ascend_checkpoint.py create mode 100644 tests/test_ascend_serving.py diff --git a/README.md b/README.md index 00ed4e17..4240b3fd 100644 --- a/README.md +++ b/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 | Existing runtime and benchmark coverage under their recorded configurations; initial single-device Ascend 910B3 BF16/fixed-four-loop synchronous path with explicit Torch attention | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.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) | diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md new file mode 100644 index 00000000..1c27dd9d --- /dev/null +++ b/docs/recipes/ouro_ascend.md @@ -0,0 +1,113 @@ +# Ouro-1.4B on Ascend 910B3: initial synchronous support + +This path reuses the native Ouro model and shared engine on one Ascend NPU. +Select **--attention-backend torch explicitly**. Weights, ordinary activations +and KV storage remain BF16; the existing Torch attention computes QK, softmax +and PV intermediates in FP32 through Torch-NPU. This is not a fused CANN +attention integration and makes no acceleration claim. + +## Supported scope + +- Ouro, fixed four-loop greedy inference, synchronous local execution, + LAST_EXITED depth-aware paged KV, prefill and incremental decode. +- Explicit num_blocks or kv_cache_memory_bytes; automatic NPU memory sizing + is not implemented. +- Cancellation, request ID reuse, EOS/length handling, HTTP completions and SSE + reuse the existing engine and serving lifecycle. +- Async scheduling, CUDA graphs, speculation, preemption, prefix caching, PD, + adaptive exit and stochastic sampling are rejected for this initial path. +- Qualification targets Ouro-1.4B on 910B3. Other checkpoints and Ascend devices + require separate validation. + +The shared scheduler, model and KV owner are retained. NPU submissions drain +before metadata buffers or request pages can be reused. CPU/CUDA execution +does not import the optional Torch-NPU package. Built-in profiling supports +CPU/CUDA; use torch_npu.profiler for NPU device traces. + +## Environment and installation + +The development environment is Linux aarch64, Python 3.11.13, Torch 2.9.0+cpu, +torch-npu 2.9.0, CANN 9.0.0, driver 26.1.0.b050 and Transformers 4.57.6. +Keep Torch and Torch-NPU matched; activate vendor CANN before the project Python. + +~~~bash +source /path/to/cann/set_env.sh +source /path/to/isolated-npu-environment/bin/activate +export ASCEND_RT_VISIBLE_DEVICES=4 # replace with your assigned physical NPU +python -c 'import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.get_device_name(0))' +python -m pip install 'numpy>=1.24' 'safetensors>=0.4.3' 'huggingface-hub>=0.25' \ + 'transformers>=4.54' 'aiohttp>=3.11,<4' setuptools wheel +python -m pip install --no-deps --no-build-isolation -e . +~~~ + +The no-deps install preserves the vendor Torch pair and avoids the project's +Linux Triton dependency. Neither Triton nor a native CUDA extension is required +for this explicit Torch backend. No InferRT or MindIE-SD integration is involved. + +Use the complete ByteDance/Ouro-1.4B checkpoint at revision +574fa66cb8bf5abdc979642d01cf2b79b16bfab1. The model.safetensors SHA256 is +58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af. + +## Start the service + +~~~bash +python -m vllm_rlt.entrypoints.serve \ + --model /path/to/Ouro-1.4B --served-model-name ouro \ + --device npu --dtype bfloat16 --attention-backend torch \ + --num-blocks 64 --block-size 16 --max-num-seqs 2 \ + --max-num-batched-tokens 32 --port 8123 +~~~ + +~~~bash +curl http://127.0.0.1:8123/health +curl http://127.0.0.1:8123/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"ouro","prompt":"The capital of France is","max_tokens":4,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}' +~~~ + +Add stream=true for SSE. Python callers use device="npu" and +attention_backend="torch"; the spawned worker registers Torch-NPU itself. +The assigned physical device becomes logical npu:0. Increase KV capacity for +longer prompts, output limits or concurrency. This small service command is a +smoke configuration, not a performance benchmark. + +## Reproduce validation + +~~~bash +python -m pip install pytest ruff pre-commit +pre-commit run --all-files +export VLLM_RLT_OURO_CHECKPOINT=/path/to/Ouro-1.4B +export VLLM_RLT_ASCEND_RESULTS=/path/to/validation-results +python -m pytest tests/test_ascend.py tests/test_ascend_checkpoint.py \ + tests/test_ascend_serving.py --run-npu -q +~~~ + +The NPU flag requires assigned ASCEND_RT_VISIBLE_DEVICES; CPU collection skips +device tests. Checks cover a CPU attention oracle, GQA and causal page +boundaries, incremental/full KV allocation, four-depth state ownership, +chunking, cancellation/reuse, EOS/length behavior, pinned checkpoint generation +and a real spawned HTTP/SSE worker. + +Metadata-only allocation changes require exact hidden/logit/KV equality. +The BF16 diagnostic bounds remain relative L2 at most 2% and elementwise +atol=0.02, rtol=0.02. These are this validation's bounds, not a universal quality +guarantee. Task quality needs separately pinned paired evaluation against the +official code; short generation does not establish broader quality. +Performance characterization should retain input lengths, concurrency, +TTFT/TPOT/throughput, memory and repetitions; the existing +[serving benchmark guide](../serving.md#reuse-the-upstream-benchmark) describes +the compatible upstream client. + +## Fused attention remains separate + +An experimental BF16 paged npu_fused_infer_attention_score adapter is outside +this change. It failed frozen whole-model numerical checks and lost a correct +answer in a ten-question smoke. Promoting attention intermediates in +npu_fusion_attention improved local mathematical agreement but still failed +the original whole-model joint gate. Those failures have not been relaxed or +hidden; this recipe does not qualify that fused backend. + +This contribution covers a bounded part of +[#69 Ascend support](https://github.com/InferMatrixLab/vllm-rlt/issues/69). +It does not complete CANN fused adaptation or other Ouro runtime features on +Ascend. Measured qualification results are recorded separately from this scope. diff --git a/pyproject.toml b/pyproject.toml index 04ccfa6f..42d86fc9 100644 --- a/pyproject.toml +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index fa5d6020..83d268e4 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/test_ascend.py b/tests/test_ascend.py new file mode 100644 index 00000000..4632c7f1 --- /dev/null +++ b/tests/test_ascend.py @@ -0,0 +1,171 @@ +"""Opt-in synchronous BF16 Ouro checks using explicit Torch attention on 910B3.""" + +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) + assert torch.npu.get_device_name(0) == "Ascend910B3" + 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() diff --git a/tests/test_ascend_checkpoint.py b/tests/test_ascend_checkpoint.py new file mode 100644 index 00000000..6b08c1d3 --- /dev/null +++ b/tests/test_ascend_checkpoint.py @@ -0,0 +1,113 @@ +"""Pinned real-weight NPU checks for explicit Torch attention, not fused CANN.""" + +import hashlib +import json +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 the pinned local Ouro-1.4B checkpoint") + resolve_device("npu") + torch.npu.set_device(0) + assert torch.npu.get_device_name(0) == "Ascend910B3" + with (Path(source) / "model.safetensors").open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + assert digest == "58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af" + 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() + + +def save_record(name, record): + folder = os.environ.get("VLLM_RLT_ASCEND_RESULTS") + if folder: + path = Path(folder) + path.mkdir(parents=True, exist_ok=True) + (path / name).write_text(json.dumps(record, indent=2)) + + +@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) + save_record( + "checkpoint-numerics.json", + dict( + tokens=tokens, + states=len(pairs), + exact=all(torch.equal(a, b) for a, b in pairs.values()), + finite=all( + bool(torch.isfinite(a).all() and torch.isfinite(b).all()) for a, b in pairs.values() + ), + ), + ) + for a, b in pairs.values(): + assert_bf16_close(a, b) + torch.testing.assert_close(a, b, atol=0, rtol=0) + + +@torch.inference_mode() +def test_checkpoint_greedy_generation_and_cleanup(checkpoint): + model, tokenizer = checkpoint + engine = StageEngine( + model, + attention_backend="torch", + 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"] + finished = {} + try: + for index, prompt in enumerate(prompts): + engine.add_request( + str(index), + tokenizer.encode(prompt, add_special_tokens=False), + SamplingParams(max_tokens=4, temperature=0, ignore_eos=True), + ) + for _ in range(200): + for output in engine.step(): + if output.finished: + assert output.exit_depths == [4] * 4 + assert output.finish_reason == "length" + finished[output.request_id] = dict( + tokens=output.token_ids, text=tokenizer.decode(output.token_ids) + ) + if not engine.has_unfinished_requests(): + break + assert set(finished) == {"0", "1"} + assert finished["0"]["tokens"] == [7042, 30, 198, 198] + assert finished["1"]["tokens"] == [827, 30, 198, 198] + assert engine.cache_manager.num_used_blocks == 0 + save_record("checkpoint-generation.json", finished) + finally: + engine.close() diff --git a/tests/test_ascend_serving.py b/tests/test_ascend_serving.py new file mode 100644 index 00000000..90f596c0 --- /dev/null +++ b/tests/test_ascend_serving.py @@ -0,0 +1,135 @@ +"""Real socket + spawned NPU worker smoke, opt-in with pinned local weights.""" + +import asyncio +import json +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.npu + + +def test_checkpoint_http_completions_stream_and_disconnect(tmp_path): + import aiohttp + + source = os.environ.get("VLLM_RLT_OURO_CHECKPOINT") + if not source: + pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to the pinned local Ouro-1.4B checkpoint") + folder = Path(os.environ.get("VLLM_RLT_ASCEND_RESULTS", tmp_path)) + folder.mkdir(parents=True, exist_ok=True) + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + command = [ + sys.executable, + "-m", + "vllm_rlt.entrypoints.serve", + "--model", + source, + "--served-model-name", + "ouro-ascend", + "--device", + "npu", + "--dtype", + "bfloat16", + "--attention-backend", + "torch", + "--num-blocks", + "64", + "--max-num-seqs", + "2", + "--max-num-batched-tokens", + "32", + "--port", + str(port), + ] + with (folder / "service.log").open("w") as log: + process = subprocess.Popen( + command, stdout=log, stderr=subprocess.STDOUT, start_new_session=True + ) + + async def exercise(): + url = f"http://127.0.0.1:{port}" + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as client: + deadline = time.monotonic() + 120 + while True: + assert process.poll() is None, (folder / "service.log").read_text() + try: + async with client.get(url + "/health") as response: + if response.status == 200: + break + except aiohttp.ClientConnectionError: + pass + assert time.monotonic() < deadline, "NPU service startup timed out" + await asyncio.sleep(0.25) + + def body(prompt, **extra): + return dict( + model="ouro-ascend", + prompt=prompt, + max_tokens=4, + temperature=0, + min_loops=4, + max_loops=4, + exit_threshold=1, + ignore_eos=True, + **extra, + ) + + async def complete(prompt): + async with client.post(url + "/v1/completions", json=body(prompt)) as response: + result = await response.json() + assert response.status == 200, result + assert result["usage"]["completion_tokens"] == 4 + return result + + prompts = ["The capital of France is", "One plus one equals"] + normal = await asyncio.gather(*(complete(prompt) for prompt in prompts)) + events = [] + async with client.post( + url + "/v1/completions", + json=body(prompts[0], stream=True, stream_options={"include_usage": True}), + ) as response: + assert response.status == 200 + async for line in response.content: + if line.startswith(b"data: "): + data = line[6:].strip().decode() + events.append(data if data == "[DONE]" else json.loads(data)) + assert events[-1] == "[DONE]" + text = "".join( + choice["text"] for event in events[:-1] for choice in event.get("choices", []) + ) + assert text == normal[0]["choices"][0]["text"] + assert any( + isinstance(event, dict) + and (event.get("usage") or {}).get("completion_tokens") == 4 + for event in events + ) + # Close a real streaming connection early, then verify the worker + # remains usable. Direct engine tests check page reclamation. + async with client.post( + url + "/v1/completions", json=body(prompts[1], stream=True) + ) as response: + assert response.status == 200 + await response.content.readline() + after_disconnect = await complete(prompts[1]) + assert after_disconnect["choices"][0]["text"] == normal[1]["choices"][0]["text"] + return dict(normal=normal, stream=events, after_disconnect=after_disconnect) + + try: + result = asyncio.run(exercise()) + (folder / "service-results.json").write_text(json.dumps(result, indent=2)) + finally: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=10) + assert process.returncode == 0, (folder / "service.log").read_text() diff --git a/tests/test_device.py b/tests/test_device.py index fa0cd828..02672ce0 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -11,6 +11,7 @@ from vllm_rlt.core.memory import plan_cache from vllm_rlt.device import resolve_device from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.profiling import ProfileConfig def test_cpu_device_resolution_does_not_import_optional_backend(monkeypatch): @@ -48,8 +49,19 @@ def test_npu_unsupported_execution_rejected_before_allocation(options): def test_npu_does_not_use_cpu_default_memory_budget(): model = npu_model_stub() with pytest.raises(ValueError, match="explicit num_blocks"): - plan_cache(model, CacheConfig(), SchedulerConfig(), ExecutionConfig(), "ascend") + plan_cache(model, CacheConfig(), SchedulerConfig(), ExecutionConfig(), "torch") blocks, info = plan_cache( - model, CacheConfig(num_blocks=17), SchedulerConfig(), ExecutionConfig(), "ascend" + model, CacheConfig(num_blocks=17), SchedulerConfig(), ExecutionConfig(), "torch" ) assert blocks == 17 and info["source"] == "blocks" + + +@pytest.mark.parametrize("backend", ["auto", "triton", "flash_attn_4"]) +def test_npu_requires_explicit_torch_before_allocation(backend): + with pytest.raises(ValueError, match="explicit attention_backend='torch'"): + StageEngine(npu_model_stub(), attention_backend=backend) + + +def test_builtin_profiler_does_not_mislabel_npu_as_cpu_only(tmp_path): + with pytest.raises(ValueError, match="torch_npu.profiler"): + ProfileConfig(str(tmp_path)).resolve_activities("npu:0") diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 3001eee9..dd13576a 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -2,6 +2,8 @@ from dataclasses import dataclass, replace from typing import TYPE_CHECKING, cast +import torch + from vllm_rlt.attention import BackendCapabilities, create_backend from vllm_rlt.config import ( CacheConfig, @@ -71,9 +73,11 @@ def __init__( execution_config or ExecutionConfig(), scheduler_config, speculative_config ) if parameter.device.type == "npu": + if attention_backend != "torch": + raise ValueError("NPU execution requires explicit attention_backend='torch'") if config.to_dict()["model_type"] != "ouro" or config.total_ut_steps != 4: raise ValueError("NPU currently supports Ouro with four loops") - if str(parameter.dtype) != "torch.bfloat16": + if parameter.dtype != torch.bfloat16: raise ValueError("NPU Ouro requires BF16 weights") if ( self.execution_config.async_scheduling diff --git a/vllm_rlt/entrypoints/cli.py b/vllm_rlt/entrypoints/cli.py index a9b8025c..5099ef26 100644 --- a/vllm_rlt/entrypoints/cli.py +++ b/vllm_rlt/entrypoints/cli.py @@ -100,6 +100,8 @@ def main(): parser.error("--toy uses built-in token ID prompts; omit --prompt") if not args.toy and not args.prompt: parser.error("provide --prompt or use --toy for a CPU smoke test") + if args.device.split(":", 1)[0] == "npu" and args.attention_backend != "torch": + parser.error("NPU execution requires explicit --attention-backend torch") profiling = profile_config_from_args(args) if profiling is not None: profiling.resolve_activities(args.device) diff --git a/vllm_rlt/entrypoints/serve.py b/vllm_rlt/entrypoints/serve.py index cd452d92..713bab30 100644 --- a/vllm_rlt/entrypoints/serve.py +++ b/vllm_rlt/entrypoints/serve.py @@ -82,7 +82,7 @@ def main(): parser.add_argument("--tokenizer") parser.add_argument("--tokenizer-revision") parser.add_argument("--served-model-name") - parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--device", choices=["cpu", "cuda", "npu"], default="cuda") parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") parser.add_argument( "--attention-backend", @@ -115,6 +115,8 @@ def main(): limits = ServingLimits(**{name: getattr(args, name) for name in defaults}) if args.device == "cpu" and args.attention_backend not in ("auto", "torch"): parser.error("CPU execution requires --attention-backend torch or auto") + if args.device == "npu" and args.attention_backend != "torch": + parser.error("NPU execution requires explicit --attention-backend torch") torch.set_num_threads(args.cpu_threads) logging.basicConfig(level=logging.INFO) app = create_app( diff --git a/vllm_rlt/profiling.py b/vllm_rlt/profiling.py index 41245be8..757bbafd 100644 --- a/vllm_rlt/profiling.py +++ b/vllm_rlt/profiling.py @@ -51,6 +51,8 @@ def __post_init__(self): object.__setattr__(self, "activities", values) def resolve_activities(self, device): + if str(device).startswith("npu"): + raise ValueError("NPU device profiling requires torch_npu.profiler") names = self.activities or (("cpu", "cuda") if str(device).startswith("cuda") else ("cpu",)) values = [getattr(torch.profiler.ProfilerActivity, name.upper()) for name in names] if not set(values) <= torch.profiler.supported_activities(): From 04d8ce818dd5d0b8ec110cfbd4fad2e003a2af30 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 01:27:27 +0800 Subject: [PATCH 03/12] docs: record Ascend validation results and quality blockers --- README.md | 2 +- docs/recipes/ouro_ascend.md | 77 +++++++++++++++++++++++++++++++++++-- 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 4240b3fd..183cc29e 100644 --- a/README.md +++ b/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; initial single-device Ascend 910B3 BF16/fixed-four-loop synchronous path with explicit Torch attention | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.md), [Ascend recipe](docs/recipes/ouro_ascend.md) | +| Ouro-1.4B | Native Ouro | Existing runtime and benchmark coverage under their recorded configurations; experimental single-device Ascend 910B3 BF16/fixed-four-loop synchronous Torch path, with paired quality regressions pending correction | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.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) | diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md index 1c27dd9d..43d58995 100644 --- a/docs/recipes/ouro_ascend.md +++ b/docs/recipes/ouro_ascend.md @@ -6,6 +6,11 @@ and KV storage remain BF16; the existing Torch attention computes QK, softmax and PV intermediates in FP32 through Torch-NPU. This is not a fused CANN attention integration and makes no acceleration claim. +**Quality status:** the selected same-device paired screen is HF 63/87 versus +native 61/87 (losses: 203 and 881; gains: none). The frozen no-net-loss gate +failed. This implementation remains experimental and is not ready for +submission as a quality-qualified backend. + ## Supported scope - Ouro, fixed four-loop greedy inference, synchronous local execution, @@ -35,8 +40,8 @@ source /path/to/cann/set_env.sh source /path/to/isolated-npu-environment/bin/activate export ASCEND_RT_VISIBLE_DEVICES=4 # replace with your assigned physical NPU python -c 'import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.get_device_name(0))' -python -m pip install 'numpy>=1.24' 'safetensors>=0.4.3' 'huggingface-hub>=0.25' \ - 'transformers>=4.54' 'aiohttp>=3.11,<4' setuptools wheel +python -m pip install 'numpy==1.26.4' 'safetensors==0.8.0' 'huggingface-hub==0.36.2' \ + 'transformers==4.57.6' 'tokenizers==0.22.2' 'aiohttp==3.14.3' setuptools wheel python -m pip install --no-deps --no-build-isolation -e . ~~~ @@ -110,4 +115,70 @@ hidden; this recipe does not qualify that fused backend. This contribution covers a bounded part of [#69 Ascend support](https://github.com/InferMatrixLab/vllm-rlt/issues/69). It does not complete CANN fused adaptation or other Ouro runtime features on -Ascend. Measured qualification results are recorded separately from this scope. +Ascend. Measured validation outcomes and remaining blockers are recorded below. + +## Measured validation (2026-11-03) + +Runtime source: 4734e1963bdf7fd4e0fd66b2612511d6d4173550, against base +989921ac1d3897c5c43e27ec8cddbc91fa953194. Later recipe updates do not change +the measured runtime sources. + +- Ruff lint/format and the CPU regression hook passed. +- Opt-in real 910B3 tests: 10 passed in 44.41s, including actual HTTP/SSE. +- Incremental/preallocated real-weight KV differential: 216 finite states, all + exactly equal; attention/chunking diagnostics retain their original joint bounds. +- Same-device paired selected GSM8K-87: official 63/87, native 61/87; + losses [203, 881], gains []; 63/87 exact token sequences. + Parse failures: {'hf_eager': 19, 'rlt_target': 19}; length limits: {'hf_eager': 0, 'rlt_target': 0}; + execution limits: {'hf_eager': 0, 'rlt_target': 0}. +- Each pair runs on the same physical 910B3; all 87 cases pool three cards. + The original ten are regression smoke; 77 cases were held out. This is + selected-case evidence, not full GSM8K or adaptive/concurrent quality. + +The quality run uses official unmodified HF eager, Transformers 4.57.6, +Tokenizers 0.22.2, lm-eval 0.4.9.2 strict scoring, three-shot raw CoT, +greedy four-loop BF16, max 1024 output tokens and source-ranked seed0 IDs. +Dataset revision is 740312add88f781978c0658806c59bc2815b9866. + +The table below is a separate **engine-only short-burst** characterization: +8 output tokens/request, ignore EOS, one shape-specific warmup and three +unprofiled repetitions. Values are medians of per-burst request means; output +throughput includes prefill and drain. TTFT starts before engine admission. +TPOT is first-to-last observed output time divided by seven. + +| Input | Concurrency | Output tok/s (range) | TTFT ms | TPOT ms | Peak allocated / reserved GiB | +|---|---:|---:|---:|---:|---:| +| 16 | 1 | 2.726 (2.700–2.901) | 1238.01 | 241.94 | 4.970 / 5.309 | +| 16 | 2 | 3.793 (3.657–3.810) | 2122.58 | 298.19 | 4.970 / 5.311 | +| 16 | 8 | 5.031 (5.016–5.086) | 7980.90 | 675.01 | 4.973 / 5.314 | +| 64 | 1 | 1.367 (1.342–1.372) | 4176.89 | 242.77 | 4.972 / 5.314 | +| 64 | 2 | 1.590 (1.569–1.599) | 7964.07 | 300.45 | 4.974 / 5.314 | +| 64 | 8 | 1.789 (1.786–1.793) | 23417.86 | 1756.95 | 4.984 / 5.314 | +| 256 | 1 | 0.370 (0.364–0.373) | 19472.72 | 307.45 | 4.978 / 5.318 | +| 256 | 2 | 0.372 (0.371–0.377) | 39959.74 | 431.61 | 4.988 / 5.318 | +| 256 | 8 | 0.379 (0.378–0.380) | 125114.88 | 6228.64 | 4.988 / 5.318 | + +Configuration: 768 blocks, block size 16, max 8 sequences / 256 batched tokens, +incremental allocation, no async/graphs/UVA/prefix caching/speculation. +Measured requests: 99, output tokens: 792, failures: 0. Live logical KV blocks +after each completed burst: 0. +Torch allocator peaks include model/cache/workspace and exclude some driver +or runtime memory. Other cards ran independent quality shards; CPU/NUMA +affinity was not pinned. No steady-state throughput, HTTP/p99, accelerator +hardware-only comparison or acceleration claim is made. + +Separate ops-only/full profiles capture a complete warmed second-token decode +iteration, with actual NPU device kernels. Full mode enables shapes, stacks +and memory; exported shapes, readable Python stacks and memory CSV rows +were checked separately from performance timing. +Full prefill plus all eight output tokens is outside these diagnostic captures. + +The conservative Torch path has substantial prefill/dispatch overhead. Fused +attention optimization remains a later step and must retain the quality gates. + +The no-net-loss quality criterion failed on this selected set. Case 203 +changes a quantity from 7 to 5 and produces 27 instead of 31. Case 881 writes +'The answer is 10 + 6 = 16'; the unchanged strict scorer extracts 10. It remains +a failure under the frozen protocol. Neither the scorer nor the gate was +relaxed. A diagnostic HF-style BF16 attention alternative also failed the +original numerical bounds and was not incorporated. From bc6725df8dd358922f730a520e04065078db6323 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 02:51:19 +0800 Subject: [PATCH 04/12] perf: use native Ascend scatter for paged KV writes --- tests/test_ascend_kv.py | 51 +++++++++++++++++++++++++++++++ vllm_rlt/core/kv_cache_manager.py | 18 +++++++++-- vllm_rlt/kernels/ascend_kv.py | 41 +++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 tests/test_ascend_kv.py create mode 100644 vllm_rlt/kernels/ascend_kv.py diff --git a/tests/test_ascend_kv.py b/tests/test_ascend_kv.py new file mode 100644 index 00000000..8445e0f0 --- /dev/null +++ b/tests/test_ascend_kv.py @@ -0,0 +1,51 @@ +"""Assigned-device regressions for native KV copies and descriptor lifetime.""" + +import pytest +import torch + +from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.device import resolve_device + +pytestmark = pytest.mark.npu + + +@pytest.mark.parametrize("rows", [1, 3]) +@torch.inference_mode() +def test_native_scatter_preserves_inputs_other_layers_and_reused_ids(rows): + device = resolve_device("npu:0") + torch.npu.set_device(device) + cache = KVCacheManager(2, 2, 64, 32, 16, 4, device=device, dtype=torch.bfloat16) + cache.key_cache.zero_() + cache.value_cache.zero_() + assert cache.allocate("a", 17) + assert cache.allocate("b", 3) + ids, depths, positions = ["a", "b", "a"][:rows], [0, 2, 3][:rows], [15, 0, 16][:rows] + batch = cache._prepare_batch(ids, depths, positions) + torch.manual_seed(99) + k = torch.randn(rows, 2, 64, device=device, dtype=torch.bfloat16) + v = torch.randn_like(k) + saved_k, saved_v = k.clone(), v.clone() + expected_k, expected_v = cache.key_cache.clone(), cache.value_cache.clone() + for layer in range(2): + expected_k[batch.write_blocks, layer, batch.write_offsets] = k + expected_v[batch.write_blocks, layer, batch.write_offsets] = v + cache._write_prepared(layer, batch, k, v) + torch.testing.assert_close(cache.key_cache, expected_k, atol=0, rtol=0) + torch.testing.assert_close(cache.value_cache, expected_v, atol=0, rtol=0) + torch.testing.assert_close(k, saved_k, atol=0, rtol=0) + torch.testing.assert_close(v, saved_v, atol=0, rtol=0) + cache.free("a") + assert cache.allocate("a", 17) + with pytest.raises((ValueError, RuntimeError)): + cache._write_prepared(0, batch, k, v) + new_batch = cache._prepare_batch(["a"], [1], [0]) + expected_k, expected_v = cache.key_cache.clone(), cache.value_cache.clone() + expected_k[new_batch.write_blocks, 1, new_batch.write_offsets] = k[:1] + expected_v[new_batch.write_blocks, 1, new_batch.write_offsets] = v[:1] + cache._write_prepared(1, new_batch, k[:1], v[:1]) + torch.testing.assert_close(cache.key_cache, expected_k, atol=0, rtol=0) + torch.testing.assert_close(cache.value_cache, expected_v, atol=0, rtol=0) + cache.free("a") + cache.free("b") + assert cache.num_used_blocks == 0 + torch.npu.synchronize() diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index b0b28be3..37c1352a 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -684,8 +684,22 @@ def _write_prepared( self._validate_tensor(v, len(batch.position_ids), "v") if not batch.rows: return - self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k[: len(batch.rows)] - self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v[: len(batch.rows)] + if self.device.type == "npu": + if not hasattr(self, "_npu_writer"): + from vllm_rlt.kernels.ascend_kv import AscendKVWriter + + self._npu_writer = AscendKVWriter() + self._npu_writer.write( + self.key_cache, + self.value_cache, + layer, + batch, + k[: len(batch.rows)], + v[: len(batch.rows)], + ) + else: + self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k[: len(batch.rows)] + self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v[: len(batch.rows)] for allocation, depth, position in batch.rows: allocation.written[self._plane(depth)][layer].add(position) diff --git a/vllm_rlt/kernels/ascend_kv.py b/vllm_rlt/kernels/ascend_kv.py new file mode 100644 index 00000000..fc41189c --- /dev/null +++ b/vllm_rlt/kernels/ascend_kv.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact KV copies through native Scatter ND on Ascend.""" + +from __future__ import annotations + +import weakref + +import torch + + +class AscendKVWriter: + """Reuse immutable traversal indices across layers without retaining the descriptor.""" + + def __init__(self): + self._metadata = None + self._indices = None + + def write(self, keys, values, layer, batch, k, v): + import torch_npu + + if self._metadata is None or self._metadata() is not batch: + owner = batch.owner + addresses = [ + ( + allocation.block_tables[owner._plane(depth)][position // owner.block_size], + position % owner.block_size, + ) + for allocation, depth, position in batch.rows + ] + self._indices = torch.tensor( + [ + [[block, index, offset] for block, offset in addresses] + for index in range(keys.shape[1]) + ], + device=keys.device, + dtype=torch.int32, + ) + self._metadata = weakref.ref(batch) + indices = self._indices[layer] + torch_npu.npu_scatter_nd_update_(keys, indices, k) + torch_npu.npu_scatter_nd_update_(values, indices, v) From 45220081159627ff4fba36a6b0b3ac10787214c8 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 02:51:19 +0800 Subject: [PATCH 05/12] feat: add explicit BF16 eager attention on Ascend --- tests/test_ascend_eager.py | 127 ++++++++++++++++++++++++++++++ tests/test_ascend_serving.py | 7 +- vllm_rlt/attention/backend.py | 14 +++- vllm_rlt/core/kv_cache_manager.py | 8 +- vllm_rlt/engine/stage_engine.py | 6 +- vllm_rlt/entrypoints/cli.py | 8 +- vllm_rlt/entrypoints/serve.py | 5 +- vllm_rlt/kernels/ascend_eager.py | 76 ++++++++++++++++++ 8 files changed, 235 insertions(+), 16 deletions(-) create mode 100644 tests/test_ascend_eager.py create mode 100644 vllm_rlt/kernels/ascend_eager.py diff --git a/tests/test_ascend_eager.py b/tests/test_ascend_eager.py new file mode 100644 index 00000000..3698f2d6 --- /dev/null +++ b/tests/test_ascend_eager.py @@ -0,0 +1,127 @@ +"""Causal/GQA semantics for the separately selected BF16 eager contract.""" + +import pytest +import torch + +from vllm_rlt.kernels.ascend_eager import AscendEagerAttention + + +@pytest.mark.parametrize("kv_heads", [1, 2, 4]) +def test_eager_groups_match_independent_causal_queries(kv_heads): + torch.manual_seed(730) + keys = torch.randn(8, 16, kv_heads, 64, dtype=torch.bfloat16) + values = torch.randn_like(keys) + q = torch.randn(7, 4, 64, dtype=torch.bfloat16) + lengths = torch.tensor([0, 1, 15, 16, 17, 3, 31], dtype=torch.int32) + tables = torch.tensor([[4, 1], [4, 1], [4, 1], [4, 1], [4, 1], [7, 2], [7, 2]]) + actual = AscendEagerAttention()(q, keys, values, tables, lengths) + expected = torch.zeros_like(q) + for row, length in enumerate(lengths.tolist()): + if not length: + continue + pos = torch.arange(length) + k = keys[tables[row, pos // 16], pos % 16].repeat_interleave(4 // kv_heads, dim=1) + v = values[tables[row, pos // 16], pos % 16].repeat_interleave(4 // kv_heads, dim=1) + scores = torch.einsum("hd,thd->ht", q[row], k) * (64**-0.5) + probs = scores.softmax(-1, dtype=torch.float32).to(q.dtype) + expected[row] = torch.einsum("ht,thd->hd", probs, v) + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + assert torch.isfinite(actual).all() + assert torch.count_nonzero(actual[0]) == 0 + + +@pytest.mark.npu +@pytest.mark.parametrize("chunk_size", [4096, 128]) +@torch.inference_mode() +def test_engine_matches_the_same_manual_prefill_chunks(chunk_size): + from tests.helpers import tiny_ouro_config + from vllm_rlt.models import OuroForCausalLM + 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 + + device = resolve_device("npu:0") + torch.npu.set_device(device) + torch.manual_seed(41) + model = OuroForCausalLM( + tiny_ouro_config( + hidden_size=256, intermediate_size=512, head_dim=64, max_position_embeddings=256 + ) + ).to(device=device, dtype=torch.bfloat16) + c = model.config + tokens = ([3, 7, 2, 11, 5, 17, 8] * 22)[:149] + cache = KVCacheManager( + c.num_hidden_layers, + c.num_key_value_heads, + c.head_dim, + 128, + 16, + 4, + device=device, + dtype=torch.bfloat16, + backend="ascend_eager", + ) + assert cache.allocate("manual", len(tokens) + 3) + expected, manual_states = [], [] + handle = model.model.norm.register_forward_hook( + lambda _m, _i, output: manual_states.append(output.detach().cpu().clone()) + ) + try: + for start in range(0, len(tokens), chunk_size): + end = min(start + chunk_size, len(tokens)) + hidden = model.prelude(torch.tensor(tokens[start:end], device=device)) + for depth in range(4): + hidden, _ = model.recurrent( + hidden, + ["manual"] * (end - start), + [depth] * (end - start), + list(range(start, end)), + cache, + compute_gate=False, + ) + for step in range(3): + expected.append(int(model.coda(hidden[-1:])[0].argmax())) + if step < 2: + pos = len(tokens) + step + hidden = model.prelude(torch.tensor([expected[-1]], device=device)) + for depth in range(4): + hidden, _ = model.recurrent( + hidden, ["manual"], [depth], [pos], cache, compute_gate=False + ) + finally: + handle.remove() + cache.free("manual") + actual_states = [] + handle = model.model.norm.register_forward_hook( + lambda _m, _i, output: actual_states.append(output.detach().cpu().clone()) + ) + engine = StageEngine( + model, + attention_backend="ascend_eager", + cache_config=CacheConfig(num_blocks=128), + scheduler_config=SchedulerConfig( + max_num_seqs=1, max_num_batched_tokens=4096, prefill_chunk_size=chunk_size + ), + ) + try: + engine.add_request( + "engine", + tokens, + SamplingParams(max_tokens=3, temperature=0, min_loops=4, max_loops=4, ignore_eos=True), + ) + result = None + while engine.has_unfinished_requests(): + for out in engine.step(): + if out.finished: + result = out + assert result is not None + assert result.token_ids == expected + assert result.exit_depths == [4] * 3 + assert len(actual_states) == len(manual_states) + for actual, reference in zip(actual_states, manual_states, strict=True): + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + assert engine.cache_manager.num_used_blocks == 0 + finally: + handle.remove() + engine.close() diff --git a/tests/test_ascend_serving.py b/tests/test_ascend_serving.py index 90f596c0..384eef37 100644 --- a/tests/test_ascend_serving.py +++ b/tests/test_ascend_serving.py @@ -15,13 +15,14 @@ pytestmark = pytest.mark.npu -def test_checkpoint_http_completions_stream_and_disconnect(tmp_path): +@pytest.mark.parametrize("backend", ["torch", "ascend_eager"]) +def test_checkpoint_http_completions_stream_and_disconnect(tmp_path, backend): import aiohttp source = os.environ.get("VLLM_RLT_OURO_CHECKPOINT") if not source: pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to the pinned local Ouro-1.4B checkpoint") - folder = Path(os.environ.get("VLLM_RLT_ASCEND_RESULTS", tmp_path)) + folder = Path(os.environ.get("VLLM_RLT_ASCEND_RESULTS", tmp_path)) / backend folder.mkdir(parents=True, exist_ok=True) with socket.socket() as listener: listener.bind(("127.0.0.1", 0)) @@ -39,7 +40,7 @@ def test_checkpoint_http_completions_stream_and_disconnect(tmp_path): "--dtype", "bfloat16", "--attention-backend", - "torch", + backend, "--num-blocks", "64", "--max-num-seqs", diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..48a06383 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -36,6 +36,12 @@ class AttentionBackend: def __call__(self, *args, **kwargs): return self.kernel(*args, **kwargs) + def attend_metadata(self, q, keys, values, batch): + method = getattr(self.kernel, "attend_metadata", None) + if method is not None: + return method(q, keys, values, batch) + return self(q, keys, values, batch.block_tables, batch.context_lengths) + def prefill(self, *args, **kwargs): return self.kernel.prefill(*args, **kwargs) @@ -68,11 +74,17 @@ def backend_capabilities(implementation) -> BackendCapabilities: def validate_backend_name(backend: str) -> None: - if backend not in {"auto", "torch", "triton", *FLASH_BACKENDS}: + if backend not in {"auto", "torch", "triton", "ascend_eager", *FLASH_BACKENDS}: raise ValueError("unknown attention backend") def _create_explicit(backend, device, dtype, head_dim, block_size): + if backend == "ascend_eager": + if device.type != "npu" or dtype != torch.bfloat16: + raise ValueError("ascend_eager requires an NPU and BF16") + from vllm_rlt.kernels.ascend_eager import AscendEagerAttention + + return AttentionBackend(AscendEagerAttention(), _TORCH_CAPABILITIES, {"backend": backend}) if backend == "triton": if device.type != "cuda": raise ValueError("the Triton attention backend requires a CUDA or ROCm device") diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 37c1352a..3a88f39e 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -755,12 +755,8 @@ def _attend_prepared( batch.cu_seqlens_q, batch.max_seqlen_q, ) - return self.attention( - q, - self.key_cache[:, layer], - self.value_cache[:, layer], - batch.block_tables, - batch.context_lengths, + return self.attention.attend_metadata( + q, self.key_cache[:, layer], self.value_cache[:, layer], batch ) @torch.no_grad() diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index dd13576a..e1b25f5c 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -73,8 +73,10 @@ def __init__( execution_config or ExecutionConfig(), scheduler_config, speculative_config ) if parameter.device.type == "npu": - if attention_backend != "torch": - raise ValueError("NPU execution requires explicit attention_backend='torch'") + if attention_backend not in ("torch", "ascend_eager"): + raise ValueError( + "NPU execution requires explicit attention_backend='torch' or 'ascend_eager'" + ) if config.to_dict()["model_type"] != "ouro" or config.total_ut_steps != 4: raise ValueError("NPU currently supports Ouro with four loops") if parameter.dtype != torch.bfloat16: diff --git a/vllm_rlt/entrypoints/cli.py b/vllm_rlt/entrypoints/cli.py index 5099ef26..29a96295 100644 --- a/vllm_rlt/entrypoints/cli.py +++ b/vllm_rlt/entrypoints/cli.py @@ -67,6 +67,7 @@ def main(): "--attention-backend", type=str.lower, choices=[ + "ascend_eager", "auto", "torch", "triton", @@ -100,8 +101,11 @@ def main(): parser.error("--toy uses built-in token ID prompts; omit --prompt") if not args.toy and not args.prompt: parser.error("provide --prompt or use --toy for a CPU smoke test") - if args.device.split(":", 1)[0] == "npu" and args.attention_backend != "torch": - parser.error("NPU execution requires explicit --attention-backend torch") + if args.device.split(":", 1)[0] == "npu" and args.attention_backend not in ( + "torch", + "ascend_eager", + ): + parser.error("NPU execution requires explicit --attention-backend torch or ascend_eager") profiling = profile_config_from_args(args) if profiling is not None: profiling.resolve_activities(args.device) diff --git a/vllm_rlt/entrypoints/serve.py b/vllm_rlt/entrypoints/serve.py index 713bab30..70a708f7 100644 --- a/vllm_rlt/entrypoints/serve.py +++ b/vllm_rlt/entrypoints/serve.py @@ -88,6 +88,7 @@ def main(): "--attention-backend", type=str.lower, choices=[ + "ascend_eager", "auto", "torch", "triton", @@ -115,8 +116,8 @@ def main(): limits = ServingLimits(**{name: getattr(args, name) for name in defaults}) if args.device == "cpu" and args.attention_backend not in ("auto", "torch"): parser.error("CPU execution requires --attention-backend torch or auto") - if args.device == "npu" and args.attention_backend != "torch": - parser.error("NPU execution requires explicit --attention-backend torch") + if args.device == "npu" and args.attention_backend not in ("torch", "ascend_eager"): + parser.error("NPU execution requires explicit --attention-backend torch or ascend_eager") torch.set_num_threads(args.cpu_threads) logging.basicConfig(level=logging.INFO) app = create_app( diff --git a/vllm_rlt/kernels/ascend_eager.py b/vllm_rlt/kernels/ascend_eager.py new file mode 100644 index 00000000..2ab08478 --- /dev/null +++ b/vllm_rlt/kernels/ascend_eager.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit BF16 eager attention on Ascend, matching released Ouro arithmetic.""" + +import math +import weakref + +import torch + + +class AscendEagerAttention: + """Batch each request/depth's queries; borrow only validated host metadata.""" + + def __init__(self): + self._metadata = None + self._groups = None + + def _plan(self, keys, lengths, device): + grouped = {} + for row, (key, length) in enumerate(zip(keys, lengths, strict=True)): + if length: + grouped.setdefault(key, []).append(row) + planned = [] + for rows in grouped.values(): + maximum = max(lengths[row] for row in rows) + indices = torch.tensor(rows, dtype=torch.long, device=device) + positions = torch.arange(maximum, device=device) + valid = positions.unsqueeze(0) < torch.tensor( + [lengths[row] for row in rows], device=device + ).unsqueeze(1) + mask = torch.zeros((len(rows), maximum), dtype=torch.bfloat16, device=device) + mask.masked_fill_(~valid, torch.finfo(torch.bfloat16).min) + initial = len(rows) == maximum and [lengths[row] for row in rows] == list( + range(1, maximum + 1) + ) + planned.append((rows[0], indices, positions, mask, initial)) + return planned + + def _forward(self, q, keys, values, tables, groups): + result = torch.zeros_like(q) + block_size = keys.shape[1] + head_groups = q.shape[1] // keys.shape[2] + for first, indices, positions, mask, initial in groups: + blocks = tables[first, positions // block_size] + offsets = positions % block_size + k, v = keys[blocks, offsets], values[blocks, offsets] + if head_groups != 1: + k = k.repeat_interleave(head_groups, dim=1) + v = v.repeat_interleave(head_groups, dim=1) + query = q.index_select(0, indices).transpose(0, 1).unsqueeze(0).contiguous() + key = k.transpose(0, 1).unsqueeze(0) + value = v.transpose(0, 1).unsqueeze(0) + # HF's initial V is a projection view; subsequent cache concatenation + # makes K/V contiguous. Preserve that operation layout on this path. + if not initial: + key, value = key.contiguous(), value.contiguous() + scores = torch.matmul(query, key.transpose(-2, -1)) * (1.0 / math.sqrt(q.shape[-1])) + # Keep the BF16 score/scaling and probability rounding boundaries. + scores = scores + mask[None, None, :, :] + probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype) + output = torch.matmul(probabilities, value).squeeze(0).transpose(0, 1) + result.index_copy_(0, indices, output) + return result + + def __call__(self, q, keys, values, tables, lengths): + groups = self._plan([tuple(row) for row in tables.tolist()], lengths.tolist(), q.device) + return self._forward(q, keys, values, tables, groups) + + def attend_metadata(self, q, keys, values, batch): + if self._metadata is None or self._metadata() is not batch: + self._groups = self._plan( + [(id(allocation), depth) for allocation, depth, _ in batch.rows], + [position + 1 for _, _, position in batch.rows], + q.device, + ) + self._metadata = weakref.ref(batch) + return self._forward(q, keys, values, batch.block_tables, self._groups) From 1461b190665a295dd97a7bca9e7a55f7f521a3ee Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 03:21:33 +0800 Subject: [PATCH 06/12] test: guard Ascend BF16 attention rounding contract --- tests/test_ascend_eager.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/test_ascend_eager.py b/tests/test_ascend_eager.py index 3698f2d6..37e5f6ce 100644 --- a/tests/test_ascend_eager.py +++ b/tests/test_ascend_eager.py @@ -35,11 +35,11 @@ def test_eager_groups_match_independent_causal_queries(kv_heads): @torch.inference_mode() def test_engine_matches_the_same_manual_prefill_chunks(chunk_size): from tests.helpers import tiny_ouro_config - from vllm_rlt.models import OuroForCausalLM 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.models import OuroForCausalLM device = resolve_device("npu:0") torch.npu.set_device(device) @@ -125,3 +125,30 @@ def test_engine_matches_the_same_manual_prefill_chunks(chunk_size): finally: handle.remove() engine.close() + + +@pytest.mark.parametrize("device_name", ["cpu", pytest.param("npu", marks=pytest.mark.npu)]) +@torch.inference_mode() +def test_bf16_probability_rounding_is_part_of_the_eager_contract(device_name): + from vllm_rlt.device import resolve_device + from vllm_rlt.kernels.paged_attention import torch_paged_attention + + device = resolve_device(device_name) + if device_name == "npu": + torch.npu.set_device(0) + keys = torch.zeros(1, 16, 1, 64, dtype=torch.bfloat16, device=device) + values = torch.zeros_like(keys) + keys[0, 0].fill_(1.03125) + keys[0, 1].fill_(1) + values[0, 0].fill_(128) + values[0, 1].fill_(-128) + query = torch.ones(1, 1, 64, dtype=torch.bfloat16, device=device) + table = torch.tensor([[0]], dtype=torch.int32, device=device) + lengths = torch.tensor([2], dtype=torch.int32, device=device) + # QK/sqrt(64) is [8.25,8]. BF16 probabilities are [.5625,.4375], + # so their weighted value is exactly16. FP32 probabilities retain + # [.5621765,.4378235] and round the final value to15.9375 instead. + actual = AscendEagerAttention()(query, keys, values, table, lengths) + torch.testing.assert_close(actual, torch.full_like(actual, 16), atol=0, rtol=0) + old_contract = torch_paged_attention(query, keys, values, table, lengths) + assert not torch.equal(actual, old_contract) From be037ed5cb9344b59830da2449b5590f3dd6c231 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 03:24:48 +0800 Subject: [PATCH 07/12] fix: keep native KV writes opt-in with Ascend eager --- tests/test_ascend_kv.py | 18 ++++++++++++++++-- vllm_rlt/core/kv_cache_manager.py | 2 +- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/tests/test_ascend_kv.py b/tests/test_ascend_kv.py index 8445e0f0..298d0c1a 100644 --- a/tests/test_ascend_kv.py +++ b/tests/test_ascend_kv.py @@ -10,11 +10,24 @@ @pytest.mark.parametrize("rows", [1, 3]) +@pytest.mark.parametrize("backend", ["torch", "ascend_eager"]) @torch.inference_mode() -def test_native_scatter_preserves_inputs_other_layers_and_reused_ids(rows): +def test_native_scatter_preserves_inputs_other_layers_and_reused_ids(rows, backend, monkeypatch): device = resolve_device("npu:0") torch.npu.set_device(device) - cache = KVCacheManager(2, 2, 64, 32, 16, 4, device=device, dtype=torch.bfloat16) + import torch_npu + + native = torch_npu.npu_scatter_nd_update_ + calls = [] + + def observed(*args, **kwargs): + calls.append(True) + return native(*args, **kwargs) + + monkeypatch.setattr(torch_npu, "npu_scatter_nd_update_", observed) + cache = KVCacheManager( + 2, 2, 64, 32, 16, 4, device=device, dtype=torch.bfloat16, backend=backend + ) cache.key_cache.zero_() cache.value_cache.zero_() assert cache.allocate("a", 17) @@ -49,3 +62,4 @@ def test_native_scatter_preserves_inputs_other_layers_and_reused_ids(rows): cache.free("b") assert cache.num_used_blocks == 0 torch.npu.synchronize() + assert len(calls) == (6 if backend == "ascend_eager" else 0) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 3a88f39e..2ee380eb 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -684,7 +684,7 @@ def _write_prepared( self._validate_tensor(v, len(batch.position_ids), "v") if not batch.rows: return - if self.device.type == "npu": + if self.device.type == "npu" and self.attention.info.get("backend") == "ascend_eager": if not hasattr(self, "_npu_writer"): from vllm_rlt.kernels.ascend_kv import AscendKVWriter From 765755769a941dae496fa8096c626f07c9776552 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 03:45:15 +0800 Subject: [PATCH 08/12] docs: record Ascend BF16 contract and measured validation --- README.md | 2 +- docs/recipes/ouro_ascend.md | 275 +++++++++++++++++------------------- 2 files changed, 132 insertions(+), 145 deletions(-) diff --git a/README.md b/README.md index 183cc29e..9b4f90db 100644 --- a/README.md +++ b/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; experimental single-device Ascend 910B3 BF16/fixed-four-loop synchronous Torch path, with paired quality regressions pending correction | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.md), [Ascend recipe](docs/recipes/ouro_ascend.md) | +| Ouro-1.4B | Native Ouro | Existing runtime and benchmark coverage under their recorded configurations; experimental single-device Ascend 910B3 BF16/fixed-four-loop synchronous path with explicit HF-style attention; selected full-prefill paired regression matches official HF (63/87 each, 87/87 exact sequences) | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.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) | diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md index 43d58995..659b8e3e 100644 --- a/docs/recipes/ouro_ascend.md +++ b/docs/recipes/ouro_ascend.md @@ -1,66 +1,67 @@ -# Ouro-1.4B on Ascend 910B3: initial synchronous support +# Ouro-1.4B on Ascend 910B3 -This path reuses the native Ouro model and shared engine on one Ascend NPU. -Select **--attention-backend torch explicitly**. Weights, ordinary activations -and KV storage remain BF16; the existing Torch attention computes QK, softmax -and PV intermediates in FP32 through Torch-NPU. This is not a fused CANN -attention integration and makes no acceleration claim. +This experimental path runs native Ouro on one NPU with BF16, synchronous +greedy decoding and four recurrent loops. Select +`--attention-backend ascend_eager` explicitly. It batches causal queries by +request/depth through Torch-NPU/CANN MatMul and Softmax primitives and uses +native ScatterND for paged KV writes. It preserves released Ouro HF eager +rounding: BF16 QK/scale, FP32 softmax rounded to BF16, then BF16 PV. -**Quality status:** the selected same-device paired screen is HF 63/87 versus -native 61/87 (losses: 203 and 881; gains: none). The frozen no-net-loss gate -failed. This implementation remains experimental and is not ready for -submission as a quality-qualified backend. +The older explicit `torch` backend remains available with its original +FP32 attention intermediates and indexed KV writes. Its arithmetic is a +different reference contract. Automatic backend selection is unchanged. ## Supported scope -- Ouro, fixed four-loop greedy inference, synchronous local execution, - LAST_EXITED depth-aware paged KV, prefill and incremental decode. -- Explicit num_blocks or kv_cache_memory_bytes; automatic NPU memory sizing - is not implemented. -- Cancellation, request ID reuse, EOS/length handling, HTTP completions and SSE - reuse the existing engine and serving lifecycle. -- Async scheduling, CUDA graphs, speculation, preemption, prefix caching, PD, - adaptive exit and stochastic sampling are rejected for this initial path. -- Qualification targets Ouro-1.4B on 910B3. Other checkpoints and Ascend devices - require separate validation. - -The shared scheduler, model and KV owner are retained. NPU submissions drain -before metadata buffers or request pages can be reused. CPU/CUDA execution -does not import the optional Torch-NPU package. Built-in profiling supports -CPU/CUDA; use torch_npu.profiler for NPU device traces. +- Native Ouro-1.4B, BF16 parameters/ordinary activations/KV, fixed four loops, + synchronous single-device greedy generation. +- Per-request/depth/layer paged KV; incremental/full allocation, causal prefill, + decode, cancellation and request-ID reuse. +- Existing spawned HTTP/SSE serving lifecycle. Explicit KV capacity is required. +- Async scheduling, graphs, speculation, PD, preemption, prefix caching, + adaptive depth and stochastic sampling are rejected for this initial path. +- Other checkpoints, Ascend models/devices and concurrent task quality need + separate qualification. The quality regression below uses a full prompt + forward; chunked generation can round differently. + +The shared scheduler/model/KV owner are retained. NPU submissions drain before +metadata or pages can be reused. CPU/CUDA do not import Torch-NPU unconditionally. +Use `torch_npu.profiler` for NPU device profiling; built-in profiling targets +CPU/CUDA. ## Environment and installation -The development environment is Linux aarch64, Python 3.11.13, Torch 2.9.0+cpu, -torch-npu 2.9.0, CANN 9.0.0, driver 26.1.0.b050 and Transformers 4.57.6. -Keep Torch and Torch-NPU matched; activate vendor CANN before the project Python. +Measured environment: Linux aarch64, Python 3.11.13, Torch 2.9.0+cpu, +Torch-NPU 2.9.0, CANN 9.0.0, driver 26.1.0.b050, Transformers 4.57.6, +Tokenizers 0.22.2 and NumPy 1.26.4. Keep the vendor Torch pair matched. ~~~bash source /path/to/cann/set_env.sh source /path/to/isolated-npu-environment/bin/activate -export ASCEND_RT_VISIBLE_DEVICES=4 # replace with your assigned physical NPU +export ASCEND_RT_VISIBLE_DEVICES=4 # use your assigned physical NPU python -c 'import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.get_device_name(0))' python -m pip install 'numpy==1.26.4' 'safetensors==0.8.0' 'huggingface-hub==0.36.2' \ 'transformers==4.57.6' 'tokenizers==0.22.2' 'aiohttp==3.14.3' setuptools wheel python -m pip install --no-deps --no-build-isolation -e . ~~~ -The no-deps install preserves the vendor Torch pair and avoids the project's -Linux Triton dependency. Neither Triton nor a native CUDA extension is required -for this explicit Torch backend. No InferRT or MindIE-SD integration is involved. +The no-deps install preserves vendor Torch and avoids the Linux Triton +dependency. No native CUDA extension, custom compiled kernel or InferRT is used. +Current upstream ScatterPA examples use an API/layout unavailable in this +Torch-NPU/CANN pair; the native writer retains the existing ND cache. -Use the complete ByteDance/Ouro-1.4B checkpoint at revision -574fa66cb8bf5abdc979642d01cf2b79b16bfab1. The model.safetensors SHA256 is -58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af. +Use ByteDance/Ouro-1.4B revision 574fa66cb8bf5abdc979642d01cf2b79b16bfab1. +Weights SHA256: +`58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. -## Start the service +## Start and call the service ~~~bash python -m vllm_rlt.entrypoints.serve \ --model /path/to/Ouro-1.4B --served-model-name ouro \ - --device npu --dtype bfloat16 --attention-backend torch \ - --num-blocks 64 --block-size 16 --max-num-seqs 2 \ - --max-num-batched-tokens 32 --port 8123 + --device npu --dtype bfloat16 --attention-backend ascend_eager \ + --num-blocks 1024 --block-size 16 --max-num-seqs 1 \ + --max-num-batched-tokens 4096 --prefill-chunk-size 4096 --port 8123 ~~~ ~~~bash @@ -70,115 +71,101 @@ curl http://127.0.0.1:8123/v1/completions \ -d '{"model":"ouro","prompt":"The capital of France is","max_tokens":4,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}' ~~~ -Add stream=true for SSE. Python callers use device="npu" and -attention_backend="torch"; the spawned worker registers Torch-NPU itself. -The assigned physical device becomes logical npu:0. Increase KV capacity for -longer prompts, output limits or concurrency. This small service command is a -smoke configuration, not a performance benchmark. +Add `stream=true` for SSE. Python callers select `device="npu"` and +`attention_backend="ascend_eager"`; the spawned worker registers Torch-NPU. +The assigned physical NPU becomes logical `npu:0`. Capacity and token budgets +must fit your prompts/outputs/concurrency; this command is a reproduction +configuration, not certification of every 4096-token input. -## Reproduce validation +Both `prefill_chunk_size` and `max_num_batched_tokens` must accommodate the +whole prompt to compare with official full-prefill. Raising only the latter +retains the scheduler's default 128-token chunks. + +## Validation (2026-11-03) + +Base 989921ac1d3897c5c43e27ec8cddbc91fa953194 rejects `--device npu`. +Runtime 45220081159627ff4fba36a6b0b3ac10787214c8 supplies the new backend. +Final runtime be037ed5cb9344b59830da2449b5590f3dd6c231 adds an arithmetic +regression and makes native writes exclusive to the new backend; selected +eager execution and original Torch execution were separately revalidated. +Later documentation changes preserve these runtime/test contents. ~~~bash -python -m pip install pytest ruff pre-commit -pre-commit run --all-files +pre-commit run --all-files --verbose export VLLM_RLT_OURO_CHECKPOINT=/path/to/Ouro-1.4B -export VLLM_RLT_ASCEND_RESULTS=/path/to/validation-results +export VLLM_RLT_ASCEND_RESULTS=/path/to/results python -m pytest tests/test_ascend.py tests/test_ascend_checkpoint.py \ - tests/test_ascend_serving.py --run-npu -q + tests/test_ascend_serving.py tests/test_ascend_kv.py \ + tests/test_ascend_eager.py --run-npu -q ~~~ -The NPU flag requires assigned ASCEND_RT_VISIBLE_DEVICES; CPU collection skips -device tests. Checks cover a CPU attention oracle, GQA and causal page -boundaries, incremental/full KV allocation, four-depth state ownership, -chunking, cancellation/reuse, EOS/length behavior, pinned checkpoint generation -and a real spawned HTTP/SSE worker. - -Metadata-only allocation changes require exact hidden/logit/KV equality. -The BF16 diagnostic bounds remain relative L2 at most 2% and elementwise -atol=0.02, rtol=0.02. These are this validation's bounds, not a universal quality -guarantee. Task quality needs separately pinned paired evaluation against the -official code; short generation does not establish broader quality. -Performance characterization should retain input lengths, concurrency, -TTFT/TPOT/throughput, memory and repetitions; the existing -[serving benchmark guide](../serving.md#reuse-the-upstream-benchmark) describes -the compatible upstream client. - -## Fused attention remains separate - -An experimental BF16 paged npu_fused_infer_attention_score adapter is outside -this change. It failed frozen whole-model numerical checks and lost a correct -answer in a ten-question smoke. Promoting attention intermediates in -npu_fusion_attention improved local mathematical agreement but still failed -the original whole-model joint gate. Those failures have not been relaxed or -hidden; this recipe does not qualify that fused backend. - -This contribution covers a bounded part of -[#69 Ascend support](https://github.com/InferMatrixLab/vllm-rlt/issues/69). -It does not complete CANN fused adaptation or other Ouro runtime features on -Ascend. Measured validation outcomes and remaining blockers are recorded below. - -## Measured validation (2026-11-03) - -Runtime source: 4734e1963bdf7fd4e0fd66b2612511d6d4173550, against base -989921ac1d3897c5c43e27ec8cddbc91fa953194. Later recipe updates do not change -the measured runtime sources. - -- Ruff lint/format and the CPU regression hook passed. -- Opt-in real 910B3 tests: 10 passed in 44.41s, including actual HTTP/SSE. -- Incremental/preallocated real-weight KV differential: 216 finite states, all - exactly equal; attention/chunking diagnostics retain their original joint bounds. -- Same-device paired selected GSM8K-87: official 63/87, native 61/87; - losses [203, 881], gains []; 63/87 exact token sequences. - Parse failures: {'hf_eager': 19, 'rlt_target': 19}; length limits: {'hf_eager': 0, 'rlt_target': 0}; - execution limits: {'hf_eager': 0, 'rlt_target': 0}. -- Each pair runs on the same physical 910B3; all 87 cases pool three cards. - The original ten are regression smoke; 77 cases were held out. This is - selected-case evidence, not full GSM8K or adaptive/concurrent quality. - -The quality run uses official unmodified HF eager, Transformers 4.57.6, -Tokenizers 0.22.2, lm-eval 0.4.9.2 strict scoring, three-shot raw CoT, -greedy four-loop BF16, max 1024 output tokens and source-ranked seed0 IDs. -Dataset revision is 740312add88f781978c0658806c59bc2815b9866. - -The table below is a separate **engine-only short-burst** characterization: -8 output tokens/request, ignore EOS, one shape-specific warmup and three -unprofiled repetitions. Values are medians of per-burst request means; output -throughput includes prefill and drain. TTFT starts before engine admission. -TPOT is first-to-last observed output time divided by seven. - -| Input | Concurrency | Output tok/s (range) | TTFT ms | TPOT ms | Peak allocated / reserved GiB | -|---|---:|---:|---:|---:|---:| -| 16 | 1 | 2.726 (2.700–2.901) | 1238.01 | 241.94 | 4.970 / 5.309 | -| 16 | 2 | 3.793 (3.657–3.810) | 2122.58 | 298.19 | 4.970 / 5.311 | -| 16 | 8 | 5.031 (5.016–5.086) | 7980.90 | 675.01 | 4.973 / 5.314 | -| 64 | 1 | 1.367 (1.342–1.372) | 4176.89 | 242.77 | 4.972 / 5.314 | -| 64 | 2 | 1.590 (1.569–1.599) | 7964.07 | 300.45 | 4.974 / 5.314 | -| 64 | 8 | 1.789 (1.786–1.793) | 23417.86 | 1756.95 | 4.984 / 5.314 | -| 256 | 1 | 0.370 (0.364–0.373) | 19472.72 | 307.45 | 4.978 / 5.318 | -| 256 | 2 | 0.372 (0.371–0.377) | 39959.74 | 431.61 | 4.988 / 5.318 | -| 256 | 8 | 0.379 (0.378–0.380) | 125114.88 | 6228.64 | 4.988 / 5.318 | - -Configuration: 768 blocks, block size 16, max 8 sequences / 256 batched tokens, -incremental allocation, no async/graphs/UVA/prefix caching/speculation. -Measured requests: 99, output tokens: 792, failures: 0. Live logical KV blocks -after each completed burst: 0. -Torch allocator peaks include model/cache/workspace and exclude some driver -or runtime memory. Other cards ran independent quality shards; CPU/NUMA -affinity was not pinned. No steady-state throughput, HTTP/p99, accelerator -hardware-only comparison or acceleration claim is made. - -Separate ops-only/full profiles capture a complete warmed second-token decode -iteration, with actual NPU device kernels. Full mode enables shapes, stacks -and memory; exported shapes, readable Python stacks and memory CSV rows -were checked separately from performance timing. -Full prefill plus all eight output tokens is outside these diagnostic captures. - -The conservative Torch path has substantial prefill/dispatch overhead. Fused -attention optimization remains a later step and must retain the quality gates. - -The no-net-loss quality criterion failed on this selected set. Case 203 -changes a quantity from 7 to 5 and produces 27 instead of 31. Case 881 writes -'The answer is 10 + 6 = 16'; the unchanged strict scorer extracts 10. It remains -a failure under the frozen protocol. Neither the scorer nor the gate was -relaxed. A diagnostic HF-style BF16 attention alternative also failed the -original numerical bounds and was not incorporated. +- Ruff/pre-commit passed; CPU 926 passed / 602 skipped. Device cases were run + separately:22 passed, comprising18 NPU and 4 CPU cases. +- Final eager vs unmodified official HF:216 selected states and2016 layer + events exactly equal; restored Torch vs original Torch is also exact. +- The arithmetic regression distinguishes an analytical output16 in BF16 + eager from15.9375 in the original FP32-intermediate path, on CPU andNPU. +- Real spawned health/completions/SSE, termination/usage, disconnect recovery + and shutdown pass for both explicit backends. +- At4522008,87 selected same-device HF/native pairs score63/87 each, + losses[]/gains[], with87/87 complete token sequences exactly equal. + Parse failures19 each; length/execution limits0 each. Every pair shares a + physical910B3, pooled across cards4/5/6. These are reused regression cases, + not a held-out/full-GSM8K estimate. No threshold or scorer was relaxed. +- Quality uses one full-prefill chunk4096 and4096 token budget. Chunked128, + concurrent/adaptive task quality and other models remain unqualified. + +Numerical bounds remain joint relative L2<=2% and atol=rtol=0.02. +The new explicit backend is checked against released HF BF16 eager. +It still fails 209/216 states against the old FP32-intermediate Torch reference; +that failure remains recorded and is not called passed. + +The earlier63/87 HF vs61/87 native observation combined different arithmetic +with native default 128-token chunking. Changing either alone did not fix +question 203. Matching both restores203/881/526 and the full87 regression. +Manual question 203 replay has15 exact logits and192 exact populated KV hashes. + +## Measured performance + +Same physical910B3, profiler disabled, P64/C1 orC8, fixed4 loops, ignore EOS, +exactly 8 outputs/request,768 blocks,max 8 sequences / 256 batched tokens,explicit +128-token chunks. One warmup/shape,3repetitions/run; A/C/C/A gives 6 observations +per arm/shape. TTFT/TPOT are medians of per-burst request means, not pooled +individual-request percentiles. Throughput includes prefill/drain; no HTTP. + +|Concurrency|Old Torch tok/s|New backend tok/s|Old/new TTFT ms|Old/new TPOT ms| +|---:|---:|---:|---:|---:| +|1|1.326|4.711|4266.35 /221.03|246.81 /210.99| +|8|1.720|12.778|24363.00 /702.72|1818.05 /583.61| + +Pair-specific throughput ratios are 3.47–3.81x at C1 and 7.11–8.05x at C8, +with non-overlapping observed ranges. This combined backend/arithmetic change +does not isolate ScatterND. A KV-only A/B regressed 1.23% and 6.28%; +the final runtime therefore keeps original writes on the Torch path. +New backend native-writing optimality is not established. + +Final runtime be037ed confirmation on the sameNPU gives medians 4.752tok/s +at C1 and13.705tok/s at C8 (three reps each, measured after ABBA). +Allocated peaks are about 4.97–4.99GiB; reserved peaks about 5.33GiB. +All bursts finish with 0 live logical KV blocks and no failed repetitions. + +A separate 9-shape candidate characterization covers 16/64/256 input tokens and +concurrency 1/2/8. These are short bursts, not saturation/HTTP/SLO or hardware-only +claims. CPU/NUMA placement is not pinned; independent CPU/device checks overlap +part of the shared-host timing. Raw ranges and controls must accompany reuse. + +Separate ops-only/full warmed decode traces confirm actual native ScatterND +dispatch and 96 softmax calls. Full mode has readable shapes, stacks and memory +exports. Complete warmed request traces additionally cover prefill and all 8 +outputs. Profile durations are diagnostic and excluded from speed claims. + +## Fused operators remain experimental + +BF16 FIA and promoted FP32 FlashAttentionScore did not pass earlier whole-model +gates. Native RMSNorm changes also failed selected states despite local speed +gains; SwiGLU/RotaryMul candidates were not adopted. +The current backend uses native primitives and batching, not fused attention. + +This is a bounded contribution toward +[#69](https://github.com/InferMatrixLab/vllm-rlt/issues/69), with task quality +qualified only for the stated regression configuration. From 8bb68fb5e89518ee7971edad59202bef6e52ef8e Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 09:59:54 +0800 Subject: [PATCH 09/12] fix: remove Ascend SKU assumptions and clarify backend contracts --- README.md | 2 +- docs/recipes/ouro_ascend.md | 8 +++- tests/test_ascend.py | 3 +- tests/test_ascend_checkpoint.py | 6 ++- tests/test_attention_module.py | 70 +++++++++++++++++++++++++++++++ vllm_rlt/attention/backend.py | 5 +++ vllm_rlt/core/kv_cache_manager.py | 2 +- 7 files changed, 89 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 9b4f90db..569373db 100644 --- a/README.md +++ b/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; experimental single-device Ascend 910B3 BF16/fixed-four-loop synchronous path with explicit HF-style attention; selected full-prefill paired regression matches official HF (63/87 each, 87/87 exact sequences) | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.md), [Ascend recipe](docs/recipes/ouro_ascend.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) | diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md index 659b8e3e..124855b7 100644 --- a/docs/recipes/ouro_ascend.md +++ b/docs/recipes/ouro_ascend.md @@ -1,4 +1,4 @@ -# Ouro-1.4B on Ascend 910B3 +# Ouro-1.4B on Ascend This experimental path runs native Ouro on one NPU with BF16, synchronous greedy decoding and four recurrent loops. Select @@ -29,6 +29,12 @@ metadata or pages can be reused. CPU/CUDA do not import Torch-NPU unconditionall Use `torch_npu.profiler` for NPU device profiling; built-in profiling targets CPU/CUDA. +Hardware selection does not inspect or whitelist device model names. +The backend requires BF16 and the Torch-NPU/CANN operations used below. +The measurements in this recipe were collected on 910B3; other Ascend devices +can run the same opt-in checks without a model-name assertion. Their numerical +and performance results require validation on the actual device/runtime pair. + ## Environment and installation Measured environment: Linux aarch64, Python 3.11.13, Torch 2.9.0+cpu, diff --git a/tests/test_ascend.py b/tests/test_ascend.py index 4632c7f1..fc4209d7 100644 --- a/tests/test_ascend.py +++ b/tests/test_ascend.py @@ -1,4 +1,4 @@ -"""Opt-in synchronous BF16 Ouro checks using explicit Torch attention on 910B3.""" +"""Opt-in synchronous BF16 Ouro checks using explicit Torch attention on Ascend.""" import pytest import torch @@ -18,7 +18,6 @@ def npu_device(): resolve_device("npu") torch.npu.set_device(0) - assert torch.npu.get_device_name(0) == "Ascend910B3" yield torch.device("npu:0") torch.npu.synchronize() torch.npu.empty_cache() diff --git a/tests/test_ascend_checkpoint.py b/tests/test_ascend_checkpoint.py index 6b08c1d3..6beb91b2 100644 --- a/tests/test_ascend_checkpoint.py +++ b/tests/test_ascend_checkpoint.py @@ -24,9 +24,11 @@ def checkpoint(): pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to the pinned local Ouro-1.4B checkpoint") resolve_device("npu") torch.npu.set_device(0) - assert torch.npu.get_device_name(0) == "Ascend910B3" + checksum = hashlib.sha256() with (Path(source) / "model.safetensors").open("rb") as stream: - digest = hashlib.file_digest(stream, "sha256").hexdigest() + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + checksum.update(block) + digest = checksum.hexdigest() assert digest == "58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af" model = OuroForCausalLM.from_pretrained(source, device="npu", dtype=torch.bfloat16) assert all(p.dtype == torch.bfloat16 for p in model.parameters()) diff --git a/tests/test_attention_module.py b/tests/test_attention_module.py index 5c4de6dd..823cf322 100644 --- a/tests/test_attention_module.py +++ b/tests/test_attention_module.py @@ -155,3 +155,73 @@ def stage(values, dtype): assert len({batch.cu_seqlens_q.data_ptr() for batch in batches}) == 1 assert batches[0].cu_seqlens_q.tolist() == [0, 2, 3] assert batches[0].max_seqlen_q == 2 + + +def test_prepared_attention_keeps_the_legacy_kernel_signature(): + cache = KVCacheManager(1, 1, 4, 4, 2, 1) + assert cache.allocate("a", 3) + batch = cache._prepare_batch(["a"] * 3, [0] * 3, [0, 1, 2]) + keys = torch.zeros(3, 1, 4) + values = torch.arange(12, dtype=torch.float32).reshape(3, 1, 4) + query = torch.zeros(3, 2, 4) + cache._write_prepared(0, batch, keys, values) + kernel = cache.attention.kernel + expected = kernel( + query, + cache.key_cache[:, 0], + cache.value_cache[:, 0], + batch.block_tables, + batch.context_lengths, + ) + seen = [] + + def legacy(*args): + seen.append(args) + return kernel(*args) + + cache.attention.kernel = legacy + actual = cache._attend_prepared(0, batch, query) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert len(seen) == 1 and len(seen[0]) == 5 + assert seen[0][0] is query + assert seen[0][3] is batch.block_tables + assert seen[0][4] is batch.context_lengths + cache.free("a") + + +def test_metadata_attention_borrows_only_a_live_validated_descriptor(): + cache = KVCacheManager(1, 1, 4, 4, 2, 1) + assert cache.allocate("a", 3) + batch = cache._prepare_batch(["a"] * 3, [0] * 3, [0, 1, 2]) + keys = torch.zeros(3, 1, 4) + values = torch.arange(12, dtype=torch.float32).reshape(3, 1, 4) + query = torch.zeros(3, 2, 4) + cache._write_prepared(0, batch, keys, values) + legacy = cache.attention.kernel + expected = legacy( + query, + cache.key_cache[:, 0], + cache.value_cache[:, 0], + batch.block_tables, + batch.context_lengths, + ) + seen = [] + + class MetadataKernel: + def __call__(self, *args): + pytest.fail("metadata-aware kernels must not read device metadata back") + + def attend_metadata(self, q, k, v, descriptor): + seen.append(descriptor) + return legacy(q, k, v, descriptor.block_tables, descriptor.context_lengths) + + cache.attention.kernel = MetadataKernel() + actual = cache._attend_prepared(0, batch, query) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert len(seen) == 1 and seen[0] is batch + cache.free("a") + assert cache.allocate("a", 3) + with pytest.raises((ValueError, RuntimeError)): + cache._attend_prepared(0, batch, query) + assert len(seen) == 1 + cache.free("a") diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 48a06383..e04d9bac 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -37,6 +37,11 @@ def __call__(self, *args, **kwargs): return self.kernel(*args, **kwargs) def attend_metadata(self, q, keys, values, batch): + """Borrow validated host metadata when supported; retain the legacy tensor API. + + Metadata-aware kernels can reuse host planning across layers without + reading device tables or lengths back to the host. + """ method = getattr(self.kernel, "attend_metadata", None) if method is not None: return method(q, keys, values, batch) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 2ee380eb..67b21e2e 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -684,7 +684,7 @@ def _write_prepared( self._validate_tensor(v, len(batch.position_ids), "v") if not batch.rows: return - if self.device.type == "npu" and self.attention.info.get("backend") == "ascend_eager": + if self.device.type == "npu" and self.backend == "ascend_eager": if not hasattr(self, "_npu_writer"): from vllm_rlt.kernels.ascend_kv import AscendKVWriter From f3ce48a026796503bec344f10d5d94176791689d Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 10:05:06 +0800 Subject: [PATCH 10/12] docs: distinguish portability checks from historical measurements --- docs/recipes/ouro_ascend.md | 82 ++++++++++++++++++++++++------------- 1 file changed, 53 insertions(+), 29 deletions(-) diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md index 124855b7..bedfeeb7 100644 --- a/docs/recipes/ouro_ascend.md +++ b/docs/recipes/ouro_ascend.md @@ -89,12 +89,20 @@ retains the scheduler's default 128-token chunks. ## Validation (2026-11-03) -Base 989921ac1d3897c5c43e27ec8cddbc91fa953194 rejects `--device npu`. -Runtime 45220081159627ff4fba36a6b0b3ac10787214c8 supplies the new backend. -Final runtime be037ed5cb9344b59830da2449b5590f3dd6c231 adds an arithmetic -regression and makes native writes exclusive to the new backend; selected -eager execution and original Torch execution were separately revalidated. -Later documentation changes preserve these runtime/test contents. +Recorded sources: + +- Base `989921a` rejects `--device npu`. +- Backend revision `4522008` was used for the paired 87-question quality + campaign and primary performance comparison below. +- Routing revision `be037ed` adds an arithmetic regression and makes native + writes exclusive to the new backend. Selected eager and original Torch + execution were separately revalidated at that revision. +- Portability/interface revision `8bb68fb` removes device-name assertions, + supports Python 3.10 checksum calculation, and adds legacy/metadata-aware + attention dispatch regressions. + +Each result retains its measured source. Earlier quality, numerical and +performance campaigns are not relabeled as tests of a newer revision. ~~~bash pre-commit run --all-files --verbose @@ -105,34 +113,50 @@ python -m pytest tests/test_ascend.py tests/test_ascend_checkpoint.py \ tests/test_ascend_eager.py --run-npu -q ~~~ -- Ruff/pre-commit passed; CPU 926 passed / 602 skipped. Device cases were run - separately:22 passed, comprising18 NPU and 4 CPU cases. -- Final eager vs unmodified official HF:216 selected states and2016 layer - events exactly equal; restored Torch vs original Torch is also exact. -- The arithmetic regression distinguishes an analytical output16 in BF16 - eager from15.9375 in the original FP32-intermediate path, on CPU andNPU. -- Real spawned health/completions/SSE, termination/usage, disconnect recovery - and shutdown pass for both explicit backends. -- At4522008,87 selected same-device HF/native pairs score63/87 each, - losses[]/gains[], with87/87 complete token sequences exactly equal. - Parse failures19 each; length/execution limits0 each. Every pair shares a - physical910B3, pooled across cards4/5/6. These are reused regression cases, - not a held-out/full-GSM8K estimate. No threshold or scorer was relaxed. -- Quality uses one full-prefill chunk4096 and4096 token budget. Chunked128, +### Portability/interface checks at `8bb68fb` + +- Ruff/pre-commit passed; CPU **928 passed / 602 skipped**. +- Assigned-device suite: **22 passed**, comprising 18 NPU and 4 CPU cases. + Both explicit backends pass spawned HTTP/SSE, termination/usage, disconnect + recovery and shutdown checks. +- Legacy kernels retain their five tensor arguments; metadata-aware kernels + receive the validated descriptor, and stale descriptors remain rejected. +- The actual fixture checksum block was also executed successfully under + Python 3.10.12, where `hashlib.file_digest` is unavailable. +- Arithmetic kernels are unchanged by this revision. The full 87-question + campaign and whole-model HF layer-event comparison were not repeated here. + +### Earlier numerical and quality evidence + +- At `be037ed`, eager versus unmodified official HF has 216 selected states + and 2016 layer events exactly equal. Restored Torch versus original Torch + is also exact. +- The analytical rounding regression distinguishes output 16 in BF16 eager + from 15.9375 in the original FP32-intermediate path, on CPU and NPU. +- At `4522008`, 87 selected same-device HF/native pairs score 63/87 each, + with no losses or gains and 87/87 complete token sequences exactly equal. + Parse failures are 19 per arm; length and execution limits are zero. + Each pair shares a physical 910B3, pooled across three cards. These are + reused regression cases, not a held-out or full-GSM8K estimate. +- Quality uses one full-prefill chunk and token budget of 4096. Chunked, concurrent/adaptive task quality and other models remain unqualified. -Numerical bounds remain joint relative L2<=2% and atol=rtol=0.02. -The new explicit backend is checked against released HF BF16 eager. -It still fails 209/216 states against the old FP32-intermediate Torch reference; -that failure remains recorded and is not called passed. +Numerical bounds remain joint relative L2 <= 2% and atol=rtol=0.02. +The explicit BF16 backend targets released HF eager. At `be037ed`, it still +fails 209/216 states against the old FP32-intermediate Torch reference; +that historical failure remains recorded and is not relabeled as passed. -The earlier63/87 HF vs61/87 native observation combined different arithmetic -with native default 128-token chunking. Changing either alone did not fix -question 203. Matching both restores203/881/526 and the full87 regression. -Manual question 203 replay has15 exact logits and192 exact populated KV hashes. +The earlier HF 63/87 versus native 61/87 observation combined different +arithmetic with native default 128-token chunking. Changing either alone did +not fix question 203. Matching both restores questions 203/881/526 and the +87-question regression. The manual question 203 replay has 15 exact logits +and 192 exact populated KV hashes. ## Measured performance +The primary comparison uses original Torch at `4734e19` and the explicit +backend at `4522008`; the later confirmation uses `be037ed`. + Same physical910B3, profiler disabled, P64/C1 orC8, fixed4 loops, ignore EOS, exactly 8 outputs/request,768 blocks,max 8 sequences / 256 batched tokens,explicit 128-token chunks. One warmup/shape,3repetitions/run; A/C/C/A gives 6 observations @@ -150,7 +174,7 @@ does not isolate ScatterND. A KV-only A/B regressed 1.23% and 6.28%; the final runtime therefore keeps original writes on the Torch path. New backend native-writing optimality is not established. -Final runtime be037ed confirmation on the sameNPU gives medians 4.752tok/s +Routing-revision be037ed confirmation on the sameNPU gives medians 4.752tok/s at C1 and13.705tok/s at C8 (three reps each, measured after ABBA). Allocated peaks are about 4.97–4.99GiB; reserved peaks about 5.33GiB. All bursts finish with 0 live logical KV blocks and no failed repetitions. From 9b14425dcca812bd07583965c24d12e72516f25e Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 10:21:22 +0800 Subject: [PATCH 11/12] test: reuse native checkpoint loading on Ascend --- tests/test_ascend_checkpoint.py | 91 ++++++++++++++------------------- 1 file changed, 38 insertions(+), 53 deletions(-) diff --git a/tests/test_ascend_checkpoint.py b/tests/test_ascend_checkpoint.py index 6beb91b2..c26e8fe8 100644 --- a/tests/test_ascend_checkpoint.py +++ b/tests/test_ascend_checkpoint.py @@ -1,7 +1,5 @@ -"""Pinned real-weight NPU checks for explicit Torch attention, not fused CANN.""" +"""Opt-in local-checkpoint KV and lifecycle regressions on Ascend.""" -import hashlib -import json import os from pathlib import Path @@ -21,29 +19,17 @@ def checkpoint(): source = os.environ.get("VLLM_RLT_OURO_CHECKPOINT") if not source: - pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to the pinned local Ouro-1.4B checkpoint") + 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) - checksum = hashlib.sha256() - with (Path(source) / "model.safetensors").open("rb") as stream: - for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): - checksum.update(block) - digest = checksum.hexdigest() - assert digest == "58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af" 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() -def save_record(name, record): - folder = os.environ.get("VLLM_RLT_ASCEND_RESULTS") - if folder: - path = Path(folder) - path.mkdir(parents=True, exist_ok=True) - (path / name).write_text(json.dumps(record, indent=2)) - - @torch.inference_mode() def test_checkpoint_incremental_pages_preserve_hidden_logits_and_kv(checkpoint): model, tokenizer = checkpoint @@ -62,54 +48,53 @@ def test_checkpoint_incremental_pages_preserve_hidden_logits_and_kv(checkpoint): 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) - save_record( - "checkpoint-numerics.json", - dict( - tokens=tokens, - states=len(pairs), - exact=all(torch.equal(a, b) for a, b in pairs.values()), - finite=all( - bool(torch.isfinite(a).all() and torch.isfinite(b).all()) for a, b in pairs.values() - ), - ), - ) 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_generation_and_cleanup(checkpoint): +def test_checkpoint_greedy_replay_and_request_cleanup(checkpoint, backend): model, tokenizer = checkpoint engine = StageEngine( model, - attention_backend="torch", + 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"] - finished = {} + 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: - for index, prompt in enumerate(prompts): - engine.add_request( - str(index), - tokenizer.encode(prompt, add_special_tokens=False), - SamplingParams(max_tokens=4, temperature=0, ignore_eos=True), - ) - for _ in range(200): - for output in engine.step(): - if output.finished: - assert output.exit_depths == [4] * 4 - assert output.finish_reason == "length" - finished[output.request_id] = dict( - tokens=output.token_ids, text=tokenizer.decode(output.token_ids) - ) - if not engine.has_unfinished_requests(): - break - assert set(finished) == {"0", "1"} - assert finished["0"]["tokens"] == [7042, 30, 198, 198] - assert finished["1"]["tokens"] == [827, 30, 198, 198] - assert engine.cache_manager.num_used_blocks == 0 - save_record("checkpoint-generation.json", finished) + # 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() From d9217196a40a754f4e7d967f47dac79765236f36 Mon Sep 17 00:00:00 2001 From: KrystalRay Date: Tue, 3 Nov 2026 10:21:22 +0800 Subject: [PATCH 12/12] docs: simplify the Ascend setup recipe --- docs/recipes/ouro_ascend.md | 224 ++++++++---------------------------- 1 file changed, 47 insertions(+), 177 deletions(-) diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md index bedfeeb7..f5e280e7 100644 --- a/docs/recipes/ouro_ascend.md +++ b/docs/recipes/ouro_ascend.md @@ -1,201 +1,71 @@ -# Ouro-1.4B on Ascend - -This experimental path runs native Ouro on one NPU with BF16, synchronous -greedy decoding and four recurrent loops. Select -`--attention-backend ascend_eager` explicitly. It batches causal queries by -request/depth through Torch-NPU/CANN MatMul and Softmax primitives and uses -native ScatterND for paged KV writes. It preserves released Ouro HF eager -rounding: BF16 QK/scale, FP32 softmax rounded to BF16, then BF16 PV. - -The older explicit `torch` backend remains available with its original -FP32 attention intermediates and indexed KV writes. Its arithmetic is a -different reference contract. Automatic backend selection is unchanged. - -## Supported scope - -- Native Ouro-1.4B, BF16 parameters/ordinary activations/KV, fixed four loops, - synchronous single-device greedy generation. -- Per-request/depth/layer paged KV; incremental/full allocation, causal prefill, - decode, cancellation and request-ID reuse. -- Existing spawned HTTP/SSE serving lifecycle. Explicit KV capacity is required. -- Async scheduling, graphs, speculation, PD, preemption, prefix caching, - adaptive depth and stochastic sampling are rejected for this initial path. -- Other checkpoints, Ascend models/devices and concurrent task quality need - separate qualification. The quality regression below uses a full prompt - forward; chunked generation can round differently. - -The shared scheduler/model/KV owner are retained. NPU submissions drain before -metadata or pages can be reused. CPU/CUDA do not import Torch-NPU unconditionally. -Use `torch_npu.profiler` for NPU device profiling; built-in profiling targets -CPU/CUDA. - -Hardware selection does not inspect or whitelist device model names. -The backend requires BF16 and the Torch-NPU/CANN operations used below. -The measurements in this recipe were collected on 910B3; other Ascend devices -can run the same opt-in checks without a model-name assertion. Their numerical -and performance results require validation on the actual device/runtime pair. - -## Environment and installation - -Measured environment: Linux aarch64, Python 3.11.13, Torch 2.9.0+cpu, -Torch-NPU 2.9.0, CANN 9.0.0, driver 26.1.0.b050, Transformers 4.57.6, -Tokenizers 0.22.2 and NumPy 1.26.4. Keep the vendor Torch pair matched. +# 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/isolated-npu-environment/bin/activate -export ASCEND_RT_VISIBLE_DEVICES=4 # use your assigned physical NPU -python -c 'import torch, torch_npu; print(torch.__version__, torch_npu.__version__); print(torch.npu.get_device_name(0))' -python -m pip install 'numpy==1.26.4' 'safetensors==0.8.0' 'huggingface-hub==0.36.2' \ - 'transformers==4.57.6' 'tokenizers==0.22.2' 'aiohttp==3.14.3' setuptools wheel +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 ~~~ -The no-deps install preserves vendor Torch and avoids the Linux Triton -dependency. No native CUDA extension, custom compiled kernel or InferRT is used. -Current upstream ScatterPA examples use an API/layout unavailable in this -Torch-NPU/CANN pair; the native writer retains the existing ND cache. +Use the checkpoint's own config and tokenizer. Model loading uses the shared +native Ouro loader, including its supported sharded checkpoint layout. -Use ByteDance/Ouro-1.4B revision 574fa66cb8bf5abdc979642d01cf2b79b16bfab1. -Weights SHA256: -`58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. +## 2. Start the server -## Start and call the service +Select an allocated physical NPU using `ASCEND_RT_VISIBLE_DEVICES`. ~~~bash +ASCEND_RT_VISIBLE_DEVICES= \ python -m vllm_rlt.entrypoints.serve \ - --model /path/to/Ouro-1.4B --served-model-name ouro \ + --model "$MODEL" --tokenizer "$MODEL" --served-model-name ouro \ --device npu --dtype bfloat16 --attention-backend ascend_eager \ - --num-blocks 1024 --block-size 16 --max-num-seqs 1 \ - --max-num-batched-tokens 4096 --prefill-chunk-size 4096 --port 8123 + --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 http://127.0.0.1:8123/health -curl http://127.0.0.1:8123/v1/completions \ +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":4,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}' + -d '{"model":"ouro","prompt":"The capital of France is","max_tokens":8,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}' ~~~ -Add `stream=true` for SSE. Python callers select `device="npu"` and -`attention_backend="ascend_eager"`; the spawned worker registers Torch-NPU. -The assigned physical NPU becomes logical `npu:0`. Capacity and token budgets -must fit your prompts/outputs/concurrency; this command is a reproduction -configuration, not certification of every 4096-token input. +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. -Both `prefill_chunk_size` and `max_num_batched_tokens` must accommodate the -whole prompt to compare with official full-prefill. Raising only the latter -retains the scheduler's default 128-token chunks. +## 4. Run checks -## Validation (2026-11-03) - -Recorded sources: - -- Base `989921a` rejects `--device npu`. -- Backend revision `4522008` was used for the paired 87-question quality - campaign and primary performance comparison below. -- Routing revision `be037ed` adds an arithmetic regression and makes native - writes exclusive to the new backend. Selected eager and original Torch - execution were separately revalidated at that revision. -- Portability/interface revision `8bb68fb` removes device-name assertions, - supports Python 3.10 checksum calculation, and adds legacy/metadata-aware - attention dispatch regressions. - -Each result retains its measured source. Earlier quality, numerical and -performance campaigns are not relabeled as tests of a newer revision. +With an assigned NPU visible to the test process: ~~~bash -pre-commit run --all-files --verbose -export VLLM_RLT_OURO_CHECKPOINT=/path/to/Ouro-1.4B -export VLLM_RLT_ASCEND_RESULTS=/path/to/results -python -m pytest tests/test_ascend.py tests/test_ascend_checkpoint.py \ - tests/test_ascend_serving.py tests/test_ascend_kv.py \ - tests/test_ascend_eager.py --run-npu -q +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 ~~~ -### Portability/interface checks at `8bb68fb` - -- Ruff/pre-commit passed; CPU **928 passed / 602 skipped**. -- Assigned-device suite: **22 passed**, comprising 18 NPU and 4 CPU cases. - Both explicit backends pass spawned HTTP/SSE, termination/usage, disconnect - recovery and shutdown checks. -- Legacy kernels retain their five tensor arguments; metadata-aware kernels - receive the validated descriptor, and stale descriptors remain rejected. -- The actual fixture checksum block was also executed successfully under - Python 3.10.12, where `hashlib.file_digest` is unavailable. -- Arithmetic kernels are unchanged by this revision. The full 87-question - campaign and whole-model HF layer-event comparison were not repeated here. - -### Earlier numerical and quality evidence - -- At `be037ed`, eager versus unmodified official HF has 216 selected states - and 2016 layer events exactly equal. Restored Torch versus original Torch - is also exact. -- The analytical rounding regression distinguishes output 16 in BF16 eager - from 15.9375 in the original FP32-intermediate path, on CPU and NPU. -- At `4522008`, 87 selected same-device HF/native pairs score 63/87 each, - with no losses or gains and 87/87 complete token sequences exactly equal. - Parse failures are 19 per arm; length and execution limits are zero. - Each pair shares a physical 910B3, pooled across three cards. These are - reused regression cases, not a held-out or full-GSM8K estimate. -- Quality uses one full-prefill chunk and token budget of 4096. Chunked, - concurrent/adaptive task quality and other models remain unqualified. - -Numerical bounds remain joint relative L2 <= 2% and atol=rtol=0.02. -The explicit BF16 backend targets released HF eager. At `be037ed`, it still -fails 209/216 states against the old FP32-intermediate Torch reference; -that historical failure remains recorded and is not relabeled as passed. - -The earlier HF 63/87 versus native 61/87 observation combined different -arithmetic with native default 128-token chunking. Changing either alone did -not fix question 203. Matching both restores questions 203/881/526 and the -87-question regression. The manual question 203 replay has 15 exact logits -and 192 exact populated KV hashes. - -## Measured performance - -The primary comparison uses original Torch at `4734e19` and the explicit -backend at `4522008`; the later confirmation uses `be037ed`. - -Same physical910B3, profiler disabled, P64/C1 orC8, fixed4 loops, ignore EOS, -exactly 8 outputs/request,768 blocks,max 8 sequences / 256 batched tokens,explicit -128-token chunks. One warmup/shape,3repetitions/run; A/C/C/A gives 6 observations -per arm/shape. TTFT/TPOT are medians of per-burst request means, not pooled -individual-request percentiles. Throughput includes prefill/drain; no HTTP. - -|Concurrency|Old Torch tok/s|New backend tok/s|Old/new TTFT ms|Old/new TPOT ms| -|---:|---:|---:|---:|---:| -|1|1.326|4.711|4266.35 /221.03|246.81 /210.99| -|8|1.720|12.778|24363.00 /702.72|1818.05 /583.61| - -Pair-specific throughput ratios are 3.47–3.81x at C1 and 7.11–8.05x at C8, -with non-overlapping observed ranges. This combined backend/arithmetic change -does not isolate ScatterND. A KV-only A/B regressed 1.23% and 6.28%; -the final runtime therefore keeps original writes on the Torch path. -New backend native-writing optimality is not established. - -Routing-revision be037ed confirmation on the sameNPU gives medians 4.752tok/s -at C1 and13.705tok/s at C8 (three reps each, measured after ABBA). -Allocated peaks are about 4.97–4.99GiB; reserved peaks about 5.33GiB. -All bursts finish with 0 live logical KV blocks and no failed repetitions. - -A separate 9-shape candidate characterization covers 16/64/256 input tokens and -concurrency 1/2/8. These are short bursts, not saturation/HTTP/SLO or hardware-only -claims. CPU/NUMA placement is not pinned; independent CPU/device checks overlap -part of the shared-host timing. Raw ranges and controls must accompany reuse. - -Separate ops-only/full warmed decode traces confirm actual native ScatterND -dispatch and 96 softmax calls. Full mode has readable shapes, stacks and memory -exports. Complete warmed request traces additionally cover prefill and all 8 -outputs. Profile durations are diagnostic and excluded from speed claims. - -## Fused operators remain experimental - -BF16 FIA and promoted FP32 FlashAttentionScore did not pass earlier whole-model -gates. Native RMSNorm changes also failed selected states despite local speed -gains; SwiGLU/RotaryMul candidates were not adopted. -The current backend uses native primitives and batching, not fused attention. - -This is a bounded contribution toward -[#69](https://github.com/InferMatrixLab/vllm-rlt/issues/69), with task quality -qualified only for the stated regression configuration. +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.