docker run -it -u root -d --net=host \
  --privileged \
  --ipc=host \
  --device=/dev/davinci_manager \
  --device=/dev/devmm_svm \
  --device=/dev/hisi_hdc \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
  -v /usr/local/dcmi:/usr/local/dcmi \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/sbin:/usr/local/sbin \
  -v /etc/hccn.conf:/etc/hccn.conf:ro \
  --name chronos_npu \
  quay.io/ascend/vllm-ascend:v0.18.0-310p-openeuler \
  /bin/bash

git clone https://github.com/amazon-science/chronos-forecasting.git

cd chronos-forecasting

# 安装 Chronos 及其依赖
pip install -e .

# 安装额外依赖（用于 DataFrame 预测接口）
pip install 'pandas[pyarrow]>=2.0,<2.4'
pip install matplotlib

# 使用 魔塔社区 下载模型
modelscope download --model amazon/chronos-2

# 检查 NPU 设备状态
npu-smi info

# 验证 torch_npu 安装
python -c "import torch; import torch_npu; print('NPU available:', torch.npu.is_available()); print('Device:', torch.npu.get_device_name(0))"

NPU available: True
Device: Ascend310P3

# 应用 310P 适配修改
cd chronos-forecasting
git apply 310p-ascend-compat.patch

Atlas 训练系列产品: 数据类型支持float16、float32.           ← 910B
Atlas A2 训练系列产品: 数据类型支持float16、float32、bfloat16. ← 910B 新版
Atlas 推理系列产品: 数据类型仅支持float16、float32.           ← 310P，⚠️ 无 bfloat16

def is_ascend_310p(device: torch.device | None = None) -> bool:
    """Check if we're running on Huawei Ascend 310P chip."""
    try:
        if hasattr(torch, "npu") and torch.npu.is_available():
            if device is not None and device.type != "npu":
                return False
            device_idx = (device.index if device is not None else None) or torch.npu.current_device()
            device_name = torch.npu.get_device_name(device_idx)
            return "310P" in device_name.upper()
    except Exception:
        pass
    return False

# ====== 修改前（原始代码，910B 上正常）======
loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
scale = torch.nan_to_num((x - loc).square().nanmean(dim=-1, keepdim=True).sqrt(), nan=1.0)

# ====== 修改后（310P 适配）======
if is_ascend_310p(x.device):
    # 310P 专用路径
    mask = ~torch.isnan(x)
    mask_sum = mask.sum(dim=-1, keepdim=True).to(x.dtype)
    loc = torch.where(mask, x, torch.zeros_like(x)).sum(dim=-1, keepdim=True) / mask_sum.clamp(min=1.0)
    loc = torch.where(mask_sum == 0, 0.0, loc)
    scale = torch.sqrt(
        torch.where(mask, (x - loc).square(), torch.zeros_like(x)).sum(dim=-1, keepdim=True)
        / mask_sum.clamp(min=1.0)
    )
    scale = torch.where(mask_sum == 0, 1.0, torch.where(scale == 0, self.eps, scale))
else:
    # 标准路径（910B/GPU 使用）
    loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
    scale = torch.nan_to_num((x - loc).square().nanmean(dim=-1, keepdim=True).sqrt(), nan=1.0)

# ====== 修改前（原始代码）======
# Attention implementation - default to "sdpa" if not specified
attn_implementation = attn_implementation or "sdpa"

# ====== 修改后（310P 适配）======
# Attention implementation - default to "eager" on Ascend NPU (310P lacks
# FlashAttention hardware; "sdpa" may fall back or error), otherwise "sdpa"
attn_implementation = attn_implementation or ("eager" if is_ascend_310p() else "sdpa")

# ====== 修改前（原始代码）======
patched_mask = torch.nan_to_num(self.patch(context_mask), nan=0.0)

# ====== 修改后（310P 适配）======
patched_mask_raw = self.patch(context_mask)
if is_ascend_310p(context.device):
    patched_mask = torch.where(torch.isnan(patched_mask_raw), 0.0, patched_mask_raw)
else:
    patched_mask = torch.nan_to_num(patched_mask_raw, nan=0.0)

if is_ascend_310p(device):
    # 310P 专用路径：使用有 NPU 算子支持的替代操作
    result = alternative_operation(x)
else:
    # 标准路径：保留原始高性能实现（910B/GPU 使用）
    result = original_operation(x)

# 修改前（910B3 默认）
pipeline = Chronos2Pipeline.from_pretrained(model_path, device_map=str(device), dtype=torch.bfloat16)

# 修改后（310P 适配）
pipeline = Chronos2Pipeline.from_pretrained(model_path, device_map=str(device), dtype=torch.float16)

