spark_model/model/
drop.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3#![allow(unused_imports, dead_code)]
4
5use parking_lot::Mutex;
6use std::collections::HashMap;
7use std::sync::Arc;
8
9use anyhow::{Result, bail};
10use atlas_core::config::{LayerType, ModelConfig};
11use spark_runtime::buffers::BufferArena;
12use spark_runtime::gpu::{DevicePtr, GpuBackend, GraphHandle, KernelHandle};
13use spark_runtime::kv_cache::PagedKvCache;
14
15use super::block_mgmt::{
16    apply_evicted_blocks, ensure_blocks_through_decode, ensure_blocks_through_prefill,
17    extract_layer_refs, reuse_prefix_match_disk_ids,
18};
19use super::ssm_pool::SsmStatePool;
20use super::ssm_snapshot::SsmSnapshotPool;
21use super::types::{PinnedMetaStaging, TransformerModel};
22use crate::layer::{
23    AttnMetadataDev, ForwardContext, GdnPrefillBuffers, LayerState, SsmLayerState, TransformerLayer,
24};
25use crate::layers::ops;
26use crate::speculative::DraftProposer;
27use crate::traits::{ChunkedPrefillPageMetadata, Model, SequenceState};
28use crate::weight_map::{DenseWeight, MtpWeights, QuantizedWeight};
29
30impl Drop for TransformerModel {
31    fn drop(&mut self) {
32        self.drop_pinned_staging();
33        // The SSM spill tier's reusable staging blob is page-locked host memory
34        // owned by the snapshot pool, which holds no `gpu` handle of its own —
35        // same ownership shape as `drop_pinned_staging`. No-op when the tier
36        // never ran (the buffer is allocated on first spill).
37        self.ssm_snapshots.free_staging(self.gpu.as_ref());
38    }
39}