Migration to allocations by LuckyIYI · Pull Request #349 · trymirai/uzu

LuckyIYI · 2026-04-16T16:55:11Z

No description provided.

… code

…ssifier cleanup

norpadon

Accidentally stumbled upon this PR and found a few things to nitpick

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 70ed067262

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

CC-Yeh · 2026-04-21T09:03:09Z

-
-        let src = &self.suffix_state.as_bytes()[src_start..src_end];
-        let dst = self.conv_state.as_bytes_mut();
-        dst.copy_from_slice(src);


dst.copy_from_slice(src);

This used to panic when length mismatch I think.

And now it silently clamps the len, is this behavior expected?

# Conflicts: # crates/backend-uzu/build/common/traitgen.rs # crates/backend-uzu/build/cpu/compiler.rs # crates/backend-uzu/src/backends/cpu/kernel/ssm/ssd_prefill.rs # crates/backend-uzu/src/backends/metal/kernel/matmul/mod.rs # crates/backend-uzu/src/encodable_block/classifier_layer.rs # crates/backend-uzu/src/encodable_block/layer/executables.rs # crates/backend-uzu/src/encodable_block/rope.rs

uuuvn · 2026-05-03T15:41:02Z

-pub(super) fn read_float_array<B: Backend, const RANK: usize>(
-    tree: &BackendParameterTree<'_, B>,
+pub(super) fn read_float_allocation<B: Backend, const RANK: usize>(
+    tree: &BackendParameterTree<B>,
    name: &str,
    expected_data_type: DataType,
 ) -> AudioResult<([usize; RANK], Array<B>)> {


Why are you renaming read_float_array to read_float_allocation when this function still returns array?

uuuvn · 2026-05-03T15:46:01Z

-    let (shape, array) = read_float_array::<B, 1>(tree, name, expected_data_type)?;
+    let (shape, allocation) = read_float_allocation::<B, 1>(tree, name, expected_data_type)?;


Same as previous one:

This is array, not an allocation

I won't flag more changes that are almost exactly like that, please read the diff before re-requesting review. A change that renames a bunch of functions and variables that were and are still arrays into allocations should never make it's way into a supposedly ready to review pull request.

uuuvn · 2026-05-03T16:02:32Z

-    pub hidden: &'a mut B::Buffer,
+    pub hidden: &'a mut Allocation<B>,
    /// Final output [d_model]
-    pub y: &'a mut B::Buffer,
+    pub y: &'a mut Allocation<B>,


Hidden is not an argument but an intermediate (the comment literally says this + it's fully overwritten in pass a without being read), y should similarly be a return of MoeExpertsSingleDecodeKernels::encode

Also looking at lsp references for this it looks like it's just a dead code (nothing outside of tests encodes this kernel), so it should be just deleted (in a separate prereq pr).

uuuvn · 2026-05-03T16:12:15Z

    type Completed: CommandBufferCompleted<CommandBuffer = Self>;
 }

+type CommandBufferBackend<C> = <C as CommandBuffer>::Backend;


Let's not have "getter types", things like Allocation<CommandBufferBackend<Self::CommandBuffer>> are much more confusing than Allocation<<Self::CommandBuffer as CommandBuffer>::Backend>

uuuvn · 2026-05-03T16:36:42Z

    pub fn encode_copy(
        &mut self,
-        src: &B::Buffer,
-        src_range: Range<usize>,
-        dst: &mut B::Buffer,
-        dst_range: Range<usize>,
+        src: &Allocation<B>,
+        src_range: impl RangeBounds<usize>,
+        dst: &mut Allocation<B>,
+        dst_range: impl RangeBounds<usize>,
    ) {
-        let size = src_range.end - src_range.start;
-        debug_assert_eq!(size, dst_range.end - dst_range.start);
+        let (src_buffer, src_allocation_range) = src.as_buffer_range();
+        let (dst_buffer, dst_allocation_range) = dst.as_buffer_range();
+        let src_range = resolve_copy_range(src_range, src_allocation_range.len(), "source");
+        let dst_range = resolve_copy_range(dst_range, dst_allocation_range.len(), "destination");
+        let byte_len = src_range.len();
+        assert_eq!(byte_len, dst_range.len(), "copy range lengths must match");
+        assert!(byte_len > 0, "zero-sized copies are not allowed");
+        let src_access_range = src_allocation_range.start + src_range.start..src_allocation_range.start + src_range.end;
+        let dst_access_range = dst_allocation_range.start + dst_range.start..dst_allocation_range.start + dst_range.end;
        self.access(&[
            Access {
-                range: src.gpu_address_subrange(src_range.clone()),
+                range: src_buffer.gpu_address_subrange(src_access_range),
                flags: AccessFlags::copy_read(),
            },
            Access {
-                range: dst.gpu_address_subrange(dst_range.clone()),
+                range: dst_buffer.gpu_address_subrange(dst_access_range),
                flags: AccessFlags::copy_write(),
            },
        ]);


This is very confusing and it implements the same range add thing both here and in backend's implementation, this is not the right way to do things.

eugenebokhan · 2026-05-07T12:57:09Z

+    pub down_biases: &'a Allocation<B>,
    /// Hidden buffer [K, d_ff] - intermediate storage (f32)
-    pub hidden: &'a mut B::DenseBuffer,
+    pub hidden: &'a mut Allocation<B>,


I think this should not be an argument any more and but a part of MoeExpertsSingleDecodeKernels::encode

eugenebokhan · 2026-05-07T12:58:17Z

+    pub hidden: &'a mut Allocation<B>,
    /// Final output [d_model]
-    pub y: &'a mut B::DenseBuffer,
+    pub y: &'a mut Allocation<B>,


It should be returned by the encode

eugenebokhan · 2026-05-07T12:59:54Z

+    pub hidden: &'a mut Allocation<B>,
+    pub output: &'a mut Allocation<B>,


hidden & the output should be created inside

eugenebokhan · 2026-05-07T13:15:37Z

+        let mut x = encoder.allocate_scratch(size_for_shape(
+            &[active_row_count, self.config.num_heads, self.config.head_dim],
+            self.data_type,
+        ))?;
+        let bc_shape = [active_row_count, self.config.num_groups, self.config.state_dim];
+        let mut b = encoder.allocate_scratch(size_for_shape(&bc_shape, self.data_type))?;
+        let mut c = encoder.allocate_scratch(size_for_shape(&bc_shape, self.data_type))?;


I believe x, b, c should be allocated inside run_conv_scan and returned bound in a struct

eugenebokhan · 2026-05-07T13:17:13Z

+        let in_proj = self.in_projection.encode(input, active_row_count, encoder)?;
+        let mut conv_inputs =
+            encoder.allocate_scratch(size_for_shape(&[active_row_count, self.config.conv_dim()], self.data_type))?;
+        let mut z = encoder.allocate_scratch(size_for_shape(
+            &[active_row_count, self.config.num_heads, self.config.head_dim],
+            self.data_type,
+        ))?;
+        let mut dt =
+            encoder.allocate_scratch(size_for_shape(&[active_row_count, self.config.num_heads], self.data_type))?;


Looks like in_proj, z and dt should be allocated inside self.run_split_inproj and returned as a struct

eugenebokhan · 2026-05-07T13:24:00Z

-            ArrayId::AttentionOutput,
-            ArrayId::Main,
        )
        .expect("Failed to create out-projection kernel");


replace expect with error

eugenebokhan · 2026-05-07T13:24:41Z

        model_dim: usize,
        decoder_layer_loader: &ParameterTree<B::Context>,
-    ) -> (Self, Option<B::DenseBuffer>) {
+    ) -> (Self, Option<Allocation<B>>) {


should return a Result

eugenebokhan · 2026-05-07T13:24:52Z

-    ) -> (Self, Option<B::DenseBuffer>) {
+    ) -> (Self, Option<Allocation<B>>) {
        if !matches!(layer_type, DecoderLayerType::ShortConv { .. }) {
            panic!("Layer {} marked as non-ShortConv but ShortConv config provided", layer_index);


replace panic! with error

eugenebokhan · 2026-05-07T13:25:07Z

-            ArrayId::Main,
-            ArrayId::SsmInProj,
        )
        .expect("Failed to create in-projection kernel");


replace expect with error

eugenebokhan · 2026-05-07T13:25:12Z

-            ArrayId::AttentionOutput,
-            ArrayId::Main,
        )
        .expect("Failed to create out-projection kernel");


replace expect with error

LuckyIYI added 12 commits April 16, 2026 00:44

wip

c7897c1

cleaning up

216bac8

Merge branch 'main' into buffers-to-allocations

89abea8

rht, hadamard to allocations

6792915

deleted generic stuff

f428e08

removed bitmask and token seeds, and cleanup

b2d1174

Refactor decode ownership to allocation-local outputs and remove dead…

4570c90

… code

fix allocation lifetimes and attention f32 scratch buffers

9b5d4ba

make arrays allocation-backed and finish allocation-native kernel/cla…

484503a

…ssifier cleanup

clenaup

7e4b6be

clean up

67f2169

review fixes

bb31978

norpadon requested changes Apr 20, 2026

View reviewed changes

LuckyIYI added 5 commits April 20, 2026 10:08

Merge branch 'main' into buffers-to-allocations

f101760

comment fixes

37f4911

proper mut for in place ops

c870878

tests clean up, some stuff reverted to main

0f58496

fixed merge resolution errors

f8c5eb9

LuckyIYI requested a review from norpadon April 20, 2026 22:08

LuckyIYI changed the title ~~Migration to allocations [wip]~~ Migration to allocations Apr 20, 2026

using slices on CPU again for shortconv

ae50ba5

LuckyIYI requested review from CC-Yeh, eugenebokhan and uuuvn April 20, 2026 22:30

fix CI

70ed067

LuckyIYI marked this pull request as ready for review April 20, 2026 23:27

chatgpt-codex-connector Bot reviewed Apr 20, 2026

View reviewed changes

Comment thread crates/backend-uzu/src/backends/common/encoder.rs Outdated

Comment thread crates/backend-uzu/src/array.rs Outdated

CI fixes + code review fixes

02e8fce

uuuvn requested changes Apr 21, 2026

View reviewed changes

CC-Yeh requested changes Apr 21, 2026

View reviewed changes

LuckyIYI added 2 commits April 21, 2026 13:56

review fixes

25f0058

LuckyIYI requested review from CC-Yeh and uuuvn April 22, 2026 09:07

uuuvn requested changes Apr 22, 2026

View reviewed changes

LuckyIYI added 10 commits April 22, 2026 13:30

fix main parity

5aa7b8d

addressing code review

11b817b

Merge branch 'main' into buffers-to-allocations

c00cabd

reverted unrelated changes

60910f2

Merge branch 'main' into buffers-to-allocations

e69519d

Merge branch 'main' into buffers-to-allocations

d0e7a2c

fix audio

7717609

array cleanup

e9c8c85

Merge branch 'main' into buffers-to-allocations

ad9b96d

restored some test + lang generator clean up

c3d830c

LuckyIYI requested a review from uuuvn May 1, 2026 23:11

uuuvn requested changes May 3, 2026

View reviewed changes

LuckyIYI added 4 commits May 5, 2026 16:20

review + self-review fixes

15a23c3

resolving conflicts with main

a444669

fix tracer

ea3e8ad

Merge branch 'main' into buffers-to-allocations

9911411

LuckyIYI requested a review from uuuvn May 6, 2026 04:02

Merge branch 'main' into buffers-to-allocations

153db09

eugenebokhan requested changes May 7, 2026

View reviewed changes

LuckyIYI added 4 commits May 7, 2026 13:49

error handling, result as return, allocation on lower level

29b85f8

Merge branch 'main' into buffers-to-allocations

6035d25

structural improvements

a5c8d92

moe temp buffers are moved form argument struct

bfc521d

LuckyIYI requested a review from eugenebokhan May 8, 2026 11:07

CI fix

c493a2f

		let (shape, array) = read_float_array::<B, 1>(tree, name, expected_data_type)?;
		let (shape, allocation) = read_float_allocation::<B, 1>(tree, name, expected_data_type)?;

		pub hidden: &'a mut Allocation<B>,
		pub output: &'a mut Allocation<B>,

Conversation

LuckyIYI commented Apr 16, 2026

Uh oh!

norpadon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants