[RFC] Add Ascend NPU Support for DeepSpec

---

## 1. Summary

This RFC proposes extending DeepSpec to run the full data-preparation → training → evaluation pipeline on Huawei Ascend NPU (910B2 and later) in addition to the currently supported NVIDIA GPUs. The change is fully backward-compatible: every device-selection site is guarded by a runtime `try/except torch_npu` check that falls back to the existing CUDA path, so GPU users see no behavioral or performance change. No public API is altered; the only user-facing additions are one optional CLI flag (`--tasks`) on `eval.py` and per-epoch checkpoint saving during training.

## 2. Motivation

### 2.1 Current Limitation

The README states: _"Hardware: the default configs and scripts assume a single node with 8 GPUs."_ All device-related code paths are hard-coded to `torch.cuda`: stream management, distributed backend (`nccl`), FSDP `device_mesh`, RNG state, attention implementation (`flex_attention`), and the spawn-based launcher all assume CUDA. As a result, DeepSpec cannot be used on the growing fleet of Ascend NPU clusters without source-level patches.

### 2.2 Why Ascend NPU

- **Hardware availability.** Ascend 910B/910C accelerators are widely deployed in domestic (Chinese) data centers and are increasingly the default accelerator available to many research teams.
- **Speculative-decoding research demand.** Several teams working on draft-model and speculative-decoding research have asked for a reference implementation that runs on NPU, since the target models they study (e.g. Qwen3) are first-class citizens in DeepSpec.
- **Ecosystem alignment.** `torch_npu` has reached feature parity with the operations DeepSpec needs (FSDP, `sdpa`, `hccl` collectives). The remaining gaps (`flex_attention`, float64 `all_reduce`) are narrow and well-understood.

### 2.3 Goals

- Run **target-cache generation**, **DSpark/Eagle3 training**, and **speculative-decoding evaluation** end-to-end on Ascend 910B and above with Qwen3-8B.
- **Zero regression** on NVIDIA GPUs: no behavior change, no performance cliff, no new required dependency for GPU users.
- Keep the implementation **single-sourced** (no `if device == "npu"` sprinkled through modeling code beyond a few well-justified sites).

## 3. Detailed Design

### 3.1 Design Principle: Runtime Device Detection

Every site that previously called a CUDA-specific API now goes through a small helper that detects the backend at runtime:

```python
def _is_npu_available() -> bool:
try:
import torch_npu # noqa: F401 — registers the NPU backend
return torch.npu.is_available()
except Exception:
return False
```

This pattern is used in `train.py`, `eval.py`, `scripts/data/prepare_target_cache.py`, `deepspec/utils/__init__.py`, `deepspec/utils/distributed.py`, `deepspec/data/cuda_prefetcher.py`, `deepspec/trainer/base_trainer.py`, `deepspec/trainer/ckpt_manager.py`, and `deepspec/eval/dspark/confidence_head.py`. When `torch_npu` is absent (the GPU case), every helper returns the original CUDA value, so GPU behavior is byte-for-byte identical.

### 3.2 Distributed Initialization

`deepspec/utils/distributed.init_dist` now picks the backend and device module based on the detected hardware:

| Platform | Backend | Device module | Process-group init |
|---|---|---|---|
| NPU | `hccl` | `torch.npu` | `init_process_group("hccl", device_id=torch.npu.current_device())` |
| CUDA | `nccl` | `torch.cuda` | `init_process_group("nccl")` (unchanged) |

The `device_id` kwarg is required by `hccl` to bind the process to its NPU; it is harmless (and ignored) on the CUDA path because that branch is never taken when NPU is absent.

### 3.3 Process Spawn: Honoring Visibility Masks

A subtle bug was uncovered during NPU bring-up: `torch.npu.device_count()` returns the **total physical NPU count** and does **not** honor `ASCEND_RT_VISIBLE_DEVICES`. The original `torch.cuda.device_count()` *does* honor `CUDA_VISIBLE_DEVICES`, so the old code happened to work on GPU. On NPU, spawning `device_count()` workers when only a subset of devices is visible causes the extra workers to fail to acquire a device and deadlock inside `hccl`'s barrier.

A new helper, `_visible_device_count()`, is introduced in `train.py`, `eval.py`, and `prepare_target_cache.py`:

```python
def _visible_device_count() -> int:
if _is_npu_available():
vis = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "").strip()
if vis:
return len([d for d in vis.split(",") if d.strip() != ""])
return torch.npu.device_count()
vis = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
if vis:
return len([d for d in vis.split(",") if d.strip() != ""])
return torch.cuda.device_count()
```

