[feature] Add Ascend NPU Support While Preserving GPU Compatibility

deepspec/data/cuda_prefetcher.py

@@ -3,9 +3,21 @@
 import torch
 
 
+def _is_npu_available():
+    try:
+        import torch_npu  # noqa: F401
+        return torch.npu.is_available()
+    except Exception:
+        return False
+
+
+def _stream_module():
+    return torch.npu if _is_npu_available() else torch.cuda
+
+
 def move_batch_to_device(batch, device):
     moved = {key: value.to(device, non_blocking=True) for key, value in batch.items()}
-    # Embedding lookup requires int64; cast on GPU to avoid bloating CPU-to-GPU transfer.
+    # Embedding lookup requires int64; cast on device to avoid bloating host-to-device transfer.
     if moved["input_ids"].dtype != torch.long:
         moved["input_ids"] = moved["input_ids"].to(torch.long)
     return moved
@@ -14,15 +26,16 @@ def move_batch_to_device(batch, device):
 class CUDAPrefetcher:
     """Overlaps DataLoader iteration and H2D transfer with compute.
 
-    Uses a background thread and a dedicated CUDA stream so that both
+    Uses a background thread and a dedicated device stream so that both
     next(dataloader) and the H2D copy for the next batch run concurrently
     with the forward/backward of the current batch.
     """
 
     def __init__(self, dataloader, device):
         self.dataloader = dataloader
         self.device = device
-        self.stream = torch.cuda.Stream(device=device)
+        self._sm = _stream_module()
+        self.stream = self._sm.Stream(device=device)
 
     def __iter__(self):
         self._iter = iter(self.dataloader)
@@ -40,7 +53,7 @@ def _fetch_and_transfer(self):
         except StopIteration:
             self._done = True
             return
-        with torch.cuda.stream(self.stream):
+        with self._sm.stream(self.stream):
             self._gpu_batch = move_batch_to_device(cpu_batch, self.device)
 
     def __next__(self):
@@ -53,7 +66,7 @@ def __next__(self):
             raise StopIteration
 
         # Make the compute stream wait until the side-stream H2D is complete.
-        current = torch.cuda.current_stream(self.device)
+        current = self._sm.current_stream(self.device)
         current.wait_stream(self.stream)
         batch = self._gpu_batch
 

deepspec/eval/dspark/confidence_head.py

@@ -92,6 +92,10 @@ def update(
         self.brier_num[:pos_count].add_(weights * (probs - targets).pow(2))
 
     def all_reduce(self) -> None:
+        # NPU hccl does not support float64 all_reduce (ERR02007). On NPU,
+        # cast float64 tensors to float32 for the reduction and restore the
+        # original dtype afterwards. GPU/nccl supports float64 natively.
+        need_f32_cast = self._on_npu()
         for tensor in (
             self.coarse_count,
             self.coarse_pred,
@@ -100,7 +104,20 @@ def all_reduce(self) -> None:
             self.fine_neg,
             self.brier_num,
         ):
-            dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+            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)
+
+    @staticmethod
+    def _on_npu() -> bool:
+        try:
+            import torch_npu  # noqa: F401
+            return torch.npu.is_available()
+        except Exception:
+            return False
 
     @staticmethod
     def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float:

deepspec/modeling/dspark/common.py

@@ -82,6 +82,7 @@ def create_dspark_attention_mask(
     seq_len: int,
     block_size: int,
     device: torch.device,
+    use_block_mask: bool = True,
 ):
     def dspark_mask_mod(b, h, q_idx, kv_idx):
         del h
@@ -96,14 +97,42 @@ def dspark_mask_mod(b, h, q_idx, kv_idx):
         return (mask_context | mask_draft) & is_valid_block
 
     bsz, num_blocks = anchor_positions.shape