# 全局 converter 补齐 — 必须在 import chronos 之前执行
import torch_npu
from compile_npu_accel import apply_compile_patches
apply_compile_patches()

# compile_pipeline — 一次性初始化，自动替换 pipeline.predict 为优化路径
from compile_npu_accel import compile_pipeline
result = compile_pipeline(pipeline, batch_size=4, seq_len=128)
pipeline = result['pipeline']

# 推理 — 完全是原生 API，一行不变
output = pipeline.predict(inputs, prediction_length=24)

def apply_patches():
    # isnan → (x != x)
    torch.isnan = lambda x: x != x
    torch.Tensor.isnan = lambda self: self != self

    # asinh/arcsinh → log(x + sqrt(x²+1))
    _asinh = lambda x: torch.log(x + torch.sqrt(x ** 2 + 1))
    torch.asinh = torch.arcsinh = _asinh

    # sinh → (exp(x) - exp(-x)) / 2
    _sinh = lambda x: (torch.exp(x) - torch.exp(-x)) / 2
    torch.sinh = torch.Tensor.sinh = _sinh

    # is_ascend_310p → 强制 True（消除条件分支）
    chronos.utils.is_ascend_310p = lambda device=None: True

    # nanmean → where + sum + div
    torch.nanmean = lambda x, dim=None, keepdim=False, **kw: (
        torch.where(~(x != x), x, torch.zeros_like(x)).sum(dim=dim, keepdim=keepdim)
        / (~(x != x)).sum(dim=dim, keepdim=keepdim).to(x.dtype).clamp(min=1.0))

    # nan_to_num → where(x != x, nan_value, x)
    torch.nan_to_num = lambda x, nan=0.0, **kw: torch.where(x != x, float(nan), x)

# unfold → reshape
import chronos.chronos2.model as model_module
model_module.Patch.forward = patched_patch_forward

for name, mod in model.named_modules():
    if isinstance(mod, bolt_module.InstanceNorm):
        mod.forward = types.MethodType(patched_instance_norm_forward, mod)
        mod.inverse = types.MethodType(patched_instance_norm_inverse, mod)

def precompute_time_encodings(model, seq_len, num_output_patches, batch_size=4):
    # context 时间编码：[-seq_len, ..., -1] / time_scale，reshape 为 patches
    ctx_enc = torch.arange(-seq_len, 0, dtype=torch.float32, device=model.device)
    ctx_enc = ctx_enc.div(time_scale).to(model.dtype).reshape(seq_len // patch_size, patch_size)
    model.register_buffer('_ctx_time_enc', ctx_enc)

    # future 时间编码：[0, ..., future_len-1] / time_scale，reshape 为 patches
    future_len = num_output_patches * model.chronos_config.output_patch_size
    fut_enc = torch.arange(0, future_len, dtype=torch.float32, device=model.device)
    fut_enc = fut_enc.div(time_scale).to(model.dtype).reshape(num_output_patches, output_patch_size)
    model.register_buffer('_fut_time_enc', fut_enc)

    # 默认 group_ids（逐行独立推理）
    model.register_buffer('_default_group_ids', torch.arange(batch_size, dtype=torch.long, device=model.device))

context_time_enc = self._ctx_time_enc.unsqueeze(0).expand(batch_size, -1, -1)

def patched_prepare_patched_context(self, context, context_mask=None):
    context_mask = (~(context != context)).to(context.dtype)
    patched_mask = torch.where(patched_mask_raw != patched_mask_raw, 0.0, patched_mask_raw)
    context_time_enc = self._ctx_time_enc.unsqueeze(0).expand(batch_size, -1, -1)
    patched_context = torch.cat([context_time_enc, patched_context, patched_mask], dim=-1)

def patched_rope_forward(self, x, position_ids):
    inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
    position_ids_expanded = position_ids[:, None, :].float()
    freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)

def patched_construct_and_invert_group_time_mask(group_ids, attention_mask, floating_type):
    group_mask_f = (group_ids[:, None] == group_ids[None, :]).to(dtype=floating_type)
    attention_mask_f = attention_mask.to(dtype=floating_type)
    group_time_mask = torch.einsum("qb, bt -> qbt", group_mask_f, attention_mask_f)
    return (1.0 - group_time_mask) * torch.finfo(floating_type).min

def patched_forward_inference(self, context, context_mask=None, num_output_patches=1,
                              group_ids=None, **kwargs):
    if context_mask is None:
        context_mask = (~(context != context)).to(torch.float32)
    encoder_outputs, loc_scale, _, num_context_patches = self.encode_inference(
        context=context, context_mask=context_mask, num_output_patches=num_output_patches)
    hidden_states = encoder_outputs[0]
    forecast_embeds = hidden_states[:, -num_output_patches:]
    quantile_preds = self.output_patch_embedding(forecast_embeds)
    quantile_preds = rearrange(quantile_preds, "b n (q p) -> b q (n p)",
                               n=num_output_patches, q=self.num_quantiles,
                               p=self.chronos_config.output_patch_size)
    quantile_preds = self.instance_norm.inverse(quantile_preds, loc_scale)
    return Chronos2Output(quantile_preds=quantile_preds)