This makes the GPU path *more* robust (it now explicitly reads `CUDA_VISIBLE_DEVICES` rather than relying on `device_count()`'s implicit honoring) and fixes the NPU deadlock. The return value is identical to the old behavior on GPU when `CUDA_VISIBLE_DEVICES` is unset.

### 3.4 Stream / Prefetcher

`deepspec/data/cuda_prefetcher.py` hard-coded `torch.cuda.Stream`. It now selects the stream module dynamically:

```python
def _stream_module():
return torch.npu if _is_npu_available() else torch.cuda
```

`Stream`, `current_stream`, and `synchronize` all resolve through this helper. No API change.

### 3.5 RNG Seeding

`deepspec/utils/__init__.py.seed_all` now seeds NPU generators when on NPU and CUDA generators otherwise. The NPU branch is wrapped in `try/except` so GPU-only environments without `torch_npu` are unaffected.

### 3.6 FSDP Device Mesh

`deepspec/trainer/base_trainer.py` constructed the `DeviceMesh` with a hard-coded `"cuda"` `device_type`. It now selects `"npu"` or `"cuda"` and queries the per-node device count from the corresponding backend. The `sharding_strategy` logic and FSDP kwargs are otherwise unchanged.

### 3.7 Checkpoint Manager: RNG State

`deepspec/trainer/ckpt_manager.py` saved and restored `torch.cuda.get_rng_state()` only. The new implementation:

- **Save:** records RNG state from whichever backend is active (NPU or CUDA), stored under a unified key `"torch_cuda_rng"` (kept for backward-compatible reading) plus the legacy per-rank state.
- **Load:** uses `checkpoint.get("torch_cuda_rng")` so older checkpoints (which lack the key) load without error; when the key is `None`, the RNG state is simply not restored.
- **Per-epoch saving (new capability, applies to both backends):** a checkpoint is now also written at the end of each epoch (`step_` plus an `epoch_` alias), in addition to the existing step-based cadence. This is useful on any backend and is gated by the existing logging config.

### 3.8 Attention: Eager Path for NPU

This is the only modeling-level divergence. `torch.nn.attention.flex_attention` is **CUDA/CPU/HPU-only**; on NPU it raises `FlexAttention is only supported on CUDA, CPU or HPU devices`.

**DSpark mask construction** (`deepspec/modeling/dspark/common.py::create_dspark_attention_mask`): a new `use_block_mask: bool = True` parameter selects the output type:

- `use_block_mask=True` (default, GPU): returns a `BlockMask` via `create_block_mask` — the original fast path.
- `use_block_mask=False` (NPU): returns a 4-D additive float mask of shape `[bsz, 1, q_len, kv_len]` with `-inf` at masked positions, consumed by the standard `eager` / `sdpa` attention kernels.

**Config selection** (`deepspec/modeling/dspark/qwen3/config.py`):

```python
TRAIN_ATTN_IMPLEMENTATION = "eager" if _is_npu_available() else "flex_attention"
```

The constant remains `flex_attention` on GPU (matching every other model config in the repo) and switches to `eager` only on NPU. `EVAL_ATTN_IMPLEMENTATION` stays `"sdpa"` for both backends, since NPU supports `scaled_dot_product_attention`.

**Modeling wiring** (`deepspec/modeling/dspark/qwen3/modeling.py`): the call site passes `use_block_mask=(config._attn_implementation == "flex_attention")`, so the mask type always matches the attention kernel that will consume it. This is a single-line change and is a no-op on GPU.

### 3.9 Evaluation: float64 `all_reduce`

`deepspec/eval/dspark/confidence_head.py` accumulates calibration histograms in `float64` for numerical fidelity and then `dist.all_reduce`s them. The Ascend `hccl` backend does **not** support `float64` `all_reduce` (runtime error `ERR02007 DIST feature not supported`); `nccl` on GPU does.

The fix is backend-scoped:

```python
need_f32_cast = self._on_npu()
for tensor in (...):
if need_f32_cast and tensor.dtype is torch.float64:
reduced = tensor.to(torch.float32)
dist.all_reduce(reduced, op=dist.ReduceOp.SUM)
tensor.copy_(reduced.to(torch.float64))
else:
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
```

On GPU, `need_f32_cast` is `False` and the original `float64` reduction runs unchanged — no precision or behavior difference. The `float32` round-trip on NPU introduces negligible error for histogram-based AUROC/ECE/Brier metrics (the bin counts are integers and the probabilities are already `float32` before accumulation).

### 3.10 Evaluation CLI: `--tasks` Flag

`eval.py` previously ran all nine benchmark datasets unconditionally. A new optional `--tasks` flag accepts a comma-separated list (e.g. `--tasks gsm8k`) and selects the matching subset. When omitted, behavior is unchanged (all tasks run). This is generally useful for quick smoke tests on any backend and is not NPU-specific.

## 4. Backward Compatibility

| Concern | Impact on GPU users |
|---|---|
| New `try/except torch_npu` imports | `torch_npu` is absent on GPU machines; the `except` branch returns the original CUDA value. No new dependency. |
| `_visible_device_count()` | Reads `CUDA_VISIBLE_DEVICES` explicitly; returns the same value as the old `torch.cuda.device_count()` when the env var is set, and `torch.cuda.device_count()` when unset. |
| Attention implementation | Stays `flex_attention` on GPU; `use_block_mask` defaults to `True`. |
| `all_reduce` dtype | Stays `float64` on GPU (nccl). |
| Checkpoint format | New RNG key is additive; old checkpoints load via `.get()`. |
| CLI | `--tasks` is optional; default = all tasks. |
| Public API | None changed. |

No `requirements.txt` change is required: `torch_npu` is an optional, separately-installed package that users on NPU already have.

## 5. Drawbacks

- **Backend detection is repeated.** The `_is_npu_available()` helper is duplicated across entry-point files rather than centralized. Centralizing it in `deepspec.utils` would be cleaner, but each entry point currently has its own minimal `_device_count`-style helper and we chose to match that local style to keep the diff small. A follow-up could consolidate these.
- **Two attention code paths to maintain.** The eager 4-D mask path is now exercised on NPU while GPU uses `flex_attention`. Bugs in one path may not surface on the other. Mitigation: the mask *semantics* are identical (same `mask_mod`), only the materialization differs; and `sdpa`/`eager` are the universally-available fallback that GPU users can also opt into by setting `TRAIN_ATTN_IMPLEMENTATION="eager"`.
- **No NPU CI.** This RFC is validated on physical 910B2 hardware only; without NPU CI, regressions on NPU could land undetected. A periodic manual smoke test is the interim mitigation.

## 6. Validation

The implementation branch was validated end-to-end on a single node with 8× Ascend 910B2 (65 GB HBM each), using Qwen3-8B as the target model:

1. **Target-cache generation** — 10,000 samples from `qwen235_metamath395k`, produced a 217 GB cache; index integrity verified (4,729,878 tokens, mean seq_len 473, size matches theoretical estimate).
2. **Training** — DSpark draft model, 3 epochs, 19 steps/epoch, `num_anchors=256`, bf16, `no_shard`. Loss decreased 5.55 → 2.52. Three epoch checkpoints saved and reloadable.
3. **Evaluation** — gsm8k (10 samples, 4-NPU), speculative decoding with temperature 1.0. Results: `accept_len = 1.03`, `verify_rate = 0.1289`, `accept_rate@0 = 0.0296`, confidence-head AUC = 0.7422. Acceptable for a 10k-sample / 3-epoch smoke test; not a quality claim.

GPU regression: not run in this session (no GPU node available), but the code paths are unchanged when `torch_npu` is absent — see Section 4.

## 7. Unresolved Questions

1. **`num_anchors` OOM workaround.** Training Qwen3-8B DSpark on 65 GB NPU required reducing `num_anchors` from 512 → 256 to avoid OOM during the `aligned_target_hidden` gather. Should this be auto-tuned based on free memory, or left as a config override? (This is orthogonal to NPU support — the same OOM would occur on a 40 GB GPU.)
2. **Multi-node NPU.** Only single-node was tested. `hccl` supports multi-node; is there demand to validate it now, or defer?
3. **Centralization.** Should `_is_npu_available()` / `_visible_device_count()` be promoted into `deepspec.utils` as the single source of truth, and the entry-point copies removed?

## 8. Future Work

- Consolidate device helpers into `deepspec.utils.hardware` (per Section 7.3).
- Add an NPU-gated test suite (`tests/npu/`) covering init, FSDP step, mask equivalence, and a 1-sample eval round-trip.
- Performance: profile and tune NPU kernels (e.g. replace eager attention with a fused NPU SDPA variant where available).
- Extend the same backend-detection pattern to DFlash and Eagle3 modeling if they turn out to need it (currently only DSpark/Qwen3 was ported).
- Document NPU setup in `README.md` (env vars, `torch_npu` install, `ASCEND_RT_VISIBLE_DEVICES` semantics) once this RFC is accepted.

## 9. Rollout

1. Merge `npu-support` into `main` behind the existing runtime detection (no feature flag needed — NPU code is inert on GPU).
2. Add a short "Hardware" subsection to `README.md` noting NPU is now supported and linking to the env-var table below.
3. Tag a release note: _"Experimental Ascend NPU support (single-node). GPU users are unaffected."_

### Environment Variables (NPU)

| Variable | Purpose | Example |
|---|---|---|
| `ASCEND_RT_VISIBLE_DEVICES` | Restrict visible NPUs (analogous to `CUDA_VISIBLE_DEVICES`) | `4,5,6,7` |
| `MASTER_ADDR` / `MASTER_PORT` | Distributed init (same as GPU) | `127.0.0.1` / `29500` |
| `PYTORCH_NPU_ALLOC_CONF` | NPU memory allocator tuning (analogous to `PYTORCH_CUDA_ALLOC_CONF`) | `expandable_segments:True` |

## 10. Open Questions for Reviewers

- Is the runtime-detection pattern (Section 3.1) acceptable, or would reviewers prefer an explicit `--device npu\|cuda` CLI flag?
- Are there objections to adding `--tasks` to `eval.py` (Section 3.10) as part of this RFC, or should it be a separate PR?

## 11. mplementation:

- branch [`npu-support`](https://github.com/sunny-infra/DeepSpec-Ascend/tree/npu-support)
- This branch has only been validated on Ascend 910B hardware at present. We will conduct additional verification on Ascend 910C before submitting the code.

---