-    return create_block_mask(
-        dspark_mask_mod,
-        B=bsz,
-        H=None,
-        Q_LEN=num_blocks * block_size,
-        KV_LEN=seq_len + num_blocks * block_size,
-        device=device,
-    )
+    q_len = num_blocks * block_size
+    kv_len = seq_len + num_blocks * block_size
+    if use_block_mask:
+        return create_block_mask(
+            dspark_mask_mod,
+            B=bsz,
+            H=None,
+            Q_LEN=q_len,
+            KV_LEN=kv_len,
+            device=device,
+        )
+    # Eager / non-flex path: materialize a 4D additive mask (bool -> 0/-inf).
+    # attn pattern: (mask_context | mask_draft) & is_valid_block
+    q_idx = torch.arange(q_len, device=device).view(1, 1, q_len, 1)
+    kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, kv_len)
+    q_block_id = q_idx // block_size  # (1,1,q_len,1)
+    # gather anchor_pos per q_block_id: (bsz,1,q_len,1)
+    anchor_pos_q = torch.gather(
+        anchor_positions, dim=1, index=q_block_id.view(bsz, -1)
+    ).view(bsz, 1, q_len, 1)
+    is_context = kv_idx < seq_len
+    mask_context = is_context & (kv_idx < anchor_pos_q)
+    is_draft = kv_idx >= seq_len
+    kv_block_id = (kv_idx - seq_len) // block_size
+    mask_draft = is_draft & (q_block_id == kv_block_id)
+    # is_valid_block per q_block_id -> (bsz,1,q_len,1)
+    is_valid_block_q = torch.gather(
+        block_keep_mask, dim=1, index=q_block_id.view(bsz, -1)
+    ).view(bsz, 1, q_len, 1)
+    keep = (mask_context | mask_draft) & is_valid_block_q  # bool (bsz,1,q_len,kv_len)
+    # Convert to additive mask: 0 where keep, -inf elsewhere. eager_attention_forward
+    # does `attn_weights + attention_mask`, so use min dtype-safe value.
+    neg_inf = torch.finfo(torch.float32).min
+    additive = torch.zeros(bsz, 1, q_len, kv_len, dtype=torch.float32, device=device)
+    additive = additive.masked_fill(~keep, neg_inf)
+    return additive
 
 
 def build_anchor_candidate_mask(

deepspec/modeling/dspark/qwen3/config.py

@@ -1,9 +1,21 @@
 import copy
 
+import torch
+
 from deepspec.modeling.dspark.common import validate_target_layer_ids
 
 
-TRAIN_ATTN_IMPLEMENTATION = "flex_attention"
+def _is_npu_available():
+    try:
+        import torch_npu  # noqa: F401
+        return torch.npu.is_available()
+    except Exception:
+        return False
+
+
+# NPU does not support flex_attention; fall back to eager. GPU keeps the
+# original flex_attention path for best performance.
+TRAIN_ATTN_IMPLEMENTATION = "eager" if _is_npu_available() else "flex_attention"
 
 
 def build_draft_config(

deepspec/modeling/dspark/qwen3/modeling.py

@@ -419,6 +419,7 @@ def forward(
             seq_len=seq_len,
             block_size=self.block_size,
             device=device,
+            use_block_mask=self.config._attn_implementation == "flex_attention",
         )
         output_hidden = self._forward_backbone(
             position_ids=full_position_ids,

deepspec/trainer/base_trainer.py

@@ -52,6 +52,14 @@
 )
 
 
+def _is_npu_available():
+    try:
+        import torch_npu  # noqa: F401
+        return torch.npu.is_available()
+    except Exception:
+        return False
+
+
 def _build_fsdp_kwargs(
     *, sharding_strategy_name: str, precision_dtype, world_size: int
 ) -> dict:
@@ -65,9 +73,14 @@ def _build_fsdp_kwargs(
         sharding_strategy=sharding_strategy,
     )
     if sharding_strategy in _HYBRID_STRATEGIES:
-        devices_per_node = torch.cuda.device_count()
+        if _is_npu_available():
+            devices_per_node = torch.npu.device_count()
+            device_type = "npu"
+        else:
+            devices_per_node = torch.cuda.device_count()
+            device_type = "cuda"
         fsdp_kwargs["device_mesh"] = init_device_mesh(
-            "cuda",
+            device_type,
             (world_size // devices_per_node, devices_per_node),
             mesh_dim_names=("replicate", "shard"),
         )
@@ -390,7 +403,14 @@ def train(self):
                     grad_norm=grad_norm.item(),
                 )
 
-                if self.global_step % int(self.args.logging.checkpointing_steps) == 0:
+                should_checkpoint = (
+                    self.global_step % int(self.args.logging.checkpointing_steps) == 0
+                    or (
+                        self.global_step % self.steps_per_epoch == 0
+                        and self.global_step < self.max_train_steps
+                    )
+                )
+                if should_checkpoint:
                     self.save_and_eval_checkpoint()
 
                 if self.suspend_controller.requested():

deepspec/trainer/ckpt_manager.py

@@ -112,7 +112,7 @@ def load_training_state(
     assert saved_local_batch_size == int(local_batch_size)
 
     torch.set_rng_state(checkpoint["torch_rng"])
-    torch.cuda.set_rng_state(checkpoint["torch_cuda_rng"])
+    _set_device_rng_state(checkpoint.get("torch_cuda_rng"))
     np.random.set_state(checkpoint["numpy_rng"])
     random.setstate(checkpoint["python_rng"])
 
@@ -185,6 +185,34 @@ def save_checkpoint(
     return checkpoint_dir
 
 
+def _get_device_rng_state():
+    """Return RNG state of the active accelerator (NPU if available, else CUDA)."""
+    try:
+        import torch_npu  # noqa: F401
+        if torch.npu.is_available():
+            return torch.npu.get_rng_state()
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        return torch.cuda.get_rng_state()
+    return None
+
+
+def _set_device_rng_state(state):
+    """Restore RNG state on the active accelerator (NPU if available, else CUDA)."""
+    if state is None:
+        return
+    try:
+        import torch_npu  # noqa: F401
+        if torch.npu.is_available():
+            torch.npu.set_rng_state(state)
+            return
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        torch.cuda.set_rng_state(state)
+
+
 def _rank_training_state_path(checkpoint_dir: str, global_rank: int) -> str:
     return os.path.join(
         checkpoint_dir,
@@ -213,7 +241,7 @@ def _serialize_training_state(
         "world_size": int(world_size),
         "local_batch_size": int(local_batch_size),
         "torch_rng": torch.get_rng_state(),
-        "torch_cuda_rng": torch.cuda.get_rng_state(),
+        "torch_cuda_rng": _get_device_rng_state(),
         "numpy_rng": np.random.get_state(),
         "python_rng": random.getstate(),
     }

deepspec/utils/__init__.py

@@ -21,7 +21,14 @@ def seed_all(seed):
     torch.manual_seed(seed)
     random.seed(seed)
     np.random.seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    try:
+        import torch_npu  # noqa: F401
+        torch.npu.manual_seed_all(seed)
+        return
+    except Exception:
+        pass
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
 def get_git_sha(detail_info=False):
     import subprocess

deepspec/utils/distributed.py

@@ -9,18 +9,47 @@
 from torch.utils.data import Sampler
 
 
+def _is_npu_available():
+    try:
+        import torch_npu  # noqa: F401
+        return torch.npu.is_available()
+    except Exception:
+        return False
+
+
+def _device_backend_info():
+    """Return (device_module, device_count_fn, set_device_fn, current_device_fn, backend)."""
+    if _is_npu_available():
+        return (
+            torch.npu,
+            torch.npu.device_count,
+            torch.npu.set_device,
+            torch.npu.current_device,
+            "hccl",
+        )
+    return (
+        torch.cuda,
+        torch.cuda.device_count,
+        torch.cuda.set_device,
+        torch.cuda.current_device,
+        "nccl",
+    )
+
+
 def init_dist(local_rank: int, timeout_minutes: int = 60):
-    local_world_size = torch.cuda.device_count()
+    dev_mod, device_count, set_device, _current_device, backend = _device_backend_info()
+    local_world_size = device_count()
     node_rank = int(os.environ["RANK"])
     node_world_size = int(os.environ["WORLD_SIZE"])
     rank = node_rank * local_world_size + local_rank
     world_size = node_world_size * local_world_size
     init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"
-    torch.cuda.set_device(local_rank)
-    device = torch.device("cuda", local_rank)
+    set_device(local_rank)
+    device_type = "npu" if _is_npu_available() else "cuda"
+    device = torch.device(device_type, local_rank)
 
     dist.init_process_group(
-        backend="nccl",
+        backend=backend,
         init_method=init_method,
         rank=rank,
         world_size=world_size,
@@ -35,7 +64,8 @@ def is_global_main_process():
 
 
 def is_local_main_process():
-    return torch.cuda.current_device() == 0
+    _dev_mod, _dc, _sd, current_device, _be = _device_backend_info()
+    return current_device() == 0
 
 
 def print_on_global_main(*args, **kwargs):