Skip to content

[Feature] Add experimental Ascend support for Ouro - #168

Open
KrystalRay wants to merge 12 commits into
InferMatrixLab:mainfrom
KrystalRay:ouro-ascend-contribution
Open

KrystalRay wants to merge 12 commits into
InferMatrixLab:mainfrom
KrystalRay:ouro-ascend-contribution

Conversation

@KrystalRay

Copy link
Copy Markdown

Required human contributor answers

1. What outcome are you trying to achieve?

Enable Ouro to run through vllm-rlt’s native runtime and OpenAI-compatible
server on a single Ascend NPU. The initial scope is BF16, fixed four-loop,
synchronous greedy inference with full prefill.

2. What is the highest-leverage variable for that outcome?

End-to-end correctness of the supported NPU execution path, particularly
attention arithmetic and the causality, isolation and lifecycle of
per-depth KV state.

3. How will this work move that variable rather than create busywork?

This work adds the device/backend integration needed to execute the existing
runtime on Ascend, together with targeted attention, KV and serving regressions.

Paired real-weight checks against official HF and actual HTTP measurements
make the supported behavior and remaining gaps reviewable. Failed checks,
including the chunk128 sequence mismatch, remain documented without relaxing
the acceptance threshold.

4. Who is the customer or user, and what do they actually need?

Ouro users with Ascend hardware need a reproducible native deployment path,
documented operating limits, correct KV behavior, and measured latency and
memory costs.

Maintainers need an optional integration with clear boundaries and repeatable
regression tests so they can assess and maintain this hardware path.

PR necessity and bugfix evidence

Yes. A PR is needed to add NPU device dispatch, attention-backend integration
and paged-KV support to the library and serving entrypoints.

This PR is primarily experimental hardware enablement. The pre-change serving
entrypoint rejected --device npu; the new path has been exercised on real
Ascend hardware, with validation results recorded in Test Plan and Test Result.

The proposed qualified scope remains full prefill. The chunk128 sequence
mismatch is an unresolved finding, and this PR does not claim to resolve it
or certify broader model quality.

Purpose

Add single-device Ascend support for native Ouro with the explicit
ascend_eager backend. The initial recipe uses BF16, four recurrent loops,
synchronous greedy decoding and full prefill.

The implementation adds:

  • Lazy Torch-NPU initialization and NPU worker dispatch.
  • Attention that follows the official HF eager rounding order, plus ScatterND
    writes to the existing paged KV cache.
  • Synchronization before request metadata or KV pages are reused.

Setup and serving commands are in the
Ascend recipe.

Test Plan

vllm-rlt Version / Commit: d9217196a40a754f4e7d967f47dac79765236f36

Tested on Ascend 910B3 with Ouro-1.4B.

pre-commit run --all-files

VLLM_RLT_OURO_CHECKPOINT=/path/to/Ouro-1.4B 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

Compare the native engine against official HF on the same device. Check complete
generated sequences, causal attention, KV page boundaries, cancellation, request
reuse and HTTP/SSE serving. Measure latency and throughput without the profiler.

Test Result

  • Pre-commit passed; CPU suite: 928 passed, 603 skipped.
  • Ascend test selection: 23 passed.
  • Three same-device HF regression cases matched complete generated sequences.
  • HTTP/SSE runs completed with no request errors.

Single-request engine measurements, generating 64 tokens after warmup.
Throughput is the median of three runs; TTFT and TPOT are p50.

Input tokens Output tok/s TTFT (ms) TPOT (ms)
64 5.12 206 195
512 5.01 280 198
2048 4.82 1011 195

@KrystalRay

Copy link
Copy Markdown
Author

Validation details and follow-up work

At commit d921719, cases 203, 881 and 526 were checked against official HF
on the same Ascend device. All three matched the complete generated token
sequences with full prefill.

Case 203 also exposed a chunking difference. Full prefill produces the final
numerical answer 31, which is correct; 128-token chunks produce 27.
The two generated token sequences first differ at the 15th output token.
Chunked-prefill stability is follow-up work.

For case 203, native ascend_eager and official HF eager were compared with
matching chunk sizes. Their final-loop output logits matched exactly, element
by element across the full 49,152-token vocabulary: 162 generation steps for
full prefill and 163 for chunk128, with a maximum absolute difference of zero
in both comparisons. The chunked HF control used an explicit four-dimensional
causal mask with unmodified model source.

The earliest observed full/chunk difference was in the first-layer Q/K/V
projections. Replaying Linear with frozen inputs and weights reproduced those
projection differences by changing only the row count.

CUDA SM120 showed the opposite answer change on the same question: official HF
produced the final numerical answer 27 with full prefill and 31 with chunk128.
The CUDA runs used Torch 2.11; the Ascend runs used Torch 2.9.

An earlier 87-question run at commit 4522008 matched all 87 complete sequences
against HF, with both implementations scoring 63/87. That campaign has not
been rerun at d921719.

Hardware testing used 910B3. Other Ascend variants still need device validation.
Adaptive exit, async scheduling, graphs, speculation, PD, preemption and prefix
caching are outside this first contribution. Fused CANN attention is separate
work.

The performance table uses a fixed 13.5 GiB KV pool, one excluded warmup and
three measured runs per shape. The machine was shared; CPU affinity and thread
settings were recorded. Latency ranges, memory use and concurrency results are
retained with the raw measurements.

HTTP/SSE measurements covered 24 requests at concurrency 1 and 48 at
concurrency 8, with zero errors.

Earlier A/B numbers compare the entire backends, whose attention arithmetic
differs. ScatterND-only trials did not improve end-to-end performance.
Profiler captures were collected separately from timing; one full capture
reported 16 incomplete memory records.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant