[feature] Add Ascend NPU Support While Preserving GPU Compatibility

deepspec/data/cuda_prefetcher.py

@@ -2,6 +2,8 @@
 
 import torch
 
+from deepspec.utils import accelerator_module
+
 
 def move_batch_to_device(batch, device):
     moved = {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -22,7 +24,8 @@ class CUDAPrefetcher:
     def __init__(self, dataloader, device):
         self.dataloader = dataloader
         self.device = device
-        self.stream = torch.cuda.Stream(device=device)
+        self.accelerator = accelerator_module()
+        self.stream = self.accelerator.Stream(device=device)
 
     def __iter__(self):
         self._iter = iter(self.dataloader)
@@ -40,7 +43,7 @@ def _fetch_and_transfer(self):
         except StopIteration:
             self._done = True
             return
-        with torch.cuda.stream(self.stream):
+        with self.accelerator.stream(self.stream):
             self._gpu_batch = move_batch_to_device(cpu_batch, self.device)
 
     def __next__(self):
@@ -53,7 +56,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.accelerator.current_stream(self.device)
         current.wait_stream(self.stream)
         batch = self._gpu_batch
 

deepspec/eval/dspark/confidence_head.py

@@ -27,6 +27,11 @@ def model_display_name(path: str) -> str:
     return os.path.basename(normalized) or normalized
 
 
+def confidence_metric_dtype(device: torch.device) -> torch.dtype:
+    device = torch.device(device)
+    return torch.float32 if device.type == "npu" else torch.float64
+
+
 class PerPositionConfidenceMetrics:
     """Per-position ECE + AUROC + Brier for cumprod predictions."""
 
@@ -41,22 +46,23 @@ def __init__(
         self.block_size = int(block_size)
         self.num_coarse_bins = int(num_coarse_bins)
         self.num_fine_bins = int(num_fine_bins)
+        self.dtype = confidence_metric_dtype(device)
         self.coarse_count = torch.zeros(
             (self.block_size, self.num_coarse_bins),
-            dtype=torch.float64,
+            dtype=self.dtype,
             device=device,
         )
         self.coarse_pred = torch.zeros_like(self.coarse_count)
         self.coarse_target = torch.zeros_like(self.coarse_count)
         self.fine_pos = torch.zeros(
             (self.block_size, self.num_fine_bins),
-            dtype=torch.float64,
+            dtype=self.dtype,
             device=device,
         )
         self.fine_neg = torch.zeros_like(self.fine_pos)
         self.brier_num = torch.zeros(
             self.block_size,
-            dtype=torch.float64,
+            dtype=self.dtype,
             device=device,
         )
 
@@ -66,12 +72,12 @@ def update(
         probs: torch.Tensor,
         targets: torch.Tensor,
     ) -> None:
-        probs = probs.reshape(-1).to(torch.float64).clamp(EPS_PROB, 1.0 - EPS_PROB)
-        targets = targets.reshape(-1).to(torch.float64)
+        probs = probs.reshape(-1).to(self.dtype).clamp(EPS_PROB, 1.0 - EPS_PROB)
+        targets = targets.reshape(-1).to(self.dtype)
         assert probs.shape == targets.shape
         pos_count = probs.numel()
         assert pos_count <= self.block_size
-        weights = torch.ones_like(probs, dtype=torch.float64)
+        weights = torch.ones_like(probs, dtype=self.dtype)
         pos_idx = torch.arange(pos_count, device=probs.device)
 
         coarse_idx = (
@@ -363,11 +369,11 @@ def observe(
         step_probs = torch.sigmoid(
             confidence_logits[:, :effective_length]
         ).squeeze(0)
-        cumprod_pred = step_probs.to(torch.float64).cumprod(dim=0)
+        cumprod_pred = step_probs.to(self.dataset_metrics.dtype).cumprod(dim=0)
         prefix_label = (
             verification.accept_prefix_mask[:, :effective_length]
             .squeeze(0)
-            .to(torch.float64)
+            .to(self.dataset_metrics.dtype)
         )
         self.dataset_metrics.update(
             probs=cumprod_pred,

deepspec/modeling/dspark/common.py

@@ -82,7 +82,34 @@ def create_dspark_attention_mask(
     seq_len: int,
     block_size: int,
     device: torch.device,
+    attn_implementation: str = "flex_attention",
 ):
+    if attn_implementation != "flex_attention":
+        bsz, num_blocks = anchor_positions.shape
+        q_len = num_blocks * block_size
+        kv_len = seq_len + q_len
+        q_idx = torch.arange(q_len, device=device)
+        kv_idx = torch.arange(kv_len, device=device)
+        q_block_ids = (q_idx // block_size).unsqueeze(0).expand(bsz, -1)
+        anchor_pos = anchor_positions.gather(1, q_block_ids).unsqueeze(-1)
+        q_block_ids = q_block_ids.unsqueeze(-1)
+        kv_idx = kv_idx.view(1, 1, kv_len)
+
+        is_context = kv_idx < seq_len
+        mask_context = is_context & (kv_idx < anchor_pos)
+        is_draft = kv_idx >= seq_len
+        kv_block_ids = (kv_idx - seq_len) // block_size
+        mask_draft = is_draft & (q_block_ids == kv_block_ids)
+        is_valid_block = block_keep_mask.gather(
+            1,
+            q_block_ids.squeeze(-1),
+        ).unsqueeze(-1)
+        dense_mask = (mask_context | mask_draft) & is_valid_block
+        empty_rows = ~dense_mask.any(dim=-1, keepdim=True)
+        self_kv_idx = int(seq_len) + q_idx.view(1, -1, 1)
+        dense_mask = dense_mask | (empty_rows & (kv_idx == self_kv_idx))
+        return dense_mask.unsqueeze(1)
+
     def dspark_mask_mod(b, h, q_idx, kv_idx):
         del h
         q_block_id = q_idx // block_size

deepspec/modeling/dspark/gemma4/config.py

@@ -1,9 +1,10 @@
 import copy
 
 from deepspec.modeling.dspark.common import validate_target_layer_ids
+from deepspec.utils import is_npu_available
 
 
-TRAIN_ATTN_IMPLEMENTATION = "flex_attention"
+TRAIN_ATTN_IMPLEMENTATION = "sdpa" if is_npu_available() else "flex_attention"
 
 
 def get_gemma4_text_config(target_config):
@@ -84,6 +85,11 @@ def build_draft_config(target_config, model_args):
     draft_config.target_text_model_type = str(draft_config.model_type)
     draft_config.num_target_layers = num_target_layers
     draft_config.num_hidden_layers = num_draft_layers
+    # The public Gemma4 DSpark draft block is dense-only. Gemma4 A4B targets are
+    # MoE, but the draft model can still train as a dense proposal network
+    # against target hidden states from the MoE model.
+    draft_config.enable_moe_block = False
+    draft_config.hidden_size_per_layer_input = 0
     draft_config.block_size = int(model_args.block_size)
     draft_config.tie_word_embeddings = False
     draft_config.layer_types = layer_types

deepspec/modeling/dspark/gemma4/modeling.py

@@ -484,6 +484,7 @@ def forward(
             seq_len=seq_len,
             block_size=self.block_size,
             device=device,
+            attn_implementation=self.config._attn_implementation,
         )
         output_hidden = self._forward_backbone(
             position_ids=full_position_ids,

deepspec/modeling/dspark/qwen3/config.py

@@ -1,9 +1,10 @@
 import copy
 
 from deepspec.modeling.dspark.common import validate_target_layer_ids
+from deepspec.utils import is_npu_available
 
 
-TRAIN_ATTN_IMPLEMENTATION = "flex_attention"
+TRAIN_ATTN_IMPLEMENTATION = "sdpa" 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,
+            attn_implementation=self.config._attn_implementation,
         )
         output_hidden = self._forward_backbone(
             position_ids=full_position_ids,

deepspec/trainer/__init__.py

@@ -1,6 +1,17 @@
 from .base_trainer import BaseTrainer
 from .dspark_trainer import Gemma4DSparkTrainer, Qwen3DSparkTrainer
-from .eagle3_trainer import Gemma4Eagle3Trainer, Qwen3Eagle3Trainer
+
+
+def __getattr__(name):
+    if name in {"Gemma4Eagle3Trainer", "Qwen3Eagle3Trainer"}:
+        from .eagle3_trainer import Gemma4Eagle3Trainer, Qwen3Eagle3Trainer
+
+        return {
+            "Gemma4Eagle3Trainer": Gemma4Eagle3Trainer,
+            "Qwen3Eagle3Trainer": Qwen3Eagle3Trainer,
+        }[name]
+    raise AttributeError(name)
+
 
 __all__ = [
     "BaseTrainer",

deepspec/trainer/base_trainer.py

@@ -15,6 +15,8 @@
 from deepspec.utils import (
     BF16Optimizer,
     StatelessResumableDistributedSampler,
+    device_count,
+    device_type,
     ensure_dir,
     init_dist,
     is_global_main_process,
@@ -65,9 +67,9 @@ def _build_fsdp_kwargs(
         sharding_strategy=sharding_strategy,
     )
     if sharding_strategy in _HYBRID_STRATEGIES:
-        devices_per_node = torch.cuda.device_count()
+        devices_per_node = device_count()
         fsdp_kwargs["device_mesh"] = init_device_mesh(
-            "cuda",
+            device_type(),
             (world_size // devices_per_node, devices_per_node),
             mesh_dim_names=("replicate", "shard"),
         )
@@ -283,6 +285,7 @@ def _wrap_with_fsdp(self, model):
             precision_dtype=self.precision_dtype,
             world_size=self.world_size,
         )
+        fsdp_kwargs["device_id"] = self.device
         return FSDP(model, **fsdp_kwargs)
 
     def _build_train_dataloader(self, start_offset_samples=0, num_samples=None):

deepspec/trainer/ckpt_manager.py

@@ -12,10 +12,12 @@
 
 from deepspec.utils import (
     ensure_dir,
+    get_rng_state,
     is_global_main_process,
     print_on_global_main,
     print_on_local_main,
     safe_symlink,
+    set_rng_state,
 )
 
 
@@ -112,7 +114,9 @@ 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"])
+    rng_state = checkpoint.get("torch_accelerator_rng", checkpoint.get("torch_cuda_rng"))
+    if rng_state is not None:
+        set_rng_state(rng_state)
     np.random.set_state(checkpoint["numpy_rng"])
     random.setstate(checkpoint["python_rng"])
 
@@ -213,7 +217,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_accelerator_rng": get_rng_state(),
         "numpy_rng": np.random.get_state(),
         "python_rng": random.getstate(),
     }

deepspec/utils/__init__.py

@@ -4,6 +4,20 @@
 import torch
 
 from .config import CustomJSONEncoder, jsonable, load_config, parse_opts_to_config
+from .device import (
+    accelerator_backend,
+    accelerator_module,
+    current_device_index,
+    device_count,
+    device_type,
+    empty_cache,
+    get_rng_state,
+    is_npu_available,
+    make_device,
+    manual_seed_all,
+    set_device,
+    set_rng_state,
+)
 from .distributed import (
     StatelessResumableDistributedSampler,
     init_dist,
@@ -21,7 +35,7 @@ def seed_all(seed):
     torch.manual_seed(seed)
     random.seed(seed)
     np.random.seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    manual_seed_all(seed)
 
 def get_git_sha(detail_info=False):
     import subprocess
@@ -67,21 +81,33 @@ def get_git_diff(rev="HEAD"):
     "BF16Optimizer",
     "CustomJSONEncoder",
     "StatelessResumableDistributedSampler",
+    "accelerator_backend",
+    "accelerator_module",
     "add_metric",
+    "current_device_index",
+    "device_count",
+    "device_type",
+    "empty_cache",
     "ensure_dir",
     "flush",
     "get_git_diff",
     "get_git_sha",
+    "get_rng_state",
+    "is_npu_available",
     "init_dist",
     "is_global_main_process",
     "is_local_main_process",
     "jsonable",
     "load_config",
+    "make_device",
+    "manual_seed_all",
     "main_process_first",
     "parse_opts_to_config",
     "print_on_global_main",
     "print_on_local_main",
     "reset",
     "safe_symlink",
+    "set_device",
+    "set_rng_state",
     "seed_all",
 ]