def patched_encode_inference(self, context, context_mask, num_output_patches):
    patched_context, attention_mask, loc_scale = self._prepare_patched_context(
        context=context, context_mask=context_mask)
    encoder_outputs = self.encoder(
        attention_mask=attention_mask, inputs_embeds=input_embeds,
        group_ids=self._default_group_ids, output_attentions=False)

def patched_prepare_patched_future_no_covariates(self, loc_scale, num_output_patches, batch_size):
    zeros = torch.zeros(batch_size, num_output_patches, output_patch_size, ...)
    future_time_enc = self._fut_time_enc.unsqueeze(0).expand(batch_size, -1, -1)
    patched_future = torch.cat([future_time_enc, zeros, torch.zeros_like(zeros)], dim=-1)

model.forward = types.MethodType(patched_forward_inference, model)
model.encode_inference = types.MethodType(patched_encode_inference, model)
compiled_model = torch.compile(model, backend='npu', dynamic=False)

def patched_pipeline_predict(self, inputs, prediction_length=24, **kwargs):
    if any(k in kwargs for k in ('covariates', 'future_covariates', 'past_covariates')):
        return self._predict_original(inputs, prediction_length=prediction_length, **kwargs)

    series_list = _normalize_inputs(inputs)
    batches = _pad_to_compiled_shape(series_list, self._compiled_seq_len,
                                      self._compiled_batch_size, self._compiled_device)

    for context_npu, n_series in batches:
        output = self._compiled_model(
            context=context_npu,
            context_mask=(~(context_npu != context_npu)).to(torch.float32),
            num_output_patches=num_output_patches,
        )
        result_cpu = output.quantile_preds[:n_series].cpu()

def compile_pipeline(pipeline, device=DEVICE, batch_size=4, seq_len=128, num_output_patches=2):
    pipeline.model = pipeline.model.to(device).eval()
    compiled_model = apply_patches_and_compile(pipeline.model, ...)
    warmup_compiled_model(compiled_model, ...)
    pipeline._predict_original = pipeline.predict
    pipeline.predict = types.MethodType(patched_pipeline_predict, pipeline)
    return {'compiled_model': compiled_model, 'pipeline': pipeline}

# Step 1: 全局 patches（必须在 import chronos 之前！）
import torch_npu
from compile_npu_accel import apply_compile_patches
apply_compile_patches()

# Step 2: 加载模型 + compile_pipeline（一次性初始化）
from chronos.chronos2 import Chronos2Pipeline
from compile_npu_accel import compile_pipeline

model_path = '/root/predict/chronos2/model_weights'
pipeline = Chronos2Pipeline.from_pretrained(model_path)
result = compile_pipeline(pipeline, batch_size=4, seq_len=128)
pipeline = result['pipeline']

# Step 3: 推理 — 完全是原生 API，一行不变
output = pipeline.predict(inputs, prediction_length=24)
# output: list[Tensor], 每个序列 (1, 21, pred_len), 在 CPU 上

# 3D Tensor — 标准格式（最常用）
output = pipeline.predict(torch.randn(4, 1, 128), prediction_length=24)

# 2D Tensor
output = pipeline.predict(torch.randn(4, 128), prediction_length=24)

# numpy ndarray
output = pipeline.predict(np.random.randn(4, 1, 128).astype(np.float32), prediction_length=24)

# 变长序列 — 自动截断 + 左填充 NaN
output = pipeline.predict([
    torch.randn(64),    # 短序列 → 左填充 NaN 到 128
    torch.randn(200),   # 长序列 → 截断到最后 128
    torch.randn(128),   # 标准长度 → 无需填充
], prediction_length=24)

# 单条序列
output = pipeline.predict(torch.randn(128), prediction_length=24)

output = pipeline.predict(torch.randn(4, 1, 128), prediction_length=24)
# output[0].shape = (1, 21, 24) — 与原始 pipeline.predict 完全一致
# 21 个 quantile, 预测长度 24 步

# 取中位数预测
predictions = output[0]          # (1, 21, 24)
median = predictions[0, 10, :]   # 第 10 个 quantile ≈ 中位数

# 有协变量 → 自动 fallback, 无需修改代码
output = pipeline.predict(inputs, prediction_length=24, covariates=...)

# 也可手动调用原始方法
original_result = pipeline._predict_original(inputs, prediction_length=24, covariates=...)