diff --git a/README.md b/README.md index 00ed4e17..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 | [Runtime](docs/runtime.md), [accuracy](docs/accuracy.md) | +| Ouro-1.4B | Native Ouro | Ascend inference support | [Runtime](docs/runtime.md), [Ascend recipe](docs/recipes/ouro_ascend.md) | | Ouro-2.6B | Shared Ouro | Shares the Ouro implementation; not separately qualified, and 1.4B quality/performance results are not inherited | [Scope notes](docs/recipes/ouro_thinking.md#6-evidence-and-limitations) | | Ouro-1.4B-Thinking | Reuses Ouro | Bounded BF16, fixed-four-loop synchronous basic integration | [Shared recipe](docs/recipes/ouro_thinking.md) | | Ouro-2.6B-Thinking | Reuses Ouro | Bounded BF16, fixed-four-loop synchronous basic integration | [Shared recipe](docs/recipes/ouro_thinking.md) | diff --git a/docs/recipes/ouro_ascend.md b/docs/recipes/ouro_ascend.md new file mode 100644 index 00000000..f5e280e7 --- /dev/null +++ b/docs/recipes/ouro_ascend.md @@ -0,0 +1,71 @@ +# Ouro on Ascend + +Run Ouro on one Ascend NPU with BF16, synchronous greedy decoding and four +recurrent loops. Use the explicit `ascend_eager` attention backend. + +## 1. Prepare the environment and model + +Use a compatible Torch, Torch-NPU and CANN environment with the project's +Python dependencies installed. Keep the vendor Torch pair; CUDA/Triton +packages are not needed for this backend. + +~~~bash +source /path/to/cann/set_env.sh +source /path/to/npu-environment/bin/activate +python -m pip install --no-deps --no-build-isolation -e . +export MODEL=/path/to/prepared/Ouro-1.4B +~~~ + +Use the checkpoint's own config and tokenizer. Model loading uses the shared +native Ouro loader, including its supported sharded checkpoint layout. + +## 2. Start the server + +Select an allocated physical NPU using `ASCEND_RT_VISIBLE_DEVICES`. + +~~~bash +ASCEND_RT_VISIBLE_DEVICES= \ +python -m vllm_rlt.entrypoints.serve \ + --model "$MODEL" --tokenizer "$MODEL" --served-model-name ouro \ + --device npu --dtype bfloat16 --attention-backend ascend_eager \ + --num-blocks 1024 --max-num-seqs 1 \ + --max-num-batched-tokens 4096 --prefill-chunk-size 4096 \ + --host 127.0.0.1 --port 8123 +~~~ + +Explicit KV capacity is required. Adjust it for the model, available memory, +prompt/output lengths and concurrency. To compare with full-prefill inference, +both the chunk size and batched-token budget must fit the whole prompt. + +## 3. Send a request + +~~~bash +curl --fail http://127.0.0.1:8123/health +curl --fail http://127.0.0.1:8123/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"ouro","prompt":"The capital of France is","max_tokens":8,"temperature":0,"min_loops":4,"max_loops":4,"exit_threshold":1}' +~~~ + +Wait for HTTP 200 before sending requests. Add `"stream": true` for SSE. +The `ascend_eager` backend follows HF-style BF16 attention rounding; +`--attention-backend torch` selects the existing FP32-intermediate reference. + +## 4. Run checks + +With an assigned NPU visible to the test process: + +~~~bash +pre-commit run --all-files +VLLM_RLT_OURO_CHECKPOINT="$MODEL" python -m pytest \ + tests/test_ascend.py tests/test_ascend_checkpoint.py \ + tests/test_ascend_eager.py tests/test_ascend_kv.py \ + tests/test_ascend_serving.py --run-npu -q +~~~ + +These checks cover causal attention, cross-page KV, request reuse and HTTP/SSE. +They do not establish model quality or performance; use the +[accuracy guide](../accuracy.md) and [serving benchmark guide](../serving.md) +for those evaluations. + +Async scheduling, graphs, speculation, PD, preemption, prefix caching and +adaptive exit are currently unsupported on this path. 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..fc4209d7 --- /dev/null +++ b/tests/test_ascend.py @@ -0,0 +1,170 @@ +"""Opt-in synchronous BF16 Ouro checks using explicit Torch attention on Ascend.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_config +from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.device import resolve_device +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.kernels.paged_attention import torch_paged_attention +from vllm_rlt.models import OuroForCausalLM + +pytestmark = pytest.mark.npu + + +@pytest.fixture +def npu_device(): + resolve_device("npu") + torch.npu.set_device(0) + yield torch.device("npu:0") + torch.npu.synchronize() + torch.npu.empty_cache() + + +def assert_bf16_close(actual, expected): + a, b = actual.detach().float().cpu(), expected.detach().float().cpu() + assert torch.isfinite(a).all() and torch.isfinite(b).all() + assert (a - b).norm() / b.norm().clamp_min(1e-8) <= 0.02 + torch.testing.assert_close(a, b, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("kv_heads", [1, 2, 4]) +def test_torch_attention_on_npu_matches_cpu_prefix_oracle(npu_device, kv_heads): + torch.manual_seed(83) + keys = torch.randn(8, 2, 16, kv_heads, 64, dtype=torch.bfloat16) + values = torch.randn_like(keys) + q = torch.randn(64, 4, 64, dtype=torch.bfloat16) + lengths = torch.tensor([0, 1, 15, 16, 17, 31, 32, 3] * 8, dtype=torch.int32) + tables = torch.tensor([[5, 2], [1, 7]] * 32, dtype=torch.int32) + expected = torch_paged_attention(q, keys[:, 1], values[:, 1], tables, lengths) + actual = torch_paged_attention( + q.to(npu_device), + keys.to(npu_device)[:, 1], + values.to(npu_device)[:, 1], + tables.to(npu_device), + lengths.to(npu_device), + ) + assert_bf16_close(actual, expected) + assert torch.count_nonzero(actual[0]).item() == 0 + + +def make_npu_model(device): + torch.manual_seed(41) + config = tiny_ouro_config(hidden_size=256, intermediate_size=512, head_dim=64) + return OuroForCausalLM(config).to(device=device, dtype=torch.bfloat16).eval() + + +@torch.inference_mode() +def recurrent_states(model, chunks, *, tokens=None, incremental=False): + c = model.config + cache = KVCacheManager( + c.num_hidden_layers, + c.num_key_value_heads, + c.head_dim, + 32, + 16, + 4, + device=next(model.parameters()).device, + dtype=torch.bfloat16, + backend="torch", + incremental_allocation=incremental, + ) + tokens = tokens if tokens is not None else ([3, 7, 2, 11, 5, 17, 8] * 3)[:21] + assert cache.allocate("a", len(tokens), initial_tokens=16 if incremental else None) + states, kv, start = {}, {}, 0 + try: + for size in chunks: + assert cache.ensure_capacity("a", start + size) + positions = list(range(start, start + size)) + hidden = model.prelude(torch.tensor(tokens[start : start + size], device=cache.device)) + for depth in range(4): + hidden, _ = model.recurrent(hidden, ["a"] * size, [depth] * size, positions, cache) + states[(start, depth)] = (hidden.clone(), model.coda(hidden).clone()) + start += size + assert start == len(tokens) + for depth in range(4): + for layer in range(c.num_hidden_layers): + assert all(cache.token_written("a", p, depth) for p in range(len(tokens))) + kv[(layer, depth)] = tuple( + t.clone() for t in cache.read(layer, "a", depth, len(tokens)) + ) + assert len({tuple(cache.get_block_table("a", d)) for d in range(4)}) == 4 + return states, kv + finally: + cache.free("a") + assert cache.num_used_blocks == 0 + + +def test_npu_chunked_prefill_and_incremental_kv(npu_device): + model = make_npu_model(npu_device) + expected, expected_kv = recurrent_states(model, [19, 1, 1]) + incremental, incremental_kv = recurrent_states(model, [19, 1, 1], incremental=True) + for key in expected: + for a, b in zip(incremental[key], expected[key], strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0) + for key in expected_kv: + for a, b in zip(incremental_kv[key], expected_kv[key], strict=True): + torch.testing.assert_close(a, b, atol=0, rtol=0) + chunked, chunked_kv = recurrent_states(model, [5, 14, 1, 1], incremental=True) + for key in ((19, d) for d in range(4)): + for a, b in zip(chunked[key], expected[key], strict=True): + assert_bf16_close(a, b) + for key in expected_kv: + for a, b in zip(chunked_kv[key], expected_kv[key], strict=True): + assert_bf16_close(a, b) + + +def test_npu_cancel_reuse_and_multiple_requests(npu_device): + engine = StageEngine( + make_npu_model(npu_device), + attention_backend="torch", + cache_config=CacheConfig(num_blocks=32, block_size=16, incremental_allocation=True), + scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=32), + ) + try: + params = SamplingParams(max_tokens=3, temperature=0, ignore_eos=True) + engine.add_request("cancel", [3, 7, 2] * 6, params) + engine.add_request("keep", [5, 9, 11], params) + engine.step() + engine.abort_request("cancel") + engine.add_request("cancel", [13, 14, 15], params) + finished = {} + for _ in range(200): + for output in engine.step(): + if output.finished: + finished[output.request_id] = output + if not engine.has_unfinished_requests(): + break + assert set(finished) == {"keep", "cancel"} + assert all( + o.exit_depths == [4] * 3 and o.finish_reason == "length" for o in finished.values() + ) + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() + + +@pytest.mark.parametrize("ignore_eos", [False, True]) +def test_npu_eos_and_length_handling(npu_device, ignore_eos): + model = make_npu_model(npu_device) + model.lm_head.weight.zero_() # Argmax 0 is the tiny model's actual EOS token. + engine = StageEngine(model, attention_backend="torch", cache_config=CacheConfig(num_blocks=16)) + try: + engine.add_request( + "eos", [3, 7, 2], SamplingParams(max_tokens=3, temperature=0, ignore_eos=ignore_eos) + ) + output = None + for _ in range(100): + for item in engine.step(): + if item.finished: + output = item + if not engine.has_unfinished_requests(): + break + assert output is not None + assert output.token_ids == [0] * (3 if ignore_eos else 1) + assert output.finish_reason == ("length" if ignore_eos else "stop") + assert engine.cache_manager.num_used_blocks == 0 + finally: + engine.close() diff --git a/tests/test_ascend_checkpoint.py b/tests/test_ascend_checkpoint.py new file mode 100644 index 00000000..c26e8fe8 --- /dev/null +++ b/tests/test_ascend_checkpoint.py @@ -0,0 +1,100 @@ +"""Opt-in local-checkpoint KV and lifecycle regressions on Ascend.""" + +import os +from pathlib import Path + +import pytest +import torch + +from tests.test_ascend import assert_bf16_close, recurrent_states +from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig +from vllm_rlt.device import resolve_device +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import OuroForCausalLM, load_tokenizer + +pytestmark = pytest.mark.npu + + +@pytest.fixture(scope="module") +def checkpoint(): + source = os.environ.get("VLLM_RLT_OURO_CHECKPOINT") + if not source: + pytest.skip("set VLLM_RLT_OURO_CHECKPOINT to a prepared local Ouro checkpoint") + if not Path(source).is_dir(): + pytest.fail("VLLM_RLT_OURO_CHECKPOINT must be a prepared local directory") + resolve_device("npu") + torch.npu.set_device(0) + model = OuroForCausalLM.from_pretrained(source, device="npu", dtype=torch.bfloat16) + assert all(p.dtype == torch.bfloat16 for p in model.parameters()) + yield model, load_tokenizer(source) + torch.npu.synchronize() + + +@torch.inference_mode() +def test_checkpoint_incremental_pages_preserve_hidden_logits_and_kv(checkpoint): + model, tokenizer = checkpoint + tokens = tokenizer.encode( + "The capital of France is Paris. The capital of Japan is Tokyo. " + "A language model predicts the next token in a sequence.", + add_special_tokens=False, + )[:21] + assert len(tokens) == 21 + expected, expected_kv = recurrent_states(model, [19, 1, 1], tokens=tokens) + actual, actual_kv = recurrent_states(model, [19, 1, 1], tokens=tokens, incremental=True) + pairs = {} + for key in expected: + for name, a, b in zip(("hidden", "logits"), actual[key], expected[key], strict=True): + pairs[f"{key}.{name}"] = (a, b) + for key in expected_kv: + for name, a, b in zip(("key", "value"), actual_kv[key], expected_kv[key], strict=True): + pairs[f"{key}.{name}"] = (a, b) + for a, b in pairs.values(): + assert_bf16_close(a, b) + torch.testing.assert_close(a, b, atol=0, rtol=0) + + +@pytest.mark.parametrize("backend", ["torch", "ascend_eager"]) +@torch.inference_mode() +def test_checkpoint_greedy_replay_and_request_cleanup(checkpoint, backend): + model, tokenizer = checkpoint + engine = StageEngine( + model, + attention_backend=backend, + cache_config=CacheConfig(num_blocks=64), + scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=32), + ) + prompts = ["The capital of France is", "One plus one equals"] + loops = model.config.total_ut_steps + params = SamplingParams( + max_tokens=4, temperature=0, min_loops=loops, max_loops=loops, ignore_eos=True + ) + expected = None + try: + # Repeat the same packed workload and reuse IDs after all pages are freed. + # Expected IDs come from this checkpoint, not a particular weight file. + for _ in range(2): + finished = {} + for index, prompt in enumerate(prompts): + engine.add_request( + str(index), tokenizer.encode(prompt, add_special_tokens=False), params + ) + for _ in range(200): + for output in engine.step(): + if output.finished: + assert len(output.token_ids) == params.max_tokens + assert all( + 0 <= token < model.config.vocab_size for token in output.token_ids + ) + assert output.exit_depths == [loops] * params.max_tokens + assert output.finish_reason == "length" + finished[output.request_id] = list(output.token_ids) + if not engine.has_unfinished_requests(): + break + assert set(finished) == {str(index) for index in range(len(prompts))} + assert engine.cache_manager.num_used_blocks == 0 + if expected is None: + expected = finished + else: + assert finished == expected + finally: + engine.close() diff --git a/tests/test_ascend_eager.py b/tests/test_ascend_eager.py new file mode 100644 index 00000000..37e5f6ce --- /dev/null +++ b/tests/test_ascend_eager.py @@ -0,0 +1,154 @@ +"""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 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) + 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() + + +@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) diff --git a/tests/test_ascend_kv.py b/tests/test_ascend_kv.py new file mode 100644 index 00000000..298d0c1a --- /dev/null +++ b/tests/test_ascend_kv.py @@ -0,0 +1,65 @@ +"""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]) +@pytest.mark.parametrize("backend", ["torch", "ascend_eager"]) +@torch.inference_mode() +def test_native_scatter_preserves_inputs_other_layers_and_reused_ids(rows, backend, monkeypatch): + device = resolve_device("npu:0") + torch.npu.set_device(device) + 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) + 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() + assert len(calls) == (6 if backend == "ascend_eager" else 0) diff --git a/tests/test_ascend_serving.py b/tests/test_ascend_serving.py new file mode 100644 index 00000000..384eef37 --- /dev/null +++ b/tests/test_ascend_serving.py @@ -0,0 +1,136 @@ +"""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 + + +@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)) / backend + 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", + backend, + "--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_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/tests/test_device.py b/tests/test_device.py new file mode 100644 index 00000000..02672ce0 --- /dev/null +++ b/tests/test_device.py @@ -0,0 +1,67 @@ +"""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 +from vllm_rlt.profiling import ProfileConfig + + +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(), "torch") + blocks, info = plan_cache( + 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/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..e04d9bac 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -36,6 +36,17 @@ class AttentionBackend: 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) + return self(q, keys, values, batch.block_tables, batch.context_lengths) + def prefill(self, *args, **kwargs): return self.kernel.prefill(*args, **kwargs) @@ -68,11 +79,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 dfcdf0df..67b21e2e 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 @@ -683,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" and self.backend == "ascend_eager": + 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) @@ -740,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/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..e1b25f5c 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, @@ -70,6 +72,30 @@ def __init__( self.execution_config = resolve_cuda_graph_rows( execution_config or ExecutionConfig(), scheduler_config, speculative_config ) + if parameter.device.type == "npu": + 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: + 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 +349,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/entrypoints/cli.py b/vllm_rlt/entrypoints/cli.py index a9b8025c..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,6 +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 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 cd452d92..70a708f7 100644 --- a/vllm_rlt/entrypoints/serve.py +++ b/vllm_rlt/entrypoints/serve.py @@ -82,12 +82,13 @@ 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", type=str.lower, choices=[ + "ascend_eager", "auto", "torch", "triton", @@ -115,6 +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 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) 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) 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/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(): 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,