metrale_model_engine/model/
impl_a3.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The embedding-scale and logit-softcap launches, and the
4//! accessors for the single-token decode logits buffer.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports, dead_code)]
10
11use parking_lot::Mutex;
12use std::collections::HashMap;
13use std::sync::Arc;
14
15use anyhow::{Result, bail};
16use metrale_cache::kv_cache::PagedKvCache;
17use metrale_config::{LayerType, ModelConfig};
18use metrale_gpu_runtime::buffers::BufferArena;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, GraphHandle, KernelHandle};
20
21use super::block_mgmt::{
22    apply_evicted_blocks, ensure_blocks_through_decode, ensure_blocks_through_prefill,
23    extract_layer_refs, reuse_prefix_match_disk_ids,
24};
25use super::ssm_pool::SsmStatePool;
26use super::ssm_snapshot::SsmSnapshotPool;
27use super::types::{PinnedMetaStaging, TransformerModel};
28use crate::traits::{ChunkedPrefillPageMetadata, Model, SequenceState};
29use metrale_model_layers::layer::{
30    AttnMetadataDev, ForwardContext, GdnPrefillBuffers, LayerState, SsmLayerState, TransformerLayer,
31};
32use metrale_model_layers::layers::ops;
33use metrale_model_layers::speculative::DraftProposer;
34use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, QuantizedWeight};
35
36impl TransformerModel {
37    /// 2026-09-25: Scale `num_tokens` BF16 embedding rows in place by
38    /// `config.embed_scale`; a no-op when the scale kernel is not loaded.
39    pub(super) fn scale_embeddings(
40        &self,
41        data: DevicePtr,
42        num_tokens: usize,
43        stream: u64,
44    ) -> Result<()> {
45        self.scale_embeddings_bf16(data, num_tokens, stream)
46    }
47
48    pub(super) fn scale_embeddings_bf16(
49        &self,
50        data: DevicePtr,
51        num_tokens: usize,
52        stream: u64,
53    ) -> Result<()> {
54        if self.embed_scale_kernel.0 == 0 {
55            return Ok(());
56        }
57        use metrale_gpu_runtime::kernel_args::KernelLaunch;
58        let n = (num_tokens * self.config.hidden_size) as u32;
59        KernelLaunch::new(self.gpu.as_ref(), self.embed_scale_kernel)
60            .grid([n.div_ceil(256), 1, 1])
61            .block([256, 1, 1])
62            .arg_ptr(data)
63            .arg_u32(n)
64            .arg_f32(self.config.embed_scale)
65            .launch(stream)
66    }
67
68    /// 2026-09-25: Softcap BF16 logits in place: `logits[i] = cap * tanh(logits[i] / cap)`.
69    /// The caller must hold a loaded `logit_softcap_kernel`;
70    /// `apply_logit_softcap_dtype` checks instead.
71    pub(super) fn apply_logit_softcap(
72        &self,
73        logits: DevicePtr,
74        num_elements: u32,
75        cap: f32,
76        stream: u64,
77    ) -> Result<()> {
78        use metrale_gpu_runtime::kernel_args::KernelLaunch;
79        let inv_cap = 1.0f32 / cap;
80        KernelLaunch::new(self.gpu.as_ref(), self.logit_softcap_kernel)
81            .grid([num_elements.div_ceil(256), 1, 1])
82            .block([256, 1, 1])
83            .arg_ptr(logits)
84            .arg_u32(num_elements)
85            .arg_f32(inv_cap)
86            .arg_f32(cap)
87            .launch(stream)
88    }
89
90    /// 2026-09-25: Softcap `logits` in place with the BF16 or the FP32 kernel,
91    /// per `is_fp32`. A no-op when that kernel is not loaded: the BF16 one is
92    /// loaded only when the config sets `final_logit_softcapping`, the FP32 one
93    /// never.
94    pub(super) fn apply_logit_softcap_dtype(
95        &self,
96        logits: DevicePtr,
97        num_elements: u32,
98        cap: f32,
99        is_fp32: bool,
100        stream: u64,
101    ) -> Result<()> {
102        use metrale_gpu_runtime::kernel_args::KernelLaunch;
103        let kernel = if is_fp32 {
104            self.logit_softcap_fp32_kernel
105        } else {
106            self.logit_softcap_kernel
107        };
108        if kernel.0 == 0 {
109            return Ok(());
110        }
111        let inv_cap = 1.0f32 / cap;
112        KernelLaunch::new(self.gpu.as_ref(), kernel)
113            .grid([num_elements.div_ceil(256), 1, 1])
114            .block([256, 1, 1])
115            .arg_ptr(logits)
116            .arg_u32(num_elements)
117            .arg_f32(inv_cap)
118            .arg_f32(cap)
119            .launch(stream)
120    }
121
122    /// 2026-09-25: Whether the single-token decode `lm_head` writes FP32 logits
123    /// to `logits_fp32_buf` (`use_fp32_logits`, which `TransformerModel::new`
124    /// sets to false). Other lm_head paths always write BF16.
125    pub fn decode_logits_fp32(&self) -> bool {
126        self.use_fp32_logits
127    }
128
129    /// 2026-09-25: The buffer the single-token decode `lm_head` writes:
130    /// `logits_fp32_buf` when `use_fp32_logits`, else the shared BF16 logits
131    /// buffer. Read it with the dtype `decode_logits_fp32` reports.
132    pub fn decode_logits_ptr(&self) -> DevicePtr {
133        if self.use_fp32_logits {
134            self.logits_fp32_buf
135        } else {
136            self.buffers.logits()
137        }
138    }
139}
metrale_model_engine/model/impl_a3.rs

metrale_model_engine/model/
impl_a3.rs
