Skip to content

append_projected_kv fusion #665

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CC-Yeh merged 3 commits into from
Jul 31, 2026
Merged

append_projected_kv fusion #665

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
41032be
state_kv_projection fusion gemm append_kv and cleanups
CC-Yeh Jul 30, 2026
cbefb23
Merge branch 'main' into dflash_improvements2
CC-Yeh Jul 31, 2026
bfd45c5
Merge branch 'main' into dflash_improvements2
CC-Yeh Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions crates/backend-uzu/src/backends/cpu/kernel/attention/qkv_norm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,7 @@ pub fn qkv_norm<
#[optional(!scale_free)] scales: Option<*const ScaleT>,
qkv_output: *mut OutputT,
batch_size: u32,
num_q_heads: u32,
num_kv_heads: u32,
total_heads: u32,
head_dim: u32,
epsilon: f32,
scale_offset: f32,
Expand All @@ -38,7 +37,7 @@ pub fn qkv_norm<
let head_dim = head_dim as usize;
let head_offset = head_offset as usize;
let head_count = head_count as usize;
let qkv_stride = (num_q_heads + 2 * num_kv_heads) as usize * head_dim;
let qkv_stride = total_heads as usize * head_dim;
let head_dim_accum = AccumT::from(head_dim).unwrap();
let epsilon = AccumT::from(epsilon).unwrap();
let scale_offset = AccumT::from(scale_offset).unwrap();
Expand Down
10 changes: 2 additions & 8 deletions crates/backend-uzu/src/backends/metal/kernel/attention/qkv_norm.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,7 @@ PUBLIC KERNEL(QKVNorm)(
const device ScaleT* scales OPTIONAL(!scale_free),
device OutputT* qkv_output,
constant uint& batch_size,
constant uint& num_q_heads,
constant uint& num_kv_heads,
constant uint& total_heads,
constant uint& head_dim,
constant float& epsilon,
constant float& scale_offset,
Expand All @@ -41,13 +40,8 @@ PUBLIC KERNEL(QKVNorm)(
if (head_count == 0u || head_dim == 0u)
return;

const uint total_heads_in_buffer = num_q_heads + 2u * num_kv_heads;
const uint logical_head_idx = head_offset + head_idx;
if (logical_head_idx >= total_heads_in_buffer)
return;

const ulong slice_offset =
(ulong)batch_idx * (ulong)total_heads_in_buffer * (ulong)head_dim + (ulong)logical_head_idx * (ulong)head_dim;
(ulong)batch_idx * (ulong)total_heads * (ulong)head_dim + (ulong)(head_offset + head_idx) * (ulong)head_dim;

const device InputT* input_data = qkv_input + slice_offset;
const device ScaleT* scales_data = scales;
Expand Down
52 changes: 33 additions & 19 deletions crates/backend-uzu/src/encodable_block/dflash.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ impl<B: Backend> DFlashState<B> {
pub struct DFlashDraft<B: Backend> {
target_feature_projection: Box<dyn Linear<B>>,
projected_feature_norm: Normalization<B>,
state_kv_projection: Box<dyn Linear<B>>,
layer_kv_dim: usize,
layers: Box<[DFlashDraftLayer<B>]>,
output_norm: Normalization<B>,
top_k: <B::Kernels as Kernels>::RadixTopKSmall,
Expand Down Expand Up @@ -119,7 +121,6 @@ impl<B: Backend> DFlashDraft<B> {
"DFlash block_size must be in 1..=attention suffix capacity",
));
}

let target_feature_projection = <dyn Linear<B>>::new(
config.model_dim * config.target_layer_ids.len(),
[config.model_dim],
Expand All @@ -139,6 +140,16 @@ impl<B: Backend> DFlashDraft<B> {
context,
)?;
let layers_tree = parameter_tree.subtree("layers")?;
let layer_kv_dim =
2 * config.layer_configs[0].attention_config.num_groups * config.layer_configs[0].attention_config.head_dim;
Comment thread
CC-Yeh marked this conversation as resolved.
let state_kv_projection = <dyn Linear<B>>::new(
config.model_dim,
[config.layer_configs.len() * layer_kv_dim],
false,
Comment thread
CC-Yeh marked this conversation as resolved.
context,
data_type,
&parameter_tree.subtree("state_kv_projection")?,
)?;
let layers = config
.layer_configs
.iter()
Expand Down Expand Up @@ -172,6 +183,8 @@ impl<B: Backend> DFlashDraft<B> {
Ok(Self {
target_feature_projection,
projected_feature_norm,
state_kv_projection,
layer_kv_dim,
layers,
output_norm,
top_k,
Expand Down Expand Up @@ -244,20 +257,27 @@ impl<B: Backend> DFlashDraft<B> {
let token_positions = (state.context_length..state.context_length + num_tokens).collect::<Box<[_]>>();
let rope = PrecalculatedRoPE::precalculate(&self.rope_config, &token_positions, encoder)?;

let dflash_layer_count = self.layers.len();
let mut normalized_features = Some(normalized_features);
for (layer_index, (layer, attention_state)) in self.layers.iter().zip(state.layer_states.iter_mut()).enumerate()
let projected_kv = self.state_kv_projection.encode(normalized_features, num_tokens, encoder)?;
let layer_kv_bytes = self.layer_kv_dim * self.data_type.size_in_bytes();
let kv_chunk = |chunk_index: usize| chunk_index * layer_kv_bytes..(chunk_index + 1) * layer_kv_bytes;
let mut layer_key_values = (0..self.layers.len())
.map(|_| encoder.allocate_scratch(num_tokens * layer_kv_bytes))
.collect::<Result<Box<[_]>, _>>()?;
for (layer_index, key_value) in layer_key_values.iter_mut().enumerate() {
for token_index in 0..num_tokens {
encoder.encode_copy(
&projected_kv,
kv_chunk(token_index * self.layers.len() + layer_index),
key_value,
kv_chunk(token_index),
);
}
}
for ((layer, attention_state), key_value) in
self.layers.iter().zip(state.layer_states.iter_mut()).zip(layer_key_values)
{
attention_state.prepare(state.context_length, num_tokens, encoder.context())?;
let kv_input = if layer_index + 1 == dflash_layer_count {
normalized_features.take().expect("normalized features available for last layer")
} else {
let source = normalized_features.as_ref().expect("normalized features available");
let mut kv_input = encoder.allocate_scratch(source.size())?;
encoder.encode_copy(source, .., &mut kv_input, ..);
kv_input
};
layer.attention.append_kv(kv_input, Some(&rope), num_tokens, attention_state, encoder)?;
layer.attention.append_projected_kv(key_value, &rope, num_tokens, attention_state, encoder)?;
}

state.context_length += num_tokens;
Expand Down Expand Up @@ -376,12 +396,6 @@ impl<B: Backend> DFlashDraftLayer<B> {
parameter_tree: &ParameterTree<B>,
data_type: DataType,
) -> Result<Self, DFlashDraftNewError<B>> {
if config.attention_config.is_causal {
return Err(DFlashDraftNewError::InvalidAttentionConfig("DFlash attention must be non-causal"));
}
if config.attention_config.is_kv_sharing {
return Err(DFlashDraftNewError::InvalidAttentionConfig("DFlash attention must not use KV sharing"));
}
let (attention, _) = Attention::new(
Comment thread
CC-Yeh marked this conversation as resolved.
model_dim,
data_type,
Expand Down
19 changes: 8 additions & 11 deletions crates/backend-uzu/src/encodable_block/mixer/attention/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use crate::{
Mixer, MixerState,
attention::{
core::{AttentionCoreNewArguments, AttentionCores},
mode::{LinearProjection, QkvProjection},
mode::LinearProjection,
qkv_norm::{QKVNorm, QKVNormError},
rope::PrecalculatedRoPE,
},
Expand All @@ -41,7 +41,8 @@ pub struct Attention<B: Backend> {
sliding_window_size: Option<usize>,
max_rope_length: Option<usize>,
data_type: DataType,
projection: QkvProjection<B>,
qkv: LinearProjection<B>,
prepare: <B::Kernels as Kernels>::AttentionPrepareKernel,
gate_projection: Option<Box<dyn Linear<B>>>,
sinks: Option<Allocation<B>>,
flat_core: AttentionCores<B>,
Expand Down Expand Up @@ -158,14 +159,6 @@ impl<B: Backend> Attention<B> {
rope_config.is_some(),
)
.map_err(AttentionNewError::Backend)?;
let projection = QkvProjection::Packed {
qkv: LinearProjection {
lin: qkv_projection,
norm: qkv_norm,
},
prepare,
};

let sinks = config
.has_sinks
.then(|| parameter_tree.leaf("sinks")?.validate(&[num_q_heads], data_type)?.read_allocation())
Expand Down Expand Up @@ -230,7 +223,11 @@ impl<B: Backend> Attention<B> {
sliding_window_size,
max_rope_length,
data_type,
projection,
qkv: LinearProjection {
lin: qkv_projection,
norm: qkv_norm,
},
prepare,
gate_projection,
sinks,
flat_core,
Expand Down
Loading
Loading